@Test
public void testGenerateConsistencyToken() throws TimeoutException, InterruptedException {

  GenerateConsistencyTokenRequest request =
      GenerateConsistencyTokenRequest.newBuilder().setName(TABLE_NAME).build();
  final AtomicInteger counter = new AtomicInteger(0);
  Answer<Void> answer =
      new Answer<Void>() {
        @Override
        public Void answer(final InvocationOnMock invocation) throws Throwable {
          Listener listener = invocation.getArgument(0, Listener.class);
          if (counter.incrementAndGet() == 1) {
            listener.onMessage(GenerateConsistencyTokenResponse.getDefaultInstance());
          } else {
            listener.onMessage(CheckConsistencyResponse.newBuilder().setConsistent(true).build());
          }
          listener.onClose(Status.OK, null);
          return null;
        }
      };
  doAnswer(answer).when(mockClientCall).start(any(Listener.class), any(Metadata.class));
  defaultClient.waitForReplication(INSTANCE_NAME.toTableName("TABLE_NAME"), 6000);
  verify(mockClientCall, times(2)).start(any(Listener.class), any(Metadata.class));
  verify(mockClientCall, times(1))
      .sendMessage(ArgumentMatchers.isA(GenerateConsistencyTokenRequest.class));
  verify(mockClientCall, times(1))
      .sendMessage(ArgumentMatchers.isA(CheckConsistencyRequest.class));
  verify(mockClientCall, times(2)).halfClose();
}