Repository: GoogleCloudPlatform/cloud-bigtable-client
Pull request #2107 (2019-02-21 14:24:23)

BulkMutationWrapper with GCJ's model
- Introduced `BulkMutationWrapper.java` to wrap existing `BulkMutation.java`.
- Have currently adapted againsted `ListenableFuture` for both `IBulkMutation#add` & `IBulkMutation#readModifyWrite`.


Patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/core/IBulkMutation.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/core/IBulkMutation.java
index f964b2a34..fa46c3b7b 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/core/IBulkMutation.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/core/IBulkMutation.java
@@ -15,38 +15,48 @@
  */
 package com.google.cloud.bigtable.core;
 
-import com.google.api.core.ApiFuture;
-import com.google.bigtable.v2.MutateRowRequest;
+import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
+import com.google.cloud.bigtable.data.v2.models.Row;
 import com.google.cloud.bigtable.data.v2.models.RowMutation;
 
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
+import com.google.cloud.bigtable.grpc.async.OperationAccountant;
+import com.google.common.util.concurrent.ListenableFuture;
 
 /**
  * Interface to support batching multiple RowMutation request into a single grpc request.
  */
 public interface IBulkMutation {
-  long MAX_RPC_WAIT_TIME_NANOS = TimeUnit.MINUTES.toNanos(12);
 
   /**
-   * Send any outstanding {@link MutateRowRequest}s and wait until all requests are complete.
+   * Send any outstanding {@link RowMutation} and wait until all requests are complete.
    */
-  void flush() throws InterruptedException, TimeoutException;
+  void flush() throws InterruptedException;
 
   void sendUnsent();
 
   /**
-   * @return false if there is any outstanding {@link MutateRowRequest} that still needs to be sent.
+   * @return false if there is any outstanding {@link RowMutation} that still needs to be sent.
    */
   boolean isFlushed();
 
   /**
-   * Adds a {@link com.google.cloud.bigtable.data.v2.models.RowMutation} to the underlying IBulkMutation
-   * mechanism.
+   * Adds a {@link RowMutation} to the underlying IBulkMutation mechanism.
    *
-   * @param rowMutation The {@link com.google.cloud.bigtable.data.v2.models.RowMutation} to add
-   * @return a {@link ApiFuture} of type {@link Void} will be set when request is
+   * @param rowMutation The {@link RowMutation} to add.
+   * @return a {@link ListenableFuture} of type {@link Void} will be set when request is
    *     successful otherwise exception will be thrown.
    */
-  ApiFuture<Void> add(RowMutation rowMutation);
+  //TODO(rahulkql): Once it is adapted to v2.models, change the return type to ApiFuture.
+  ListenableFuture<Void> add(RowMutation rowMutation);
+
+  /**
+   * Performs a {@link IBigtableDataClient#readModifyWriteRowAsync(ReadModifyWriteRow)} on the
+   * {@link ReadModifyWriteRow}. This method may block if
+   * {@link OperationAccountant#registerOperation(ListenableFuture)} blocks.
+   *
+   * @param request The {@link ReadModifyWriteRow} to send.
+   * @return a {@link ListenableFuture} which can be listened to for completion events.
+   */
+  //TODO(rahulkql): Once it is adapted to v2.models, change the return type to ApiFuture.
+  ListenableFuture<Row> readModifyWrite(ReadModifyWriteRow request);
 }
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableSession.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableSession.java
index 7c55f6ecc..8cb614fa4 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableSession.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableSession.java
@@ -28,8 +28,10 @@ import com.google.cloud.bigtable.config.Logger;
 import com.google.cloud.bigtable.config.RetryOptions;
 import com.google.cloud.bigtable.core.IBigtableDataClient;
 import com.google.cloud.bigtable.core.IBigtableTableAdminClient;
+import com.google.cloud.bigtable.core.IBulkMutation;
 import com.google.cloud.bigtable.data.v2.internal.RequestContext;
 import com.google.cloud.bigtable.grpc.async.BulkMutation;
+import com.google.cloud.bigtable.grpc.async.BulkMutationWrapper;
 import com.google.cloud.bigtable.grpc.async.BulkRead;
 import com.google.cloud.bigtable.grpc.async.ResourceLimiter;
 import com.google.cloud.bigtable.grpc.async.ResourceLimiterStats;
@@ -340,6 +342,16 @@ public class BigtableSession implements Closeable {
         options.getBulkOptions());
   }
 
+  /**
+   * <p>createBulkMutationWrapper.</p>
+   *
+   * @param tableName a {@link BigtableTableName} object.
+   * @return a {@link IBigtableDataClient} object.
+   */
+  public IBulkMutation createBulkMutationWrapper(BigtableTableName tableName) {
+    return new BulkMutationWrapper(createBulkMutation(tableName), getDataRequestContext());
+  }
+
   /**
    * <p>createBulkRead.</p>
    *
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/BulkMutationWrapper.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/BulkMutationWrapper.java
new file mode 100644
index 000000000..94aefdb73
--- /dev/null
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/BulkMutationWrapper.java
@@ -0,0 +1,78 @@
+/*
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.bigtable.grpc.async;
+
+import com.google.bigtable.v2.MutateRowResponse;
+import com.google.cloud.bigtable.core.IBulkMutation;
+import com.google.cloud.bigtable.data.v2.internal.RequestContext;
+import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
+import com.google.cloud.bigtable.data.v2.models.Row;
+import com.google.cloud.bigtable.data.v2.models.RowMutation;
+import com.google.common.base.Function;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.MoreExecutors;
+
+/**
+ * This class wraps existing {@link com.google.cloud.bigtable.grpc.async.BulkMutation} with
+ * Google-cloud-java's model.
+ */
+public class BulkMutationWrapper implements IBulkMutation {
+
+  private final BulkMutation delegate;
+  private final RequestContext requestContext;
+
+  public BulkMutationWrapper(BulkMutation bulkMutation, RequestContext requestContext) {
+    this.delegate = bulkMutation;
+    this.requestContext = requestContext;
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  public void flush() throws InterruptedException {
+    delegate.flush();
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  public void sendUnsent() {
+    delegate.sendUnsent();
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  public boolean isFlushed() {
+    return delegate.isFlushed();
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  public ListenableFuture<Void> add(RowMutation rowMutation) {
+    return Futures.transform(delegate.add(rowMutation.toBulkProto(requestContext).getEntries(0)),
+        new Function<MutateRowResponse, Void>() {
+          @Override
+          public Void apply(MutateRowResponse response) {
+            return null;
+          }
+        }, MoreExecutors.directExecutor());
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  public ListenableFuture<Row> readModifyWrite(ReadModifyWriteRow request) {
+    return delegate.readModifyWrite(request);
+  }
+}
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BigtableBufferedMutatorHelper.java b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BigtableBufferedMutatorHelper.java
index b062bf20c..af318af8a 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BigtableBufferedMutatorHelper.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BigtableBufferedMutatorHelper.java
@@ -15,9 +15,8 @@
  */
 package com.google.cloud.bigtable.hbase;
 
-import com.google.bigtable.v2.MutateRowsRequest;
+import com.google.cloud.bigtable.core.IBulkMutation;
 import com.google.cloud.bigtable.data.v2.internal.RequestContext;
-import com.google.cloud.bigtable.data.v2.models.RowMutation;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
@@ -37,7 +36,6 @@ import com.google.cloud.bigtable.config.BigtableOptions;
 import com.google.cloud.bigtable.config.Logger;
 import com.google.cloud.bigtable.grpc.BigtableSession;
 import com.google.cloud.bigtable.grpc.BigtableTableName;
-import com.google.cloud.bigtable.grpc.async.BulkMutation;
 import com.google.cloud.bigtable.hbase.adapters.HBaseRequestAdapter;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
@@ -66,7 +64,7 @@ public class BigtableBufferedMutatorHelper {
   private boolean closed = false;
 
   private final HBaseRequestAdapter adapter;
-  private final BulkMutation bulkMutation;
+  private final IBulkMutation bulkMutation;
   private final BigtableOptions options;
   private final RequestContext requestContext;
 
@@ -86,7 +84,7 @@ public class BigtableBufferedMutatorHelper {
     this.configuration = configuration;
     this.options = session.getOptions();
     BigtableTableName tableName = this.adapter.getBigtableTableName();
-    this.bulkMutation = session.createBulkMutation(tableName);
+    this.bulkMutation = session.createBulkMutationWrapper(tableName);
     this.requestContext = session.getDataRequestContext();
   }
 
@@ -179,7 +177,7 @@ public class BigtableBufferedMutatorHelper {
         return Futures.immediateFailedFuture(
             new IllegalArgumentException("Cannot perform a mutation on a null object."));
       } else {
-        return bulkMutation.add(toEntry(adapter.adaptEntry(mutation)));
+        return bulkMutation.add(adapter.adaptEntry(mutation));
       }
     } finally {
       closedReadLock.unlock();
@@ -198,9 +196,9 @@ public class BigtableBufferedMutatorHelper {
         future = Futures.immediateFailedFuture(
           new IllegalArgumentException("Cannot perform a mutation on a null object."));
       } else if (mutation instanceof Put) {
-        future = bulkMutation.add(toEntry(adapter.adaptEntry((Put) mutation)));
+        future = bulkMutation.add(adapter.adaptEntry((Put) mutation));
       } else if (mutation instanceof Delete) {
-        future = bulkMutation.add(toEntry(adapter.adaptEntry((Delete) mutation)));
+        future = bulkMutation.add(adapter.adaptEntry((Delete) mutation));
       } else if (mutation instanceof Increment) {
         future = bulkMutation.readModifyWrite(adapter.adapt((Increment) mutation));
       } else if (mutation instanceof Append) {
@@ -217,10 +215,6 @@ public class BigtableBufferedMutatorHelper {
     return future;
   }
 
-  private MutateRowsRequest.Entry toEntry(RowMutation rowMutation) {
-    return rowMutation.toBulkProto(requestContext).getEntries(0);
-  }
-
   /**
    * <p>hasInflightRequests.</p>
    *


Test patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestBulkMutationWrapper.java b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestBulkMutationWrapper.java
new file mode 100644
index 000000000..08d05ac81
--- /dev/null
+++ b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestBulkMutationWrapper.java
@@ -0,0 +1,102 @@
+/*
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.bigtable.grpc.async;
+
+import com.google.bigtable.v2.MutateRowResponse;
+import com.google.bigtable.v2.MutateRowsRequest;
+import com.google.cloud.bigtable.core.IBulkMutation;
+import com.google.cloud.bigtable.data.v2.internal.RequestContext;
+import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
+import com.google.cloud.bigtable.data.v2.models.Row;
+import com.google.cloud.bigtable.data.v2.models.RowCell;
+import com.google.cloud.bigtable.data.v2.models.RowMutation;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.protobuf.ByteString;
+import java.util.Collections;
+import java.util.concurrent.Future;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+@RunWith(JUnit4.class) public class TestBulkMutationWrapper {
+
+  private BulkMutation mockDelegate;
+  private RequestContext requestContext =
+      RequestContext.create("ProjectId", "instanceId", "appProfileId");
+  private IBulkMutation bulkWrapper;
+
+  @Before
+  public void setUp() {
+    mockDelegate = Mockito.mock(BulkMutation.class);
+    bulkWrapper = new BulkMutationWrapper(mockDelegate, requestContext);
+  }
+
+  @Test
+  public void testFlush() throws InterruptedException {
+    doNothing().when(mockDelegate).flush();
+    bulkWrapper.flush();
+    verify(mockDelegate).flush();
+  }
+
+  @Test
+  public void testSendUnsend() {
+    doNothing().when(mockDelegate).sendUnsent();
+    bulkWrapper.sendUnsent();
+    verify(mockDelegate).sendUnsent();
+  }
+
+  @Test
+  public void testIsFlushed() {
+    when(mockDelegate.isFlushed()).thenReturn(true);
+    assertTrue(bulkWrapper.isFlushed());
+    verify(mockDelegate).isFlushed();
+  }
+
+  @Test
+  public void testAddMutate() {
+    RowMutation rowMutation = RowMutation.create("tableId", "key");
+    MutateRowsRequest.Entry requestProto = rowMutation.toBulkProto(requestContext).getEntries(0);
+    when(mockDelegate.add(requestProto))
+        .thenReturn(Futures.immediateFuture(MutateRowResponse.getDefaultInstance()));
+    Future<Void> response = bulkWrapper.add(rowMutation);
+    try {
+      response.get();
+    } catch (Exception ex) {
+      throw new AssertionError("Assertion failed for BulkMutationWrapper#add(RowMutation)");
+    }
+    verify(mockDelegate).add(requestProto);
+  }
+
+  @Test
+  public void testReadModifyRow() {
+    ReadModifyWriteRow readModifyRow = ReadModifyWriteRow.create("tableId", "test-key");
+    ListenableFuture<Row> expectedResponse = Futures.immediateFuture(
+        Row.create(ByteString.copyFromUtf8("test-key"), Collections.<RowCell>emptyList()));
+    when(mockDelegate.readModifyWrite(readModifyRow)).thenReturn(expectedResponse);
+    Future<Row> actualResponse = bulkWrapper.readModifyWrite(readModifyRow);
+    assertEquals(expectedResponse, actualResponse);
+    verify(mockDelegate).readModifyWrite(readModifyRow);
+  }
+}
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBatchExecutor.java b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBatchExecutor.java
index 467bf1682..bfb918f6e 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBatchExecutor.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBatchExecutor.java
@@ -24,14 +24,14 @@ import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import com.google.bigtable.v2.MutateRowsRequest;
 import com.google.cloud.bigtable.config.BigtableOptions;
+import com.google.cloud.bigtable.core.IBulkMutation;
 import com.google.cloud.bigtable.data.v2.internal.RequestContext;
 import com.google.cloud.bigtable.data.v2.models.Query;
 import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
+import com.google.cloud.bigtable.data.v2.models.RowMutation;
 import com.google.cloud.bigtable.grpc.BigtableSession;
 import com.google.cloud.bigtable.grpc.BigtableTableName;
-import com.google.cloud.bigtable.grpc.async.BulkMutation;
 import com.google.cloud.bigtable.grpc.async.BulkRead;
 import com.google.cloud.bigtable.grpc.scanner.FlatRow;
 import com.google.cloud.bigtable.grpc.scanner.FlatRow.Cell;
@@ -121,7 +121,7 @@ public class TestBatchExecutor {
   private BulkRead mockBulkRead;
 
   @Mock
-  private BulkMutation mockBulkMutation;
+  private IBulkMutation mockBulkMutation;
 
   @Mock
   private ListenableFuture mockFuture;
@@ -142,10 +142,10 @@ public class TestBatchExecutor {
         .create(options.getProjectId(), options.getInstanceId(), options.getAppProfileId());
 
     MockitoAnnotations.initMocks(this);
-    when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class))).thenReturn(mockFuture);
+    when(mockBulkMutation.add(any(RowMutation.class))).thenReturn(mockFuture);
     when(mockBulkMutation.readModifyWrite(any(ReadModifyWriteRow.class))).thenReturn(mockFuture);
     when(mockBigtableSession.getDataRequestContext()).thenReturn(requetsContext);
-    when(mockBigtableSession.createBulkMutation(any(BigtableTableName.class))).thenReturn(mockBulkMutation);
+    when(mockBigtableSession.createBulkMutationWrapper(any(BigtableTableName.class))).thenReturn(mockBulkMutation);
     when(mockBigtableSession.createBulkRead(any(BigtableTableName.class))).thenReturn(mockBulkRead);
     doAnswer(new Answer<Void>() {
       @Override
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBigtableBufferedMutator.java b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBigtableBufferedMutator.java
index 2392cb56b..47e2284d8 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBigtableBufferedMutator.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBigtableBufferedMutator.java
@@ -22,6 +22,8 @@ import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.google.cloud.bigtable.core.IBulkMutation;
+import com.google.cloud.bigtable.data.v2.models.RowMutation;
 import java.io.IOException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -44,12 +46,9 @@ import org.junit.runners.JUnit4;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import com.google.bigtable.v2.MutateRowResponse;
-import com.google.bigtable.v2.MutateRowsRequest;
 import com.google.cloud.bigtable.config.BigtableOptions;
 import com.google.cloud.bigtable.grpc.BigtableSession;
 import com.google.cloud.bigtable.grpc.BigtableTableName;
-import com.google.cloud.bigtable.grpc.async.BulkMutation;
 import com.google.cloud.bigtable.hbase.adapters.HBaseRequestAdapter;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.SettableFuture;
@@ -69,7 +68,7 @@ public class TestBigtableBufferedMutator {
   private BigtableSession mockSession;
 
   @Mock
-  private BulkMutation mockBulkMutation;
+  private IBulkMutation mockBulkMutation;
 
   @SuppressWarnings("rawtypes")
   private SettableFuture future = SettableFuture.create();
@@ -82,7 +81,7 @@ public class TestBigtableBufferedMutator {
   @Before
   public void setUp() {
     MockitoAnnotations.initMocks(this);
-    when(mockSession.createBulkMutation(any(BigtableTableName.class))).thenReturn(mockBulkMutation);
+    when(mockSession.createBulkMutationWrapper(any(BigtableTableName.class))).thenReturn(mockBulkMutation);
     when(mockSession.getDataRequestContext()).thenReturn(RequestContext.create("p", "i", "a"));
   }
 
@@ -109,18 +108,18 @@ public class TestBigtableBufferedMutator {
 
   @Test
   public void testPut() throws IOException{
-    when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class))).thenReturn(future);
+    when(mockBulkMutation.add(any(RowMutation.class))).thenReturn(future);
     BigtableBufferedMutator underTest = createMutator(new Configuration(false));
     underTest.mutate(SIMPLE_PUT);
-    verify(mockBulkMutation, times(1)).add(any(MutateRowsRequest.Entry.class));
+    verify(mockBulkMutation, times(1)).add(any(RowMutation.class));
   }
 
   @Test
   public void testDelete() throws IOException {
-    when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class))).thenReturn(future);
+    when(mockBulkMutation.add(any(RowMutation.class))).thenReturn(future);
     BigtableBufferedMutator underTest = createMutator(new Configuration(false));
     underTest.mutate(new Delete(EMPTY_BYTES));
-    verify(mockBulkMutation, times(1)).add(any(MutateRowsRequest.Entry.class));
+    verify(mockBulkMutation, times(1)).add(any(RowMutation.class));
   }
 
   @Test
@@ -145,8 +144,8 @@ public class TestBigtableBufferedMutator {
 
   @Test
   public void testInvalidPut() throws Exception {
-    when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class)))
-        .thenReturn(Futures.<MutateRowResponse> immediateFailedFuture(new RuntimeException()));
+    when(mockBulkMutation.add(any(RowMutation.class)))
+        .thenReturn(Futures.<Void> immediateFailedFuture(new RuntimeException()));
     BigtableBufferedMutator underTest = createMutator(new Configuration(false));
     underTest.mutate(SIMPLE_PUT);
     // Leave some time for the async worker to handle the request.
@@ -160,10 +159,10 @@ public class TestBigtableBufferedMutator {
   @Test
   public void testBulkSingleRequests() throws IOException, InterruptedException {
     Configuration config = new Configuration(false);
-    when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class))).thenReturn(future);
+    when(mockBulkMutation.add(any(RowMutation.class))).thenReturn(future);
     final BigtableBufferedMutator underTest = createMutator(config);
     underTest.mutate(SIMPLE_PUT);
-    verify(mockBulkMutation, times(1)).add(any(MutateRowsRequest.Entry.class));
+    verify(mockBulkMutation, times(1)).add(any(RowMutation.class));
     underTest.flush();
     verify(mockBulkMutation, times(1)).flush();
   }
@@ -171,13 +170,13 @@ public class TestBigtableBufferedMutator {
   @Test
   public void testBulkMultipleRequests() throws IOException, InterruptedException {
     Configuration config = new Configuration(false);
-    when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class))).thenReturn(future);
+    when(mockBulkMutation.add(any(RowMutation.class))).thenReturn(future);
     BigtableBufferedMutator underTest = createMutator(config);
     int count = 30;
     for (int i = 0; i < count; i++) {
       underTest.mutate(SIMPLE_PUT);
     }
-    verify(mockBulkMutation, times(count)).add(any(MutateRowsRequest.Entry.class));
+    verify(mockBulkMutation, times(count)).add(any(RowMutation.class));
     underTest.flush();
     verify(mockBulkMutation, times(1)).flush();
   }