PR comment 2

Author: jinseopkim0

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

@@ -76,6 +76,17 @@ protected BigtableTableAdminClientV2(BigtableTableAdminStub stub) {
     this.optimizeRestoredTableOperationBaseCallable = null;
   }
 
+  @com.google.common.annotations.VisibleForTesting
+  BigtableTableAdminClientV2(
+      BigtableTableAdminStub stub,
+      AwaitConsistencyCallable awaitConsistencyCallable,
+      OperationCallable<Void, Empty, OptimizeRestoredTableMetadata>
+          optimizeRestoredTableOperationBaseCallable) {
+    super(stub);
+    this.awaitConsistencyCallable = awaitConsistencyCallable;
+    this.optimizeRestoredTableOperationBaseCallable = optimizeRestoredTableOperationBaseCallable;
+  }
+
   private AwaitConsistencyCallable createAwaitConsistencyCallable(
       BigtableTableAdminStubSettings settings) throws IOException {
     ClientContext clientContext = ClientContext.create(settings);
@@ -324,12 +335,8 @@ private UnaryCallable<ConsistencyRequest, Void> getAwaitConsistencyCallable() {
     if (optimizeRestoredTableOperationBaseCallable != null) {
       return optimizeRestoredTableOperationBaseCallable;
     }
-    // Fallback for tests or stub-based initialization
-    if (getStub()
-        instanceof com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub) {
-      return ((com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub) getStub())
-          .awaitOptimizeRestoredTableCallable();
-    }
-    throw new IllegalStateException("OptimizeRestoredTableCallable not initialized.");
+    throw new IllegalStateException(
+        "OptimizeRestoredTableCallable not initialized. BigtableTableAdminClientV2 must be "
+            + "initialized via settings to use this functionality.");
   }
 }

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

@@ -26,6 +26,8 @@
 import com.google.cloud.bigtable.admin.v2.models.ConsistencyRequest;
 import com.google.cloud.bigtable.admin.v2.models.OptimizeRestoredTableOperationToken;
 import com.google.cloud.bigtable.admin.v2.models.RestoredTableResult;
+import com.google.cloud.bigtable.admin.v2.stub.AwaitConsistencyCallable;
+import com.google.cloud.bigtable.admin.v2.stub.BigtableTableAdminStub;
 import com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub;
 import com.google.protobuf.Empty;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -47,9 +49,9 @@ public class BigtableTableAdminClientV2Test {
   private static final String TABLE_NAME =
       "projects/my-project/instances/my-instance/tables/my-table";
 
-  @Mock private EnhancedBigtableTableAdminStub mockStub;
+  @Mock private BigtableTableAdminStub mockStub;
 
-  @Mock private UnaryCallable<ConsistencyRequest, Void> mockAwaitConsistencyCallable;
+  @Mock private AwaitConsistencyCallable mockAwaitConsistencyCallable;
 
   @Mock
   private OperationCallable<Void, Empty, OptimizeRestoredTableMetadata>
@@ -59,13 +61,14 @@ public class BigtableTableAdminClientV2Test {
 
   @Before
   public void setUp() {
-    client = BigtableTableAdminClientV2.createClient(mockStub);
+    client =
+        new BigtableTableAdminClientV2(
+            mockStub, mockAwaitConsistencyCallable, mockOptimizeRestoredTableCallable);
   }
 
   @Test
   public void testWaitForConsistencyWithToken() {
     // Setup
-    Mockito.when(mockStub.awaitConsistencyCallable()).thenReturn(mockAwaitConsistencyCallable);
 
     String token = "my-token";
     ConsistencyRequest expectedRequest =
@@ -91,8 +94,6 @@ public void testWaitForConsistencyWithToken() {
   @Test
   public void testAwaitOptimizeRestoredTable() throws Exception {
     // Setup
-    Mockito.when(mockStub.awaitOptimizeRestoredTableCallable())
-        .thenReturn(mockOptimizeRestoredTableCallable);
 
     String optimizeToken = "my-optimization-token";