igorbernstein2 commented on code in PR #16939:
URL: https://github.com/apache/beam/pull/16939#discussion_r873875195


##########
sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImplTest.java:
##########
@@ -126,6 +154,354 @@ public void testRead() throws IOException {
     verifyMetricWasSet("google.bigtable.v2.ReadRows", "ok", 1);
   }
 
+  /**
+   * This test ensures that protobuf creation and interactions with {@link 
BigtableDataClient} work
+   * as expected. This test checks that a single row is returned from the 
future.
+   *
+   * @throws IOException
+   */
+  @Test
+  public void testReadSingleRangeBelowSegmentLimit() throws Exception {
+    ByteKey start = ByteKey.copyFrom("a".getBytes(StandardCharsets.UTF_8));
+    ByteKey end = ByteKey.copyFrom("b".getBytes(StandardCharsets.UTF_8));
+
+    
when(mockBigtableSource.getRanges()).thenReturn(Arrays.asList(ByteKeyRange.of(start,
 end)));
+    FlatRow expectedRow = 
FlatRow.newBuilder().withRowKey(ByteString.copyFromUtf8("a")).build();
+
+    when(mockBigtableDataClient.readFlatRows(any(ReadRowsRequest.class), 
any()))
+        .thenAnswer(mockReadRowsAnswer(Arrays.asList(expectedRow)));
+
+    BigtableService.Reader underTest =
+        BigtableServiceImpl.BigtableSegmentReaderImpl.create(mockSession, 
mockBigtableSource);

Review Comment:
   Yes, this will require another helper to create the read rows request from a 
list of ranges. But it makes the test easier to understand...when debugging a 
test failure, I have one less class to worry about. Overall I think it will 
actually be less noise: you get to remove one mock expectation and replace it 
with a plain old function call that creates a request



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to