use of com.google.cloud.bigtable.grpc.scanner.FlatRow in project beam by apache.
the class BigtableServiceImplTest method testReadByteLimitBuffer.
/**
* This test checks that the buffer will stop filling up once the byte limit is reached. It will
* cancel the ScanHandler after reached the limit. This test completes one fill and contains one
* Row after the first buffer has been completed. The test cheaks the current available memory in
* the JVM and uses a percent of it to mock the original behavior. Range: [b00000, b00010)
*
* @throws IOException
*/
@Test
public void testReadByteLimitBuffer() throws IOException {
long segmentByteLimit = DEFAULT_ROW_SIZE * (SEGMENT_SIZE / 2);
int numOfRowsInsideBuffer = (int) (segmentByteLimit / DEFAULT_ROW_SIZE) + 1;
RowRange mockRowRange = generateRowRange(generateByteString(DEFAULT_PREFIX, 0), generateByteString(DEFAULT_PREFIX, SEGMENT_SIZE));
OngoingStubbing<?> stub = when(mockBigtableDataClient.readFlatRows(any(ReadRowsRequest.class), any()));
List<List<FlatRow>> expectedResults = ImmutableList.of(generateLargeSegmentResult(DEFAULT_PREFIX, 0, numOfRowsInsideBuffer), generateSegmentResult(DEFAULT_PREFIX, numOfRowsInsideBuffer, SEGMENT_SIZE - numOfRowsInsideBuffer), ImmutableList.of());
expectRowResults(stub, expectedResults);
BigtableService.Reader underTest = new BigtableServiceImpl.BigtableSegmentReaderImpl(mockSession, TABLE_ID, RowSet.newBuilder().addRowRanges(mockRowRange).build(), SEGMENT_SIZE, segmentByteLimit, RowFilter.getDefaultInstance(), mockCallMetric);
List<Row> actualResults = new ArrayList<>();
Assert.assertTrue(underTest.start());
do {
actualResults.add(underTest.getCurrentRow());
} while (underTest.advance());
Assert.assertEquals(expectedResults.stream().flatMap(Collection::stream).map(i -> FlatRowConverter.convert(i)).collect(Collectors.toList()), actualResults);
underTest.close();
Mockito.verify(mockCallMetric, Mockito.times(3)).call("ok");
}
use of com.google.cloud.bigtable.grpc.scanner.FlatRow in project beam by apache.
the class BigtableServiceImplTest method testReadSingleRangeBelowSegmentLimit.
/**
* 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 {
RowSet.Builder ranges = RowSet.newBuilder();
ranges.addRowRanges(generateRowRange(generateByteString(DEFAULT_PREFIX, 0), generateByteString(DEFAULT_PREFIX, 1)));
FlatRow expectedRow = FlatRow.newBuilder().withRowKey(ByteString.copyFromUtf8("a")).build();
when(mockBigtableDataClient.readFlatRows(any(ReadRowsRequest.class), any())).thenAnswer(mockReadRowsAnswer(Arrays.asList(expectedRow)));
BigtableService.Reader underTest = new BigtableServiceImpl.BigtableSegmentReaderImpl(mockSession, TABLE_ID, ranges.build(), SEGMENT_SIZE, DEFAULT_BYTE_SEGMENT_SIZE, RowFilter.getDefaultInstance(), mockCallMetric);
underTest.start();
Assert.assertEquals(FlatRowConverter.convert(expectedRow), underTest.getCurrentRow());
Assert.assertFalse(underTest.advance());
underTest.close();
Mockito.verify(mockCallMetric, Mockito.times(2)).call("ok");
}
Aggregations