use of com.google.cloud.bigtable.grpc.async.BulkMutation in project beam by apache.
the class BigtableServiceImplTest method testWrite.
/**
* This test ensures that protobuf creation and interactions with {@link BulkMutation} work as
* expected.
*
* @throws IOException
* @throws InterruptedException
*/
@Test
public void testWrite() throws IOException, InterruptedException {
BigtableService.Writer underTest = new BigtableServiceImpl.BigtableWriterImpl(mockSession, TABLE_NAME);
Mutation mutation = Mutation.newBuilder().setSetCell(SetCell.newBuilder().setFamilyName("Family").build()).build();
ByteString key = ByteString.copyFromUtf8("key");
SettableFuture<MutateRowResponse> fakeResponse = SettableFuture.create();
when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class))).thenReturn(fakeResponse);
underTest.writeRecord(KV.of(key, ImmutableList.of(mutation)));
Entry expected = MutateRowsRequest.Entry.newBuilder().setRowKey(key).addMutations(mutation).build();
verify(mockBulkMutation, times(1)).add(expected);
underTest.close();
verify(mockBulkMutation, times(1)).flush();
}
Aggregations