package com.atlassian.sal.core.rdbms;
   
   import com.atlassian.sal.api.rdbms.ConnectionCallback;
   import com.atlassian.sal.spi.HostConnectionAccessor;
   import org.junit.Before;
   import org.junit.Rule;
   import org.junit.rules.ExpectedException;
   import org.mockito.Mock;
   import org.mockito.MockitoAnnotations;
  import org.mockito.invocation.InvocationOnMock;
  import org.mockito.stubbing.Answer;
  
  import java.sql.Connection;
  
  import static org.mockito.Matchers.any;
  import static org.mockito.Mockito.when;
  
  public class TestDefaultTransactionalExecutorBase {
      @Rule
      public ExpectedException expectedException = ExpectedException.none();
  
      protected DefaultTransactionalExecutor defaultTransactionalExecutor;
  
      @Mock
      protected HostConnectionAccessor hostConnectionAccessor;
      @Mock
      protected Connection connection;
      @Mock
      protected ConnectionCallback<Object> callback;
      @Mock
      protected Object result;
  
      protected WrappedConnection wrappedConnection;
  
      protected Throwable callbackThrows;
  
      protected static class CallbackException extends RuntimeException {
          public CallbackException(final String message) {
              super(message);
          }
      }
  
      @Before
      public void before() {
          MockitoAnnotations.initMocks(this);
  
          defaultTransactionalExecutor = new DefaultTransactionalExecutor(hostConnectionAccessor, false, false);
  
          callbackThrows = null;
  
          // grab a hold of the wrapped connection each time
          when(callback.execute(any(WrappedConnection.class))).thenAnswer(new Answer<Object>() {
              @Override
              public Object answer(final InvocationOnMock invocation) throws Throwable {
                  wrappedConnection = (WrappedConnection) invocation.getArguments()[0];
                  if (callbackThrows != null) {
                      throw callbackThrows;
                  } else {
                      return result;
                  }
              }
          });
      }
  }