fix tests

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java

@@ -32,6 +32,7 @@
 import com.google.cloud.bigtable.data.v2.stub.metrics.ErrorCountPerConnectionMetricTracker;
 import com.google.cloud.bigtable.data.v2.stub.metrics.MetricsProvider;
 import com.google.cloud.bigtable.data.v2.stub.metrics.NoopMetricsProvider;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.MoreObjects;
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.opentelemetry.GrpcOpenTelemetry;
@@ -54,7 +55,8 @@ public class BigtableClientContext {
   private static final Logger logger = Logger.getLogger(BigtableClientContext.class.getName());
 
   private static final String DEFAULT_DATA_JWT_AUDIENCE = "https://bigtable.googleapis.com/";
-  private static final String DATA_JWT_OVERRIDE_NAME = "bigtable.data-jwt-audience";
+
+  @VisibleForTesting static final String DATA_JWT_OVERRIDE_NAME = "bigtable.data-jwt-audience";
 
   @Nullable private final OpenTelemetry openTelemetry;
   @Nullable private final OpenTelemetrySdk internalOpenTelemetry;
@@ -231,12 +233,12 @@ private static OpenTelemetry getOpenTelemetryFromMetricsProvider(
   }
 
   private static void patchCredentials(EnhancedBigtableStubSettings.Builder settings)
-          throws IOException {
+      throws IOException {
     // Default jwt audience is always the service name unless it's override to
     // test / staging for testing
     String audience =
-            MoreObjects.firstNonNull(
-                    System.getProperty(DATA_JWT_OVERRIDE_NAME), DEFAULT_DATA_JWT_AUDIENCE);
+        MoreObjects.firstNonNull(
+            System.getProperty(DATA_JWT_OVERRIDE_NAME), DEFAULT_DATA_JWT_AUDIENCE);
 
     URI audienceUri = null;
     try {

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactoryTest.java

@@ -182,7 +182,10 @@ public void testNewClientsShareTransportChannel() throws Exception {
 
       // Make sure that only 1 instance is created by each provider
       Mockito.verify(transportChannelProvider, Mockito.times(1)).getTransportChannel();
-      Mockito.verify(credentialsProvider, Mockito.times(1)).getCredentials();
+      // getCredentials was called twice, in patchCredentials and when creating the fixed
+      // credentials
+      // in BigtableClientContext
+      Mockito.verify(credentialsProvider, Mockito.times(2)).getCredentials();
       Mockito.verify(executorProvider, Mockito.times(1)).getExecutor();
       Mockito.verify(watchdogProvider, Mockito.times(1)).getWatchdog();
     }
@@ -270,7 +273,9 @@ public void testCreateWithRefreshingChannel() throws Exception {
     factory.createForInstance("other-project", "other-instance");
 
     // Make sure that only 1 instance is created by each provider
-    Mockito.verify(credentialsProvider, Mockito.times(1)).getCredentials();
+    // getCredentials was called twice, in patchCredentials and when creating the fixed credentials
+    // in BigtableClientContext
+    Mockito.verify(credentialsProvider, Mockito.times(2)).getCredentials();
     Mockito.verify(executorProvider, Mockito.times(1)).getExecutor();
     Mockito.verify(watchdogProvider, Mockito.times(1)).getWatchdog();
 

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java

@@ -993,7 +993,6 @@ public void enableRetryInfoFalseValueTest() throws IOException {
     "appProfileId",
     "isRefreshingChannel",
     "primedTableIds",
-    "jwtAudienceMapping",
     "enableRoutingCookie",
     "enableRetryInfo",
     "enableSkipTrailers",

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java

@@ -93,7 +93,6 @@
 import com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory;
 import com.google.cloud.bigtable.data.v2.stub.sql.SqlServerStream;
 import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Queues;
 import com.google.common.io.BaseEncoding;
 import com.google.protobuf.ByteString;
@@ -211,9 +210,9 @@ public void testJwtAudience()
 
     // Create a stub with overridden audience
     String expectedAudience = "http://localaudience";
+    System.setProperty(BigtableClientContext.DATA_JWT_OVERRIDE_NAME, expectedAudience);
     EnhancedBigtableStubSettings settings =
         defaultSettings.toBuilder()
-            .setJwtAudienceMapping(ImmutableMap.of("localhost", expectedAudience))
             .setCredentialsProvider(FixedCredentialsProvider.create(jwtCreds))
             .build();
     try (EnhancedBigtableStub stub = EnhancedBigtableStub.create(settings)) {