From aa5fef2bfe476b76dec199081dd76437e65aaf82 Mon Sep 17 00:00:00 2001
From: Mateusz Rzeszutek <mrzeszutek@splunk.com>
Date: Wed, 17 Feb 2021 13:22:14 +0100
Subject: [PATCH] Improve library tests' stability (#2300)

* Improve library tests' stability

* Code review follow-up
---
 .../v1_0/ApiGatewayProxyRequestTest.java      |  4 +--
 .../v1_0/ParentContextExtractorTest.java      |  2 ++
 .../testing/LibraryTestRunner.java            | 31 +++++++++++++------
 3 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/instrumentation/aws-lambda-1.0/library/src/test/java/io/opentelemetry/instrumentation/awslambda/v1_0/ApiGatewayProxyRequestTest.java b/instrumentation/aws-lambda-1.0/library/src/test/java/io/opentelemetry/instrumentation/awslambda/v1_0/ApiGatewayProxyRequestTest.java
index 222be30e9c..b663ce9838 100644
--- a/instrumentation/aws-lambda-1.0/library/src/test/java/io/opentelemetry/instrumentation/awslambda/v1_0/ApiGatewayProxyRequestTest.java
+++ b/instrumentation/aws-lambda-1.0/library/src/test/java/io/opentelemetry/instrumentation/awslambda/v1_0/ApiGatewayProxyRequestTest.java
@@ -21,12 +21,12 @@ import io.opentelemetry.extension.trace.propagation.B3Propagator;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 class ApiGatewayProxyRequestTest {
 
-  @AfterEach
+  @BeforeEach
   void resetOpenTelemetry() {
     GlobalOpenTelemetry.resetForTest();
   }
diff --git a/instrumentation/aws-lambda-1.0/library/src/test/java/io/opentelemetry/instrumentation/awslambda/v1_0/ParentContextExtractorTest.java b/instrumentation/aws-lambda-1.0/library/src/test/java/io/opentelemetry/instrumentation/awslambda/v1_0/ParentContextExtractorTest.java
index dfa2baa0a3..7bb0888cda 100644
--- a/instrumentation/aws-lambda-1.0/library/src/test/java/io/opentelemetry/instrumentation/awslambda/v1_0/ParentContextExtractorTest.java
+++ b/instrumentation/aws-lambda-1.0/library/src/test/java/io/opentelemetry/instrumentation/awslambda/v1_0/ParentContextExtractorTest.java
@@ -32,6 +32,8 @@ public class ParentContextExtractorTest {
 
   @BeforeClass
   public static void setUp() {
+    // reset opentelemetry global instance in case any of the previous tests set it
+    GlobalOpenTelemetry.resetForTest();
     GlobalOpenTelemetry.set(
         OpenTelemetry.getPropagating(ContextPropagators.create(B3Propagator.getInstance())));
   }
diff --git a/testing-common/src/main/java/io/opentelemetry/instrumentation/testing/LibraryTestRunner.java b/testing-common/src/main/java/io/opentelemetry/instrumentation/testing/LibraryTestRunner.java
index a110983046..8adc7414dd 100644
--- a/testing-common/src/main/java/io/opentelemetry/instrumentation/testing/LibraryTestRunner.java
+++ b/testing-common/src/main/java/io/opentelemetry/instrumentation/testing/LibraryTestRunner.java
@@ -5,6 +5,7 @@
 
 package io.opentelemetry.instrumentation.testing;
 
+import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.propagation.ContextPropagators;
@@ -27,20 +28,24 @@ import java.util.List;
  */
 public final class LibraryTestRunner implements InstrumentationTestRunner {
 
-  protected static final InMemorySpanExporter testExporter;
+  private static final OpenTelemetrySdk openTelemetry;
+  private static final InMemorySpanExporter testExporter;
   private static boolean forceFlushCalled;
   private static final LibraryTestRunner INSTANCE = new LibraryTestRunner();
 
   static {
+    GlobalOpenTelemetry.resetForTest();
+
     testExporter = InMemorySpanExporter.create();
-    OpenTelemetrySdk.builder()
-        .setTracerProvider(
-            SdkTracerProvider.builder()
-                .addSpanProcessor(new FlushTrackingSpanProcessor())
-                .addSpanProcessor(SimpleSpanProcessor.create(testExporter))
-                .build())
-        .setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
-        .buildAndRegisterGlobal();
+    openTelemetry =
+        OpenTelemetrySdk.builder()
+            .setTracerProvider(
+                SdkTracerProvider.builder()
+                    .addSpanProcessor(new FlushTrackingSpanProcessor())
+                    .addSpanProcessor(SimpleSpanProcessor.create(testExporter))
+                    .build())
+            .setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
+            .buildAndRegisterGlobal();
   }
 
   public static InstrumentationTestRunner instance() {
@@ -48,7 +53,13 @@ public final class LibraryTestRunner implements InstrumentationTestRunner {
   }
 
   @Override
-  public void beforeTestClass() {}
+  public void beforeTestClass() {
+    // just in case: if there was any test that modified the global instance, reset it
+    if (GlobalOpenTelemetry.get() != openTelemetry) {
+      GlobalOpenTelemetry.resetForTest();
+      GlobalOpenTelemetry.set(openTelemetry);
+    }
+  }
 
   @Override
   public void afterTestClass() {}