metrics test refactoring

Author: surbhigarg92

google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetryBuiltInMetricsTracerTest.java

@@ -28,8 +28,6 @@
 import com.google.api.gax.longrunning.OperationTimedPollAlgorithm;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.tracing.ApiTracerFactory;
-import com.google.api.gax.tracing.MetricsTracerFactory;
-import com.google.api.gax.tracing.OpenTelemetryMetricsRecorder;
 import com.google.cloud.NoCredentials;
 import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
@@ -65,54 +63,37 @@
 
 @RunWith(JUnit4.class)
 public class OpenTelemetryBuiltInMetricsTracerTest extends AbstractNettyMockServerTest {
-
   private static final Statement SELECT_RANDOM = Statement.of("SELECT * FROM random");
-
   private static final Statement UPDATE_RANDOM = Statement.of("UPDATE random SET foo=1 WHERE id=1");
   private static InMemoryMetricReader metricReader;
-
-  private static OpenTelemetry openTelemetry;
-
-  private static Map<String, String> attributes;
-
-  private static Attributes expectedCommonBaseAttributes;
-  private static Attributes expectedCommonRequestAttributes;
-  private static ApiTracerFactory metricsTracerFactory;
+  private static Map<String, String> attributes =
+      BuiltInMetricsProvider.INSTANCE.createClientAttributes();
+  private static Attributes expectedCommonBaseAttributes =
+      Attributes.builder()
+          .put(BuiltInMetricsConstant.CLIENT_NAME_KEY, "spanner-java/")
+          .put(BuiltInMetricsConstant.CLIENT_UID_KEY, attributes.get("client_uid"))
+          .put(BuiltInMetricsConstant.INSTANCE_ID_KEY, "i")
+          .put(BuiltInMetricsConstant.DATABASE_KEY, "d")
+          .put(BuiltInMetricsConstant.DIRECT_PATH_ENABLED_KEY, "false")
+          .build();;
+  private static Attributes expectedCommonRequestAttributes =
+      Attributes.builder().put(BuiltInMetricsConstant.DIRECT_PATH_USED_KEY, "false").build();;
 
   private static final long MIN_LATENCY = 0;
 
   private DatabaseClient client;
 
-  @BeforeClass
-  public static void setup() {
+  public ApiTracerFactory createMetricsTracerFactory() {
     metricReader = InMemoryMetricReader.create();
 
-    BuiltInMetricsProvider provider = BuiltInMetricsProvider.INSTANCE;
-
     SdkMeterProviderBuilder meterProvider =
         SdkMeterProvider.builder().registerMetricReader(metricReader);
-
     BuiltInMetricsConstant.getAllViews().forEach(meterProvider::registerView);
+    OpenTelemetry openTelemetry =
+        OpenTelemetrySdk.builder().setMeterProvider(meterProvider.build()).build();
 
-    String client_name = "spanner-java/";
-    openTelemetry = OpenTelemetrySdk.builder().setMeterProvider(meterProvider.build()).build();
-    attributes = provider.createClientAttributes();
-
-    expectedCommonBaseAttributes =
-        Attributes.builder()
-            .put(BuiltInMetricsConstant.CLIENT_NAME_KEY, client_name)
-            .put(BuiltInMetricsConstant.CLIENT_UID_KEY, attributes.get("client_uid"))
-            .put(BuiltInMetricsConstant.INSTANCE_ID_KEY, "i")
-            .put(BuiltInMetricsConstant.DATABASE_KEY, "d")
-            .put(BuiltInMetricsConstant.DIRECT_PATH_ENABLED_KEY, "false")
-            .build();
-
-    expectedCommonRequestAttributes =
-        Attributes.builder().put(BuiltInMetricsConstant.DIRECT_PATH_USED_KEY, "false").build();
-    metricsTracerFactory =
-        new BuiltInMetricsTracerFactory(
-            new BuiltInMetricsRecorder(openTelemetry, BuiltInMetricsConstant.METER_NAME),
-            attributes);
+    return new BuiltInMetricsTracerFactory(
+        new BuiltInMetricsRecorder(openTelemetry, BuiltInMetricsConstant.METER_NAME), attributes);
   }
 
   @BeforeClass
@@ -123,9 +104,9 @@ public static void setupResults() {
   }
 
   @After
-  public void clearRequests() {
+  public void clearRequests() throws IOException {
     mockSpanner.clearRequests();
-    metricReader.forceFlush();
+    metricReader.close();
   }
 
   @Override
@@ -159,7 +140,7 @@ public void createSpannerInstance() {
             // Setting this to false so that Spanner Options does not register Metrics Tracer
             // factory again.
             .setBuiltInMetricsEnabled(false)
-            .setApiTracerFactory(metricsTracerFactory)
+            .setApiTracerFactory(createMetricsTracerFactory())
             .build()
             .getService();
     client = spanner.getDatabaseClient(DatabaseId.of("test-project", "i", "d"));
@@ -209,6 +190,8 @@ public void testMetricsSingleUseQuery() {
     long gfeLatencyValue = getAggregatedValue(gfeLatencyMetricData, expectedAttributes);
     assertEquals(fakeServerTiming.get(), gfeLatencyValue, 0);
 
+    assertFalse(
+        checkIfMetricExists(metricReader, BuiltInMetricsConstant.GFE_CONNECTIVITY_ERROR_NAME));
     assertFalse(checkIfMetricExists(metricReader, BuiltInMetricsConstant.AFE_LATENCIES_NAME));
     assertFalse(
         checkIfMetricExists(metricReader, BuiltInMetricsConstant.AFE_CONNECTIVITY_ERROR_NAME));
@@ -284,10 +267,6 @@ public void testNoNetworkConnection() {
               return null;
             });
 
-    ApiTracerFactory metricsTracerFactory =
-        new MetricsTracerFactory(
-            new OpenTelemetryMetricsRecorder(openTelemetry, BuiltInMetricsConstant.METER_NAME),
-            attributes);
     Spanner spanner =
         builder
             .setProjectId("test-project")
@@ -305,7 +284,7 @@ public void testNoNetworkConnection() {
             // Setting this to false so that Spanner Options does not register Metrics Tracer
             // factory again.
             .setBuiltInMetricsEnabled(false)
-            .setApiTracerFactory(metricsTracerFactory)
+            .setApiTracerFactory(createMetricsTracerFactory())
             .build()
             .getService();
     String instance = "i";
@@ -370,7 +349,7 @@ public void testNoServerTimingHeader() throws IOException, InterruptedException
             // Setting this to false so that Spanner Options does not register Metrics Tracer
             // factory again.
             .setBuiltInMetricsEnabled(false)
-            .setApiTracerFactory(metricsTracerFactory)
+            .setApiTracerFactory(createMetricsTracerFactory())
             .build()
             .getService();
     DatabaseClient databaseClientNoHeader =
@@ -391,6 +370,8 @@ public void testNoServerTimingHeader() throws IOException, InterruptedException
     MetricData gfeConnectivityMetricData =
         getMetricData(metricReader, BuiltInMetricsConstant.GFE_CONNECTIVITY_ERROR_NAME);
     assertThat(getAggregatedValue(gfeConnectivityMetricData, expectedAttributes)).isEqualTo(1);
+    assertFalse(checkIfMetricExists(metricReader, BuiltInMetricsConstant.AFE_LATENCIES_NAME));
+    assertFalse(checkIfMetricExists(metricReader, BuiltInMetricsConstant.GFE_LATENCIES_NAME));
     assertFalse(
         checkIfMetricExists(metricReader, BuiltInMetricsConstant.AFE_CONNECTIVITY_ERROR_NAME));
     spannerNoHeader.close();