fix tests

Author: igorbernstein2

google-cloud-clients/google-cloud-bigtable-admin/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClient.java

@@ -480,7 +480,7 @@ public ApiFuture<ConsistencyToken> generateConsistencyTokenAsync(final String ta
         new ApiFunction<GenerateConsistencyTokenResponse, ConsistencyToken>() {
           @Override
           public ConsistencyToken apply(GenerateConsistencyTokenResponse proto) {
-            TableName tableName = TableName.of(instanceName.getProject(), instanceName.getProject(), tableId);
+            TableName tableName = TableName.of(instanceName.getProject(), instanceName.getInstance(), tableId);
             return ConsistencyToken.of(tableName, proto.getConsistencyToken());
           }
         });

google-cloud-clients/google-cloud-bigtable-admin/src/main/java/com/google/cloud/bigtable/admin/v2/models/ConsistencyToken.java

@@ -16,12 +16,12 @@
 package com.google.cloud.bigtable.admin.v2.models;
 
 import com.google.api.core.InternalApi;
+import com.google.api.core.InternalExtensionOnly;
+import com.google.auto.value.AutoValue;
 import com.google.bigtable.admin.v2.CheckConsistencyRequest;
 import com.google.bigtable.admin.v2.GenerateConsistencyTokenResponse;
 import com.google.bigtable.admin.v2.InstanceName;
 import com.google.bigtable.admin.v2.TableName;
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 
 /**
@@ -30,38 +30,25 @@
  * <p>Cannot be created. They are obtained by invoking {@link
  * com.google.cloud.bigtable.admin.v2.BigtableTableAdminClient#generateConsistencyToken(String)}
  */
-public final class ConsistencyToken {
-  private final TableName tableName;
-  private final String token;
-
+@InternalExtensionOnly
+@AutoValue
+public abstract class ConsistencyToken {
   public static ConsistencyToken of(TableName tableName, String token) {
-    return new ConsistencyToken(tableName, token);
+    return new AutoValue_ConsistencyToken(tableName, token);
   }
 
-  public ConsistencyToken(TableName tableName, String token) {
-    this.tableName = tableName;
-    this.token = token;
-  }
+  abstract TableName getTableName();
+  abstract String getToken();
 
   @InternalApi
   public CheckConsistencyRequest toProto(InstanceName instanceName) {
     Preconditions.checkArgument(
-        instanceName.equals(InstanceName.of(tableName.getProject(), tableName.getInstance())),
+        instanceName.equals(InstanceName.of(getTableName().getProject(), getTableName().getInstance())),
         "Consistency tokens are only valid within a single instance.");
 
     return CheckConsistencyRequest.newBuilder()
-        .setName(tableName.toString())
-        .setConsistencyToken(token)
+        .setName(getTableName().toString())
+        .setConsistencyToken(getToken())
         .build();
   }
-
-  @VisibleForTesting
-  String getToken() {
-    return token;
-  }
-
-  @Override
-  public String toString() {
-    return MoreObjects.toStringHelper(this).add("token", token).toString();
-  }
 }

google-cloud-clients/google-cloud-bigtable-admin/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientTest.java

@@ -18,14 +18,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.mockito.Matchers.any;
-import java.util.List;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.ArgumentCaptor;
-import org.mockito.Mock;
-import org.mockito.Mockito;
-import org.mockito.runners.MockitoJUnitRunner;
+
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
 import com.google.api.gax.rpc.UnaryCallable;
@@ -41,12 +34,19 @@
 import com.google.bigtable.admin.v2.ListTablesResponse;
 import com.google.bigtable.admin.v2.Table;
 import com.google.bigtable.admin.v2.TableName;
+import com.google.cloud.bigtable.admin.v2.models.ConsistencyToken;
 import com.google.cloud.bigtable.admin.v2.models.CreateTableRequest;
 import com.google.cloud.bigtable.admin.v2.models.ModifyColumnFamiliesRequest;
-import com.google.cloud.bigtable.admin.v2.models.ConsistencyToken;
 import com.google.cloud.bigtable.admin.v2.stub.BigtableTableAdminStub;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.Empty;
+import java.util.List;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.runners.MockitoJUnitRunner;
 
 @RunWith(MockitoJUnitRunner.class)
 public class BigtableTableAdminClientTest {
@@ -239,51 +239,64 @@ public void dropRowRangeAsync() {
   }
 
   @Test
-  public void generateAndCheckConsistency() {
-    GenerateConsistencyTokenResponse genResp =
-        GenerateConsistencyTokenResponse.newBuilder().build();
-    Mockito.when(
-            mockGenerateConsistencyTokenCallable.call(
-                adminClient.composeGenerateConsistencyTokenRequest("tableId")))
-        .thenReturn(genResp);
-
-    ConsistencyToken consistencyToken = adminClient.generateConsistencyToken("tableId");
-    Mockito.verify(mockGenerateConsistencyTokenCallable)
-        .call(adminClient.composeGenerateConsistencyTokenRequest("tableId"));
-
-    ArgumentCaptor<CheckConsistencyRequest> requestCaptor =
-        ArgumentCaptor.forClass(CheckConsistencyRequest.class);
-    CheckConsistencyResponse consistencyResp = CheckConsistencyResponse.newBuilder().build();
-    Mockito.when(mockCheckConsistencyCallable.call(any(CheckConsistencyRequest.class)))
-        .thenReturn(consistencyResp);
-
-    adminClient.isConsistent(consistencyToken);
-    Mockito.verify(mockCheckConsistencyCallable).call(requestCaptor.capture());
+  public void testGenerateConsistencyToken() {
+    TableName tableName = TableName.of("[PROJECT]", "[INSTANCE]", "[TABLE]");
+
+    GenerateConsistencyTokenRequest expectedRequest = GenerateConsistencyTokenRequest.newBuilder()
+        .setName(tableName.toString())
+        .build();
+
+    GenerateConsistencyTokenResponse fakeResponse = GenerateConsistencyTokenResponse.newBuilder()
+        .setConsistencyToken("fake-token")
+        .build();
+
+    Mockito.when(mockGenerateConsistencyTokenCallable.futureCall(expectedRequest))
+        .thenReturn(ApiFutures.immediateFuture(fakeResponse));
+
+    ConsistencyToken actualToken = adminClient.generateConsistencyToken(tableName.getTable());
+
+    assertThat(actualToken).isEqualTo(ConsistencyToken.of(tableName, "fake-token"));
   }
 
   @Test
-  public void generateAndCheckConsistencyAsync() throws Exception {
-    ApiFuture<GenerateConsistencyTokenResponse> genResp =
-        ApiFutures.immediateFuture(GenerateConsistencyTokenResponse.newBuilder().build());
-    Mockito.when(
-            mockGenerateConsistencyTokenCallable.futureCall(
-                adminClient.composeGenerateConsistencyTokenRequest("tableId")))
-        .thenReturn(genResp);
-
-    ApiFuture<ConsistencyToken> consistencyTokenFuture =
-        adminClient.generateConsistencyTokenAsync("tableId");
-    Mockito.verify(mockGenerateConsistencyTokenCallable)
-        .futureCall(adminClient.composeGenerateConsistencyTokenRequest("tableId"));
-
-    ArgumentCaptor<CheckConsistencyRequest> requestCaptor =
-        ArgumentCaptor.forClass(CheckConsistencyRequest.class);
-    ApiFuture<CheckConsistencyResponse> consistencyResp =
-        ApiFutures.immediateFuture(CheckConsistencyResponse.newBuilder().build());
-    Mockito.when(mockCheckConsistencyCallable.futureCall(any(CheckConsistencyRequest.class)))
-        .thenReturn(consistencyResp);
-
-    adminClient.isConsistentAsync(consistencyTokenFuture.get());
-    Mockito.verify(mockCheckConsistencyCallable).futureCall(requestCaptor.capture());
+  public void testGenerateConsistencyTokenAsync() throws Exception {
+    TableName tableName = TableName.of("[PROJECT]", "[INSTANCE]", "[TABLE]");
+
+    GenerateConsistencyTokenRequest expectedRequest = GenerateConsistencyTokenRequest.newBuilder()
+        .setName(tableName.toString())
+        .build();
+
+    GenerateConsistencyTokenResponse fakeResponse = GenerateConsistencyTokenResponse.newBuilder()
+        .setConsistencyToken("fake-token")
+        .build();
+
+    Mockito.when(mockGenerateConsistencyTokenCallable.futureCall(expectedRequest))
+        .thenReturn(ApiFutures.immediateFuture(fakeResponse));
+
+    ApiFuture<ConsistencyToken> actualTokenFuture = adminClient.generateConsistencyTokenAsync("[TABLE]");
+
+    assertThat(actualTokenFuture.get()).isEqualTo(ConsistencyToken.of(tableName, "fake-token"));
+  }
+
+  @Test
+  public void testCheckConsistencyToken() {
+    TableName tableName = TableName.of("[PROJECT]", "[INSTANCE]", "[TABLE]");
+
+    CheckConsistencyRequest expectedRequest = CheckConsistencyRequest.newBuilder()
+        .setName(tableName.toString())
+        .setConsistencyToken("fake-token")
+        .build();
+
+    CheckConsistencyResponse fakeResponse = CheckConsistencyResponse.newBuilder()
+        .setConsistent(true)
+        .build();
+
+    Mockito.when(mockCheckConsistencyCallable.futureCall(expectedRequest))
+        .thenReturn(ApiFutures.immediateFuture(fakeResponse));
+
+    boolean result = adminClient.isConsistent(ConsistencyToken.of(tableName, "fake-token"));
+
+    assertThat(result).isTrue();
   }
 
   @Test