DataClientVeneerApi implements DataClientWrapper { @Override public ApiFuture<List<Result>> readRowsAsync(Query request) { return delegate.readRowsCallable(RESULT_ADAPTER).all().futureCall(request); } DataClientVeneerApi(BigtableDataClient delegate); @Override BulkMutationWrapper createBulkMutation(String tableId); @Override BulkReadWrapper createBulkRead(String tableId); @Override ApiFuture<Void> mutateRowAsync(RowMutation rowMutation); @Override ApiFuture<Result> readModifyWriteRowAsync(ReadModifyWriteRow readModifyWriteRow); @Override ApiFuture<Boolean> checkAndMutateRowAsync(ConditionalRowMutation conditionalRowMutation); @Override ApiFuture<List<KeyOffset>> sampleRowKeysAsync(String tableId); @Override ApiFuture<Result> readRowAsync(
      String tableId, ByteString rowKey, @Nullable Filters.Filter filter); @Override ResultScanner readRows(Query request); @Override ApiFuture<List<Result>> readRowsAsync(Query request); @Override void readRowsAsync(Query request, StreamObserver<Result> observer); @Override void close();  }
@Test public void testReadRowsAsync() throws Exception { Query query = Query.create(TABLE_ID).rowKey(ROW_KEY); when(mockDataClient.readRowsCallable(Mockito.<RowResultAdapter>any())) .thenReturn(mockStreamingCallable); when(mockStreamingCallable.all()).thenReturn(mockUnaryCallable); List<Result> expectedResult = ImmutableList.of(Result.EMPTY_RESULT, EXPECTED_RESULT); when(mockUnaryCallable.futureCall(query)) .thenReturn(ApiFutures.immediateFuture(expectedResult)); List<Result> actualResult = dataClientWrapper.readRowsAsync(query).get(); assertEquals(expectedResult.size(), actualResult.size()); assertResult(Result.EMPTY_RESULT, actualResult.get(0)); assertResult(EXPECTED_RESULT, actualResult.get(1)); verify(mockDataClient).readRowsCallable(Mockito.<RowResultAdapter>any()); verify(mockStreamingCallable).all(); verify(mockUnaryCallable).futureCall(query); }
@Test public void testReadRowsAsyncWithStreamOb() { final Exception readException = new Exception(); Query request = Query.create(TABLE_ID).rowKey(ROW_KEY); StreamObserver<Result> resultStreamOb = new StreamObserver<Result>() { @Override public void onNext(Result result) { assertResult(EXPECTED_RESULT, result); } @Override public void onError(Throwable throwable) { assertEquals(readException, throwable); } @Override public void onCompleted() {} }; when(mockDataClient.readRowsCallable(Mockito.<RowResultAdapter>any())) .thenReturn(mockStreamingCallable); doAnswer( new Answer() { int count = 0; @Override public Object answer(InvocationOnMock invocationOnMock) { ResponseObserver<Result> resObserver = invocationOnMock.getArgument(1); resObserver.onStart(null); resObserver.onResponse(EXPECTED_RESULT); if (count == 0) { resObserver.onComplete(); } else { resObserver.onError(readException); } count++; return null; } }) .when(mockStreamingCallable) .call(Mockito.<Query>any(), Mockito.<ResponseObserver<Result>>any()); dataClientWrapper.readRowsAsync(request, resultStreamOb); dataClientWrapper.readRowsAsync(request, resultStreamOb); verify(mockDataClient, times(2)).readRowsCallable(Mockito.<RowResultAdapter>any()); verify(mockStreamingCallable, times(2)) .call(Mockito.<Query>any(), Mockito.<ResponseObserver<Result>>any()); }