Search in sources :

Example 1 with BulkMutation

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();
}
Also used : MutateRowResponse(com.google.bigtable.v2.MutateRowResponse) Entry(com.google.bigtable.v2.MutateRowsRequest.Entry) ByteString(com.google.protobuf.ByteString) Mutation(com.google.bigtable.v2.Mutation) BulkMutation(com.google.cloud.bigtable.grpc.async.BulkMutation) Test(org.junit.Test)

Aggregations

MutateRowResponse (com.google.bigtable.v2.MutateRowResponse)1 Entry (com.google.bigtable.v2.MutateRowsRequest.Entry)1 Mutation (com.google.bigtable.v2.Mutation)1 BulkMutation (com.google.cloud.bigtable.grpc.async.BulkMutation)1 ByteString (com.google.protobuf.ByteString)1 Test (org.junit.Test)1