private Operation<Void, UpdateDatabaseDdlMetadata> mockOperation(boolean error) {
  @SuppressWarnings("unchecked")
  Operation<Void, UpdateDatabaseDdlMetadata> op = mock(Operation.class);
  when(op.getName()).then(new Returns("TEST_OPERATION"));
  when(op.isDone()).then(new Answer<Boolean>() {
    @Override
    public Boolean answer(InvocationOnMock invocation) throws Throwable {
      return reportDone;
    }
  });
  when(op.reload()).then(new Returns(op));
  if (error)
    when(op.getResult()).thenThrow(SpannerExceptionFactory
        .newSpannerException(ErrorCode.INVALID_ARGUMENT, "Some exception"));
  else
    when(op.getResult()).then(new Returns(null));
  UpdateDatabaseDdlMetadata metadata = UpdateDatabaseDdlMetadata.getDefaultInstance();
  when(op.getMetadata()).then(new Returns(metadata));

  return op;
}
 

@Test
public void testAutoBatch() throws SQLException, NoSuchFieldException, SecurityException,
    IllegalArgumentException, IllegalAccessException {
  CloudSpannerConnection connection = createConnection();
  Mockito.doCallRealMethod().when(connection).addAutoBatchedDdlOperation(Mockito.anyString());
  Mockito.doCallRealMethod().when(connection).clearAutoBatchedDdlOperations();
  Mockito.doCallRealMethod().when(connection).getAutoBatchedDdlOperations();
  Mockito.when(connection.isAutoBatchDdlOperations()).then(new Returns(true));
  Field field = CloudSpannerConnection.class.getDeclaredField("autoBatchedDdlOperations");
  field.setAccessible(true);
  field.set(connection, new ArrayList<String>());

  CloudSpannerStatement statement = connection.createStatement();
  assertEquals(0, connection.getAutoBatchedDdlOperations().size());
  statement.execute(BATCH_DDL);
  assertEquals(1, connection.getAutoBatchedDdlOperations().size());
  statement.execute(BATCH_DML);
  assertEquals(1, connection.getAutoBatchedDdlOperations().size());
  statement.execute("EXECUTE_DDL_BATCH");
  assertEquals(0, connection.getAutoBatchedDdlOperations().size());
}
 

@Test
public void shouldGetResultsForMethods() throws Throwable {
    invocationContainerImpl.setInvocationForPotentialStubbing(new InvocationMatcher(simpleMethod));
    invocationContainerImpl.addAnswer(new Returns("simpleMethod"));
    
    Invocation differentMethod = new InvocationBuilder().differentMethod().toInvocation();
    invocationContainerImpl.setInvocationForPotentialStubbing(new InvocationMatcher(differentMethod));
    invocationContainerImpl.addAnswer(new ThrowsException(new MyException()));
    
    assertEquals("simpleMethod", invocationContainerImpl.answerTo(simpleMethod));
    
    try {
        invocationContainerImpl.answerTo(differentMethod);
        fail();
    } catch (MyException e) {}
}
 

@Test
public void should_get_results_for_methods() throws Throwable {
    invocationContainerImpl.setInvocationForPotentialStubbing(new InvocationMatcher(simpleMethod));
    invocationContainerImpl.addAnswer(new Returns("simpleMethod"));

    Invocation differentMethod = new InvocationBuilder().differentMethod().toInvocation();
    invocationContainerImpl.setInvocationForPotentialStubbing(new InvocationMatcher(differentMethod));
    invocationContainerImpl.addAnswer(new ThrowsException(new MyException()));

    assertEquals("simpleMethod", invocationContainerImpl.answerTo(simpleMethod));

    try {
        invocationContainerImpl.answerTo(differentMethod);
        fail();
    } catch (MyException e) {}
}
 

@Test
public void should_get_results_for_methods_stub_only() throws Throwable {
    invocationContainerImplStubOnly.setInvocationForPotentialStubbing(new InvocationMatcher(simpleMethod));
    invocationContainerImplStubOnly.addAnswer(new Returns("simpleMethod"));

    Invocation differentMethod = new InvocationBuilder().differentMethod().toInvocation();
    invocationContainerImplStubOnly.setInvocationForPotentialStubbing(new InvocationMatcher(differentMethod));
    invocationContainerImplStubOnly.addAnswer(new ThrowsException(new MyException()));

    assertEquals("simpleMethod", invocationContainerImplStubOnly.answerTo(simpleMethod));

    try {
        invocationContainerImplStubOnly.answerTo(differentMethod);
        fail();
    } catch (MyException e) {}
}
 

@Test
//see issue 331
public void handle_result_must_not_be_null_for_primitives() throws Throwable {
    //given:
    MockCreationSettings settings = (MockCreationSettings) new MockSettingsImpl().defaultAnswer(new Returns(null));
    InternalMockHandler handler = new MockHandlerFactory().create(settings);

    mock.intReturningMethod();
    Invocation invocation = super.getLastInvocation();

    //when:
    Object result = handler.handle(invocation);

    //then null value is not a valid result for a primitive
    assertNotNull(result);
    assertEquals(0, result);
}
 

@Test
//see issue 331
public void valid_handle_result_is_permitted() throws Throwable {
    //given:
    MockCreationSettings settings = (MockCreationSettings) new MockSettingsImpl().defaultAnswer(new Returns(123));
    InternalMockHandler handler = new MockHandlerFactory().create(settings);

    mock.intReturningMethod();
    Invocation invocation = super.getLastInvocation();

    //when:
    Object result = handler.handle(invocation);

    //then
    assertEquals(123, result);
}
 

@Test
public void testShowDDLOperations() throws SQLException, NoSuchFieldException, SecurityException,
    IllegalArgumentException, IllegalAccessException {
  final String ddl =
      "CREATE TABLE FOO (ID INT64 NOT NULL, NAME STRING(100) NOT NULL) PRIMARY KEY (ID)";

  Field adminClientField = CloudSpannerConnection.class.getDeclaredField("adminClient");
  adminClientField.setAccessible(true);
  DatabaseAdminClient adminClient = mock(DatabaseAdminClient.class);
  @SuppressWarnings("unchecked")
  Operation<Void, UpdateDatabaseDdlMetadata> operation = mock(Operation.class);
  when(operation.reload()).then(new Returns(operation));
  when(operation.getName()).then(new Returns("test"));
  when(adminClient.updateDatabaseDdl(any(), any(), any(), any())).then(new Returns(operation));
  adminClientField.set(connection, adminClient);
  Statement statement = connection.createStatement();
  assertFalse(statement.execute("SET_CONNECTION_PROPERTY AsyncDdlOperations=true"));
  assertEquals(1, statement.getUpdateCount());
  try (ResultSet rs = statement.executeQuery("SHOW_DDL_OPERATIONS")) {
    assertFalse(rs.next());
  }
  statement.execute(ddl);
  try (ResultSet rs = statement.executeQuery("SHOW_DDL_OPERATIONS")) {
    assertTrue(rs.next());
    assertEquals("test", rs.getString("NAME"));
    assertNotNull(rs.getTimestamp("TIME_STARTED"));
    assertEquals(ddl, rs.getString("STATEMENT"));
    assertFalse(rs.getBoolean("DONE"));
    assertNull(rs.getString("EXCEPTION"));
    assertFalse(rs.next());
  }
}
 

@Test
public void testExecuteBatchReadOnly() throws SQLException, NoSuchFieldException,
    SecurityException, IllegalArgumentException, IllegalAccessException {
  for (int testRun = 0; testRun < 2; testRun++) {
    final int numberOfPartitions = 6;
    BatchClient batchClient = mock(BatchClient.class);
    BatchReadOnlyTransaction tx = mock(BatchReadOnlyTransaction.class);
    List<Partition> partitions = new ArrayList<>(numberOfPartitions);
    for (int i = 0; i < numberOfPartitions; i++)
      partitions.add(mock(Partition.class));
    when(tx.partitionQuery(any(), any())).then(new Returns(partitions));
    when(batchClient.batchReadOnlyTransaction(TimestampBound.strong())).then(new Returns(tx));
    Field field = CloudSpannerTransaction.class.getDeclaredField("batchClient");
    field.setAccessible(true);
    field.set(connection.getTransaction(), batchClient);
    connection.setBatchReadOnly(true);
    Statement statement;
    if (testRun % 2 == 0) {
      statement = connection.createStatement();
      assertTrue(statement.execute(SELECT_ALL_FROM_FOO));
    } else {
      PreparedStatement ps = connection.prepareStatement(SELECT_ALL_FROM_FOO);
      assertTrue(ps.execute());
      statement = ps;
    }
    List<ResultSet> resultSets = new ArrayList<>();
    do {
      resultSets.add(statement.getResultSet());
    } while (statement.getMoreResults());
    assertEquals(numberOfPartitions, resultSets.size());
  }
}
 

@Test
public void testExecuteNormal() throws SQLException, NoSuchFieldException, SecurityException,
    IllegalArgumentException, IllegalAccessException {
  for (int testRun = 0; testRun < 2; testRun++) {
    CloudSpannerTransaction tx = mock(CloudSpannerTransaction.class);
    com.google.cloud.spanner.ResultSet rs = mock(com.google.cloud.spanner.ResultSet.class);
    when(tx.executeQuery(any())).then(new Returns(rs));
    Field field = CloudSpannerConnection.class.getDeclaredField("transaction");
    field.setAccessible(true);
    field.set(connection, tx);
    Statement statement;
    if (testRun % 2 == 0) {
      statement = connection.createStatement();
      assertTrue(statement.execute(SELECT_ALL_FROM_FOO));
    } else {
      PreparedStatement ps = connection.prepareStatement(SELECT_ALL_FROM_FOO);
      assertTrue(ps.execute());
      statement = ps;
    }
    List<ResultSet> resultSets = new ArrayList<>();
    do {
      resultSets.add(statement.getResultSet());
    } while (statement.getMoreResults());
    assertEquals(1, resultSets.size());
    connection.commit();
  }
}
 

public static CloudSpannerConnection create(String url) throws SQLException {
  ConnectionProperties properties = ConnectionProperties.parse(url);
  CloudSpannerConnection connection = mock(CloudSpannerConnection.class);
  when(connection.getSimulateMajorVersion()).then(new Returns(null));
  when(connection.getSimulateMinorVersion()).then(new Returns(null));
  when(connection.getUrl()).thenAnswer(new Returns(url));
  when(connection.getProductName()).thenAnswer(new Returns(properties.productName));
  when(connection.getNodeCount()).thenAnswer(new Returns(1));
  return connection;
}
 

public static CloudSpannerConnection createConnection() throws SQLException {
  CloudSpannerConnection connection = Mockito.mock(CloudSpannerConnection.class);
  Mockito.doCallRealMethod().when(connection).setAutoCommit(Mockito.anyBoolean());
  Mockito.when(connection.getAutoCommit()).thenCallRealMethod();
  connection.setAutoCommit(false);
  Mockito.doCallRealMethod().when(connection).setBatchReadOnly(Mockito.anyBoolean());
  Mockito.when(connection.isBatchReadOnly()).thenCallRealMethod();
  Mockito.doCallRealMethod().when(connection).setReadOnly(Mockito.anyBoolean());
  Mockito.when(connection.isReadOnly()).thenCallRealMethod();

  Mockito.when(connection.isAllowExtendedMode()).thenAnswer(new Returns(true));
  Mockito.when(connection.createArrayOf(Mockito.anyString(), Mockito.any())).thenCallRealMethod();
  CloudSpannerDatabaseMetaData metadata = createMetaData();
  Mockito.when(connection.getMetaData()).thenReturn(metadata);
  CloudSpannerTransaction transaction = Mockito.mock(CloudSpannerTransaction.class);
  Mockito.when(transaction.executeQuery(Mockito.any()))
      .thenReturn(Mockito.mock(com.google.cloud.spanner.ResultSet.class));
  Mockito.when(transaction.partitionQuery(Mockito.any(), Mockito.any())).thenReturn(
      Arrays.asList(mock(Partition.class), mock(Partition.class), mock(Partition.class)));
  Mockito.when(connection.getTransaction()).thenReturn(transaction);

  TableKeyMetaData tableFoo = Mockito.mock(TableKeyMetaData.class);
  Mockito.when(tableFoo.getKeyColumns()).thenAnswer(new Returns(Arrays.asList("ID")));
  Mockito
      .when(connection.getTable(Mockito
          .matches(Pattern.compile("FOO", Pattern.CASE_INSENSITIVE | Pattern.UNICODE_CASE))))
      .thenAnswer(new Returns(tableFoo));

  TableKeyMetaData tableBar = Mockito.mock(TableKeyMetaData.class);
  Mockito.when(tableBar.getKeyColumns()).thenAnswer(new Returns(Arrays.asList("ID1", "ID2")));
  Mockito
      .when(connection.getTable(Mockito
          .matches(Pattern.compile("BAR", Pattern.CASE_INSENSITIVE | Pattern.UNICODE_CASE))))
      .thenAnswer(new Returns(tableBar));
  Mockito.when(connection.getLogger()).thenAnswer(new Returns(new Logger()));

  mockXAMethods(connection);

  return connection;
}
 

private static void mockXAMethods(CloudSpannerConnection connection) throws SQLException {
  String checkTable = null;
  try {
    Field field = CloudSpannerXAConnection.class.getDeclaredField("CHECK_TABLE_EXISTENCE");
    field.setAccessible(true);
    checkTable = (String) field.get(null);
  } catch (Exception e) {
    throw new RuntimeException(e);
  }
  CloudSpannerPreparedStatement ps = Mockito.mock(CloudSpannerPreparedStatement.class);
  CloudSpannerResultSet rs = Mockito.mock(CloudSpannerResultSet.class);
  Mockito.when(rs.next()).thenReturn(true, false);
  Mockito.when(connection.prepareStatement(checkTable)).thenAnswer(new Returns(ps));
  Mockito.when(ps.executeQuery()).thenAnswer(new Returns(rs));
}
 

@Test
void shouldReturnEmptyInstancesListOnTimeoutBlockingClient() {
	environment.setProperty(SERVICE_DISCOVERY_TIMEOUT, "100ms");
	when(discoveryClient.getInstances(SERVICE_ID)).thenAnswer(
			new AnswersWithDelay(200, new Returns(Lists.list(instance("1host", false),
					instance("2host-secure", true), instance("3host", false)))));
	StepVerifier
			.create(new DiscoveryClientServiceInstanceListSupplier(discoveryClient,
					environment).get())
			.expectSubscription().expectNext(Lists.emptyList()).thenCancel()
			.verify(VERIFICATION_TIMEOUT);
}
 

private TestSubject createTestSubject() {
  MockTransactionRunner runner = new MockTransactionRunner();
  DatabaseClient dbClient = mock(DatabaseClient.class);
  when(dbClient.readWriteTransaction()).then(new Returns(runner));
  return new TestSubject(new TransactionThread(dbClient, new Logger()), runner.mock);
}
 

public Stubber doReturn(Object toBeReturned) {
    answers.add(new Returns(toBeReturned));
    return this;
}
 

public OngoingStubbing<T> thenReturn(T value) {
    return thenAnswer(new Returns(value));
}
 

public DeprecatedOngoingStubbing<T> toReturn(T value) {
    return toAnswer(new Returns(value));
}
 

@Test
public void shouldFinishStubbingOnAddingReturnValue() throws Exception {
    state.stubbingStarted();
    invocationContainerImpl.addAnswer(new Returns("test"));
    state.validateState();
}
 

public OngoingStubbing<T> thenReturn(T value) {
    return thenAnswer(new Returns(value));
}
 

public DeprecatedOngoingStubbing<T> toReturn(T value) {
    return toAnswer(new Returns(value));
}
 

@Test
public void should_finish_stubbing_on_adding_return_value() throws Exception {
    state.stubbingStarted();
    invocationContainerImpl.addAnswer(new Returns("test"));
    state.validateState();
}
 

/**
 * Use doReturn() in those rare occasions when you cannot use {@link Mockito#when(Object)}.
 * <p>
 * <b>Beware that {@link Mockito#when(Object)} is always recommended for stubbing because it is argument type-safe 
 * and more readable</b> (especially when stubbing consecutive calls). 
 * <p>
 * Here are those rare occasions when doReturn() comes handy:
 * <p>
 * 
 * 1. When spying real objects and calling real methods on a spy brings side effects  
 * 
 * <pre>
 *   List list = new LinkedList();
 *   List spy = spy(list);
 *   
 *   //Impossible: real method is called so spy.get(0) throws IndexOutOfBoundsException (the list is yet empty)
 *   when(spy.get(0)).thenReturn("foo");
 *   
 *   //You have to use doReturn() for stubbing:
 *   doReturn("foo").when(spy).get(0);
 * </pre>
 * 
 * 2. Overriding a previous exception-stubbing:
 * 
 * <pre>
 *   when(mock.foo()).thenThrow(new RuntimeException());
 *   
 *   //Impossible: the exception-stubbed foo() method is called so RuntimeException is thrown. 
 *   when(mock.foo()).thenReturn("bar");
 *   
 *   //You have to use doReturn() for stubbing:
 *   doReturn("bar").when(mock).foo();
 * </pre>
 * 
 * Above scenarios shows a tradeoff of Mockito's ellegant syntax. Note that the scenarios are very rare, though. 
 * Spying should be sporadic and overriding exception-stubbing is very rare. Not to mention that in general
 * overridding stubbing is a potential code smell that points out too much stubbing.
 * <p>
 * See examples in javadoc for {@link Mockito} class
 * 
 * @param toBeReturned to be returned when the stubbed method is called
 * @return stubber - to select a method for stubbing
 */
public static Stubber doReturn(Object toBeReturned) {
    return MOCKITO_CORE.doAnswer(new Returns(toBeReturned));
}