8484import com .google .cloud .bigtable .data .v2 .models .TableId ;
8585import com .google .cloud .bigtable .data .v2 .models .sql .ResultSetMetadata ;
8686import com .google .cloud .bigtable .data .v2 .models .sql .Statement ;
87+ import com .google .cloud .bigtable .data .v2 .stub .metrics .NoopMetricsProvider ;
8788import com .google .cloud .bigtable .data .v2 .stub .sql .ExecuteQueryCallable ;
8889import com .google .cloud .bigtable .data .v2 .stub .sql .SqlServerStream ;
8990import com .google .common .collect .ImmutableMap ;
9798import io .grpc .CallOptions ;
9899import io .grpc .Context ;
99100import io .grpc .Deadline ;
100- import io .grpc .ManagedChannel ;
101101import io .grpc .ManagedChannelBuilder ;
102102import io .grpc .Metadata ;
103103import io .grpc .Metadata .Key ;
@@ -172,6 +172,7 @@ public void setUp() throws IOException, IllegalAccessException, InstantiationExc
172172 .setInstanceId (INSTANCE_ID )
173173 .setAppProfileId (APP_PROFILE_ID )
174174 .setCredentialsProvider (NoCredentialsProvider .create ())
175+ .setMetricsProvider (NoopMetricsProvider .INSTANCE )
175176 .build ()
176177 .getStubSettings ();
177178
@@ -187,9 +188,6 @@ public void tearDown() {
187188 @ Test
188189 public void testJwtAudience ()
189190 throws InterruptedException , IOException , NoSuchAlgorithmException , ExecutionException {
190- // close default stub - need to create custom one
191- enhancedBigtableStub .close ();
192-
193191 // Create fake jwt creds
194192 KeyPairGenerator keyGen = KeyPairGenerator .getInstance ("RSA" );
195193 KeyPair keyPair = keyGen .genKeyPair ();
@@ -210,9 +208,10 @@ public void testJwtAudience()
210208 .setJwtAudienceMapping (ImmutableMap .of ("localhost" , expectedAudience ))
211209 .setCredentialsProvider (FixedCredentialsProvider .create (jwtCreds ))
212210 .build ();
213- enhancedBigtableStub = EnhancedBigtableStub .create (settings );
211+ try (EnhancedBigtableStub stub = EnhancedBigtableStub .create (settings )) {
212+ stub .readRowCallable ().futureCall (Query .create ("fake-table" )).get ();
213+ }
214214 // Send rpc and grab the credentials sent
215- enhancedBigtableStub .readRowCallable ().futureCall (Query .create ("fake-table" )).get ();
216215 Metadata metadata = metadataInterceptor .headers .take ();
217216
218217 String authValue = metadata .get (Key .of ("Authorization" , Metadata .ASCII_STRING_MARSHALLER ));
@@ -226,9 +225,6 @@ public void testJwtAudience()
226225 @ Test
227226 public void testBatchJwtAudience ()
228227 throws InterruptedException , IOException , NoSuchAlgorithmException , ExecutionException {
229- // close default stub - need to create custom one
230- enhancedBigtableStub .close ();
231-
232228 // Create fake jwt creds
233229 KeyPairGenerator keyGen = KeyPairGenerator .getInstance ("RSA" );
234230 KeyPair keyPair = keyGen .genKeyPair ();
@@ -241,31 +237,30 @@ public void testBatchJwtAudience()
241237 .setPrivateKeyId ("fake-private-key" )
242238 .build ();
243239
244- // Create a fixed channel that will ignore the default endpoint and connect to the emulator
245- ManagedChannel emulatorChannel =
246- ManagedChannelBuilder .forAddress ("localhost" , server .getPort ()).usePlaintext ().build ();
240+ EnhancedBigtableStubSettings settings =
241+ EnhancedBigtableStubSettings .newBuilder ()
242+ .setProjectId ("fake-project" )
243+ .setInstanceId ("fake-instance" )
244+ .setEndpoint ("batch-bigtable.googleapis.com:443" )
245+ .setCredentialsProvider (FixedCredentialsProvider .create (jwtCreds ))
246+ .setMetricsProvider (NoopMetricsProvider .INSTANCE )
247+ // Use a fixed channel that will ignore the default endpoint and connect to the emulator
248+ .setTransportChannelProvider (
249+ FixedTransportChannelProvider .create (
250+ GrpcTransportChannel .create (
251+ ManagedChannelBuilder .forAddress ("localhost" , server .getPort ())
252+ .usePlaintext ()
253+ .build ())))
254+ // Channel refreshing doesn't work with FixedTransportChannelProvider. Disable it for
255+ // the test
256+ .setRefreshingChannel (false )
257+ .build ();
247258
248259 Metadata metadata ;
249- try {
250- EnhancedBigtableStubSettings settings =
251- EnhancedBigtableStubSettings .newBuilder ()
252- .setProjectId ("fake-project" )
253- .setInstanceId ("fake-instance" )
254- .setEndpoint ("batch-bigtable.googleapis.com:443" )
255- .setCredentialsProvider (FixedCredentialsProvider .create (jwtCreds ))
256- .setTransportChannelProvider (
257- FixedTransportChannelProvider .create (
258- GrpcTransportChannel .create (emulatorChannel )))
259- // Channel refreshing doesn't work with FixedTransportChannelProvider. Disable it for
260- // the test
261- .setRefreshingChannel (false )
262- .build ();
263- enhancedBigtableStub = EnhancedBigtableStub .create (settings );
260+ try (EnhancedBigtableStub stub = EnhancedBigtableStub .create (settings )) {
264261 // Send rpc and grab the credentials sent
265- enhancedBigtableStub .readRowCallable ().futureCall (Query .create ("fake-table" )).get ();
262+ stub .readRowCallable ().futureCall (Query .create ("fake-table" )).get ();
266263 metadata = metadataInterceptor .headers .take ();
267- } finally {
268- emulatorChannel .shutdown ();
269264 }
270265
271266 String authValue = metadata .get (Key .of ("Authorization" , Metadata .ASCII_STRING_MARSHALLER ));
@@ -278,7 +273,6 @@ public void testBatchJwtAudience()
278273
279274 @ Test
280275 public void testFeatureFlags () throws InterruptedException , IOException , ExecutionException {
281-
282276 enhancedBigtableStub .readRowCallable ().futureCall (Query .create ("fake-table" )).get ();
283277 Metadata metadata = metadataInterceptor .headers .take ();
284278
0 commit comments