feat(junit): Add context extension

Author: PerfectSlayer

dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/AbstractInstrumentationTest.java

@@ -20,7 +20,7 @@
 import datadog.trace.core.DDSpan;
 import datadog.trace.core.PendingTrace;
 import datadog.trace.core.TraceCollector;
-import de.thetaphi.forbiddenapis.SuppressForbidden;
+import datadog.trace.junit.utils.context.AllowContextTestingExtension;
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.Instrumentation;
 import java.util.List;
@@ -31,7 +31,6 @@
 import java.util.function.Predicate;
 import net.bytebuddy.agent.ByteBuddyAgent;
 import org.junit.jupiter.api.AfterEach;
-import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.opentest4j.AssertionFailedError;
@@ -42,7 +41,7 @@
  * current implementation is inspired and kept close to it Groovy / Spock counterpart, the {@code
  * InstrumentationSpecification}.
  */
-@ExtendWith(TestClassShadowingExtension.class)
+@ExtendWith({TestClassShadowingExtension.class, AllowContextTestingExtension.class})
 public abstract class AbstractInstrumentationTest {
   static final Instrumentation INSTRUMENTATION = ByteBuddyAgent.getInstrumentation();
 
@@ -55,19 +54,6 @@ public abstract class AbstractInstrumentationTest {
   protected ClassFileTransformer activeTransformer;
   protected ClassFileTransformerListener transformerLister;
 
-  @SuppressForbidden // Class.forName() used to dynamically configure context if present
-  @BeforeAll
-  static void allowContextTesting() {
-    // Allow re-registration of context managers so each test can use a fresh tracer.
-    // This mirrors DDSpecification.allowContextTesting() for the Spock test framework.
-    try {
-      Class.forName("datadog.context.ContextManager").getMethod("allowTesting").invoke(null);
-      Class.forName("datadog.context.ContextBinder").getMethod("allowTesting").invoke(null);
-    } catch (Throwable ignore) {
-      // don't block testing if context types aren't available
-    }
-  }
-
   @BeforeEach
   public void init() {
     // If this fails, it's likely the result of another test loading Config before it can be

utils/junit-utils/src/main/java/datadog/trace/junit/utils/context/AllowContextTesting.java

@@ -0,0 +1,16 @@
+package datadog.trace.junit.utils.context;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+/**
+ * Enables context testing by allowing re-registration of {@code ContextManager} and {@code
+ * ContextBinder} singletons.
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.TYPE)
+@ExtendWith(AllowContextTestingExtension.class)
+public @interface AllowContextTesting {}

utils/junit-utils/src/main/java/datadog/trace/junit/utils/context/AllowContextTestingExtension.java

@@ -0,0 +1,27 @@
+package datadog.trace.junit.utils.context;
+
+import de.thetaphi.forbiddenapis.SuppressForbidden;
+import org.junit.jupiter.api.extension.BeforeAllCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+
+/**
+ * JUnit 5 extension that allows re-registration of context managers so each test can use a fresh
+ * tracer. This is needed because {@code ContextManager} and {@code ContextBinder} are singletons
+ * that normally reject re-registration.
+ *
+ * <p>Auto-registered when using {@link AllowContextTesting}. Can also be used explicitly via
+ * {@code @ExtendWith(AllowContextTestingExtension.class)}.
+ */
+@SuppressForbidden
+public class AllowContextTestingExtension implements BeforeAllCallback {
+
+  @Override
+  public void beforeAll(ExtensionContext context) {
+    try {
+      Class.forName("datadog.context.ContextManager").getMethod("allowTesting").invoke(null);
+      Class.forName("datadog.context.ContextBinder").getMethod("allowTesting").invoke(null);
+    } catch (Throwable ignore) {
+      // don't block testing if context types aren't available
+    }
+  }
+}

utils/test-utils/src/main/java/datadog/trace/test/util/DDJavaSpecification.java

@@ -4,8 +4,8 @@
 
 import datadog.environment.EnvironmentVariables;
 import datadog.trace.junit.utils.config.WithConfigExtension;
+import datadog.trace.junit.utils.context.AllowContextTestingExtension;
 import de.thetaphi.forbiddenapis.SuppressForbidden;
-import java.lang.reflect.Method;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
@@ -16,32 +16,23 @@
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.extension.ExtendWith;
 
-@ExtendWith(WithConfigExtension.class)
+@ExtendWith({WithConfigExtension.class, AllowContextTestingExtension.class})
 @SuppressForbidden
 public class DDJavaSpecification {
 
   private static final long CHECK_TIMEOUT_MS = 3000;
 
-  static final String CONTEXT_BINDER = "datadog.context.ContextBinder";
-  static final String CONTEXT_MANAGER = "datadog.context.ContextManager";
-
-  private static Boolean contextTestingAllowed;
-
   protected boolean assertThreadsEachCleanup = true;
   private static volatile boolean ignoreThreadCleanup;
 
   @BeforeAll
   static void beforeAll() {
-    allowContextTesting();
     assertTrue(
         EnvironmentVariables.getAll().entrySet().stream()
             .noneMatch(e -> e.getKey().startsWith("DD_")));
     assertTrue(
         systemPropertiesExceptAllowed().entrySet().stream()
             .noneMatch(e -> e.getKey().toString().startsWith("dd.")));
-    assertTrue(
-        contextTestingAllowed,
-        "Context not ready for testing. Ensure all test classes extend DDJavaSpecification");
 
     if (getDDThreads().isEmpty()) {
       ignoreThreadCleanup = false;
@@ -52,27 +43,6 @@ static void beforeAll() {
     }
   }
 
-  static void allowContextTesting() {
-    if (contextTestingAllowed == null) {
-      try {
-        Class<?> binderClass = Class.forName(CONTEXT_BINDER);
-        Method binderAllowTesting = binderClass.getDeclaredMethod("allowTesting");
-        binderAllowTesting.setAccessible(true);
-        Class<?> managerClass = Class.forName(CONTEXT_MANAGER);
-        Method managerAllowTesting = managerClass.getDeclaredMethod("allowTesting");
-        managerAllowTesting.setAccessible(true);
-        contextTestingAllowed =
-            (Boolean) binderAllowTesting.invoke(null) && (Boolean) managerAllowTesting.invoke(null);
-      } catch (ClassNotFoundException e) {
-        // don't block testing if these types aren't found (project doesn't use context API)
-        contextTestingAllowed =
-            CONTEXT_BINDER.equals(e.getMessage()) || CONTEXT_MANAGER.equals(e.getMessage());
-      } catch (Throwable ignore) {
-        contextTestingAllowed = false;
-      }
-    }
-  }
-
   @AfterAll
   static void afterAll() {
     assertTrue(
@@ -135,5 +105,4 @@ static void checkThreads() {
       System.out.println(names);
     }
   }
-
 }