cleanup after test in case of failure

Author: ldetmer

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpcTest.java

@@ -1111,7 +1111,7 @@ public void testHttpTracingEnabled_JsonResponseException_SetsAttributes() throws
           "Invalid request", rpcSpan.getAttributes().get(BigQueryTelemetryTracer.STATUS_MESSAGE));
       assertNull(rpcSpan.getAttributes().get(BigQueryTelemetryTracer.EXCEPTION_TYPE));
     }
-    
+
     @Test
     public void testGetUriTemplateValueTelemetry() throws Exception {
       setMockResponse(
@@ -1145,37 +1145,38 @@ public void testGetUriTemplateValueTelemetry() throws Exception {
     public void testResendCountOnRetry() throws Exception {
       // Manually set attempt count to simulate being inside a retry loop
       com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(2);
+      try {
+        setMockResponse(
+            "{\"kind\":\"bigquery#dataset\",\"id\":\""
+                + PROJECT_ID
+                + ":"
+                + DATASET_ID
+                + "\",\"datasetReference\":{\"projectId\":\""
+                + PROJECT_ID
+                + "\",\"datasetId\":\""
+                + DATASET_ID
+                + "\"}}");
 
-      setMockResponse(
-          "{\"kind\":\"bigquery#dataset\",\"id\":\""
-              + PROJECT_ID
-              + ":"
-              + DATASET_ID
-              + "\",\"datasetReference\":{\"projectId\":\""
-              + PROJECT_ID
-              + "\",\"datasetId\":\""
-              + DATASET_ID
-              + "\"}}");
-
-      rpc.getDatasetSkipExceptionTranslation(PROJECT_ID, DATASET_ID, new HashMap<>());
-
-      List<io.opentelemetry.sdk.trace.data.SpanData> spans = spanExporter.getFinishedSpanItems();
-      assertThat(spans).isNotEmpty();
-      io.opentelemetry.sdk.trace.data.SpanData rpcSpan =
-          spans.stream()
-              .filter(s -> s.getName().equals("com.google.cloud.bigquery.BigQueryRpc.getDataset"))
-              .findFirst()
-              .orElse(null);
-      assertNotNull(rpcSpan);
-      assertEquals(
-          2L,
-          rpcSpan
-              .getAttributes()
-              .get(
-                  com.google.cloud.bigquery.telemetry.HttpTracingRequestInitializer
-                      .HTTP_REQUEST_RESEND_COUNT));
+        rpc.getDatasetSkipExceptionTranslation(PROJECT_ID, DATASET_ID, new HashMap<>());
 
-      com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(0);
+        List<io.opentelemetry.sdk.trace.data.SpanData> spans = spanExporter.getFinishedSpanItems();
+        assertThat(spans).isNotEmpty();
+        io.opentelemetry.sdk.trace.data.SpanData rpcSpan =
+            spans.stream()
+                .filter(s -> s.getName().equals("com.google.cloud.bigquery.BigQueryRpc.getDataset"))
+                .findFirst()
+                .orElse(null);
+        assertNotNull(rpcSpan);
+        assertEquals(
+            2L,
+            rpcSpan
+                .getAttributes()
+                .get(
+                    com.google.cloud.bigquery.telemetry.HttpTracingRequestInitializer
+                        .HTTP_REQUEST_RESEND_COUNT));
+      } finally {
+        com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(0);
+      }
     }
   }
 
@@ -1285,36 +1286,37 @@ public void testHttpTracingDisabled_GoogleJsonResponseException_DoesNotSetAttrib
     public void testResendCountNotSetWhenDisabled() throws Exception {
       // Manually set attempt count to simulate being inside a retry loop
       com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(2);
+      try {
+        setMockResponse(
+            "{\"kind\":\"bigquery#dataset\",\"id\":\""
+                + PROJECT_ID
+                + ":"
+                + DATASET_ID
+                + "\",\"datasetReference\":{\"projectId\":\""
+                + PROJECT_ID
+                + "\",\"datasetId\":\""
+                + DATASET_ID
+                + "\"}}");
 
-      setMockResponse(
-          "{\"kind\":\"bigquery#dataset\",\"id\":\""
-              + PROJECT_ID
-              + ":"
-              + DATASET_ID
-              + "\",\"datasetReference\":{\"projectId\":\""
-              + PROJECT_ID
-              + "\",\"datasetId\":\""
-              + DATASET_ID
-              + "\"}}");
-
-      rpc.getDatasetSkipExceptionTranslation(PROJECT_ID, DATASET_ID, new HashMap<>());
-
-      List<io.opentelemetry.sdk.trace.data.SpanData> spans = spanExporter.getFinishedSpanItems();
-      assertThat(spans).isNotEmpty();
-      io.opentelemetry.sdk.trace.data.SpanData rpcSpan =
-          spans.stream()
-              .filter(s -> s.getName().equals("com.google.cloud.bigquery.BigQueryRpc.getDataset"))
-              .findFirst()
-              .orElse(null);
-      assertNotNull(rpcSpan);
-      assertNull(
-          rpcSpan
-              .getAttributes()
-              .get(
-                  com.google.cloud.bigquery.telemetry.HttpTracingRequestInitializer
-                      .HTTP_REQUEST_RESEND_COUNT));
+        rpc.getDatasetSkipExceptionTranslation(PROJECT_ID, DATASET_ID, new HashMap<>());
 
-      com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(0);
+        List<io.opentelemetry.sdk.trace.data.SpanData> spans = spanExporter.getFinishedSpanItems();
+        assertThat(spans).isNotEmpty();
+        io.opentelemetry.sdk.trace.data.SpanData rpcSpan =
+            spans.stream()
+                .filter(s -> s.getName().equals("com.google.cloud.bigquery.BigQueryRpc.getDataset"))
+                .findFirst()
+                .orElse(null);
+        assertNotNull(rpcSpan);
+        assertNull(
+            rpcSpan
+                .getAttributes()
+                .get(
+                    com.google.cloud.bigquery.telemetry.HttpTracingRequestInitializer
+                        .HTTP_REQUEST_RESEND_COUNT));
+      } finally {
+        com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(0);
+      }
     }
   }
 

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/telemetry/HttpTracingRequestInitializerTest.java

@@ -338,22 +338,24 @@ public void testAddResponseBodySizeToSpan_NullLength() throws IOException {
   @Test
   public void testResendCountIsSetFromBigQueryRetryAlgorithm() throws IOException {
     com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(3);
-    HttpTransport transport = createTransport();
-    HttpRequest request = buildGetRequest(transport, initializer, BASE_URL);
-
-    HttpResponse response = request.execute();
-    response.disconnect();
+    try {
+      HttpTransport transport = createTransport();
+      HttpRequest request = buildGetRequest(transport, initializer, BASE_URL);
 
-    spanScope.close();
-    parentSpan.end();
+      HttpResponse response = request.execute();
+      response.disconnect();
 
-    List<SpanData> spans = spanExporter.getFinishedSpanItems();
-    assertEquals(1, spans.size());
-    SpanData span = spans.get(0);
-    assertEquals(
-        3L, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_RESEND_COUNT));
+      spanScope.close();
+      parentSpan.end();
 
-    com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(0);
+      List<SpanData> spans = spanExporter.getFinishedSpanItems();
+      assertEquals(1, spans.size());
+      SpanData span = spans.get(0);
+      assertEquals(
+          3L, span.getAttributes().get(HttpTracingRequestInitializer.HTTP_REQUEST_RESEND_COUNT));
+    } finally {
+      com.google.cloud.bigquery.BigQueryRetryAlgorithm.setCurrentAttempt(0);
+    }
   }
 
   private static HttpTransport createTransport() {