Add Config and Context JUnit extensions (#11076)

feat(junit): Add config extension

feat(junit): Add context extension

feat(junit): Move table tests into junit-utils module

feat(core): Improve span links tests

Co-authored-by: bruce.bujon <bruce.bujon@datadoghq.com>

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

dd-trace-api/src/test/java/datadog/trace/api/DDTraceApiTableTestConverters.java

@@ -1,6 +1,6 @@
 package datadog.trace.api;
 
-import datadog.trace.test.util.TableTestTypeConverters;
+import datadog.trace.junit.utils.tabletest.TableTestTypeConverters;
 import org.tabletest.junit.TypeConverter;
 
 /** TableTest converters shared by dd-trace-api test classes for unparsable constants. */

dd-trace-core/src/main/java/datadog/trace/core/propagation/W3CHttpCodec.java

@@ -32,11 +32,6 @@
 class W3CHttpCodec {
   private static final Logger log = LoggerFactory.getLogger(W3CHttpCodec.class);
 
-  static final String TRACE_PARENT_KEY = "traceparent";
-  static final String TRACE_STATE_KEY = "tracestate";
-  static final String OT_BAGGAGE_PREFIX = "ot-baggage-";
-  private static final String E2E_START_KEY = OT_BAGGAGE_PREFIX + DDTags.TRACE_START_TIME;
-
   private static final int TRACE_PARENT_TID_START = 2 + 1;
   private static final int TRACE_PARENT_TID_END = TRACE_PARENT_TID_START + 32;
   private static final int TRACE_PARENT_SID_START = TRACE_PARENT_TID_END + 1;
@@ -45,6 +40,12 @@ class W3CHttpCodec {
   private static final int TRACE_PARENT_FLAGS_SAMPLED = 1;
   private static final int TRACE_PARENT_LENGTH = TRACE_PARENT_FLAGS_START + 2;
 
+  // Package-protected for testing
+  static final String TRACE_PARENT_KEY = "traceparent";
+  static final String TRACE_STATE_KEY = "tracestate";
+  static final String OT_BAGGAGE_PREFIX = "ot-baggage-";
+  static final String E2E_START_KEY = OT_BAGGAGE_PREFIX + DDTags.TRACE_START_TIME;
+
   private W3CHttpCodec() {
     // This class should not be created. This also makes code coverage checks happy.
   }

dd-trace-core/src/test/java/datadog/trace/core/DDSpanLinkTest.java

@@ -3,33 +3,35 @@
 import static datadog.trace.api.DDTags.SPAN_LINKS;
 import static datadog.trace.bootstrap.instrumentation.api.AgentSpanLink.DEFAULT_FLAGS;
 import static datadog.trace.bootstrap.instrumentation.api.AgentSpanLink.SAMPLED_FLAG;
+import static datadog.trace.bootstrap.instrumentation.api.ContextVisitors.stringValuesMap;
 import static datadog.trace.bootstrap.instrumentation.api.SpanAttributes.EMPTY;
 import static datadog.trace.core.propagation.HttpCodecTestHelper.TRACE_PARENT_KEY;
 import static datadog.trace.core.propagation.HttpCodecTestHelper.TRACE_STATE_KEY;
+import static datadog.trace.core.propagation.HttpCodecTestHelper.newW3cHttpCodecExtractor;
+import static java.util.Collections.emptyList;
+import static java.util.stream.Collectors.toList;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.type.CollectionType;
 import datadog.trace.api.Config;
 import datadog.trace.api.DDSpanId;
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.DynamicConfig;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
-import datadog.trace.bootstrap.instrumentation.api.ContextVisitors;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer.SpanBuilder;
 import datadog.trace.bootstrap.instrumentation.api.SpanAttributes;
 import datadog.trace.bootstrap.instrumentation.api.SpanLink;
 import datadog.trace.common.writer.ListWriter;
 import datadog.trace.core.propagation.ExtractedContext;
 import datadog.trace.core.propagation.HttpCodec;
-import datadog.trace.core.propagation.HttpCodecTestHelper;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -41,19 +43,21 @@ class DDSpanLinkTest extends DDCoreJavaSpecification {
 
   private static final int SPAN_LINK_TAG_MAX_LENGTH = 25_000;
   private static final ObjectMapper JSON_MAPPER = new ObjectMapper();
+  private static final CollectionType SPAN_LINK_LIST_TYPE =
+      JSON_MAPPER.getTypeFactory().constructCollectionType(List.class, SpanLinkAsTag.class);
 
   private ListWriter writer;
   private CoreTracer tracer;
 
   @BeforeEach
   void setup() {
-    writer = new ListWriter();
-    tracer = tracerBuilder().writer(writer).build();
+    this.writer = new ListWriter();
+    this.tracer = tracerBuilder().writer(this.writer).build();
   }
 
   @AfterEach
   void cleanupTest() {
-    writer.clear();
+    this.writer.clear();
   }
 
   @TableTest({
@@ -65,15 +69,15 @@ void createSpanLinkFromExtractedContext(boolean sampled, String traceFlags, Stri
     String traceId = "11223344556677889900aabbccddeeff";
     String spanId = "123456789abcdef0";
     String traceState = "dd=s:" + sample + ";o:some;t.dm:-4";
+
     Map<String, String> headers = new HashMap<>();
     headers.put(TRACE_PARENT_KEY.toUpperCase(), "00-" + traceId + "-" + spanId + "-" + traceFlags);
     headers.put(TRACE_STATE_KEY.toUpperCase(), traceState);
     HttpCodec.Extractor extractor =
-        HttpCodecTestHelper.W3CHttpCodecNewExtractor(
+        newW3cHttpCodecExtractor(
             Config.get(), () -> DynamicConfig.create().apply().captureTraceConfig());
 
-    ExtractedContext context =
-        (ExtractedContext) extractor.extract(headers, ContextVisitors.stringValuesMap());
+    ExtractedContext context = (ExtractedContext) extractor.extract(headers, stringValuesMap());
     SpanLink link = DDSpanLink.from(context);
 
     assertEquals(DDTraceId.fromHex(traceId), link.traceId());
@@ -85,19 +89,18 @@ void createSpanLinkFromExtractedContext(boolean sampled, String traceFlags, Stri
   @Test
   void testSpanLinkEncodingTagMaxSize() throws Exception {
     int tooManyLinkCount = 300;
-    AgentTracer.SpanBuilder builder = tracer.buildSpan("test", "operation");
+    SpanBuilder builder = tracer.buildSpan("test", "operation");
     List<SpanLink> links =
         IntStream.range(0, tooManyLinkCount)
             .mapToObj(this::createLink)
-            .collect(Collectors.toList());
-
-    for (SpanLink link : links) {
-      builder.withLink(link);
-    }
+            .peek(builder::withLink)
+            .collect(toList());
     AgentSpan span = builder.start();
     span.finish();
-    writer.waitForTraces(1);
-    String spanLinksTag = (String) writer.get(0).get(0).getTag(SPAN_LINKS);
+    this.writer.waitForTraces(1);
+
+    assertEquals(1, this.writer.get(0).size());
+    String spanLinksTag = (String) this.writer.get(0).get(0).getTag(SPAN_LINKS);
     List<SpanLinkAsTag> decodedSpanLinks = deserializeSpanLinks(spanLinksTag);
 
     assertTrue(spanLinksTag.length() < SPAN_LINK_TAG_MAX_LENGTH);
@@ -112,20 +115,18 @@ void testSpanLinkEncodingTagMaxSize() throws Exception {
 
   @Test
   void testSpanLinksEncodingOmittedEmptyKeys() throws Exception {
-    AgentTracer.SpanBuilder builder = tracer.buildSpan("test", "operation");
     SpanLink link =
         new DDSpanLink(
             DDTraceId.fromHex("11223344556677889900aabbccddeeff"),
             DDSpanId.fromHex("123456789abcdef0"),
             DEFAULT_FLAGS,
             "",
             EMPTY);
+    this.tracer.buildSpan("test", "operation").withLink(link).start().finish();
+    this.writer.waitForTraces(1);
 
-    AgentSpan span = builder.withLink(link).start();
-    span.finish();
-    writer.waitForTraces(1);
+    assertEquals(1, this.writer.get(0).size());
     String spanLinksTag = (String) writer.get(0).get(0).getTag(SPAN_LINKS);
-
     assertEquals(
         "[{\"span_id\":\"123456789abcdef0\",\"trace_id\":\"11223344556677889900aabbccddeeff\"}]",
         spanLinksTag);
@@ -140,7 +141,7 @@ void testSpanLinksEncodingOmittedEmptyKeys() throws Exception {
   })
   @ParameterizedTest(name = "add span link at any time [{index}]")
   void addSpanLinkAtAnyTime(boolean beforeStart, boolean afterStart) throws Exception {
-    AgentTracer.SpanBuilder builder = tracer.buildSpan("test", "operation");
+    SpanBuilder builder = this.tracer.buildSpan("test", "operation");
     List<SpanLink> links = new ArrayList<>();
 
     if (beforeStart) {
@@ -155,12 +156,11 @@ void addSpanLinkAtAnyTime(boolean beforeStart, boolean afterStart) throws Except
       links.add(link);
     }
     span.finish();
-    writer.waitForTraces(1);
-    String spanLinksTag = (String) writer.get(0).get(0).getTag(SPAN_LINKS);
-    List<SpanLinkAsTag> decodedSpanLinks =
-        spanLinksTag == null
-            ? java.util.Collections.emptyList()
-            : deserializeSpanLinks(spanLinksTag);
+    this.writer.waitForTraces(1);
+
+    assertEquals(1, this.writer.get(0).size());
+    String spanLinksTag = (String) this.writer.get(0).get(0).getTag(SPAN_LINKS);
+    List<SpanLinkAsTag> decodedSpanLinks = deserializeSpanLinks(spanLinksTag);
 
     int expectedLinkCount = (beforeStart ? 1 : 0) + (afterStart ? 1 : 0);
     assertEquals(expectedLinkCount, decodedSpanLinks.size());
@@ -171,14 +171,15 @@ void addSpanLinkAtAnyTime(boolean beforeStart, boolean afterStart) throws Except
 
   @Test
   void filterNullLinks() throws Exception {
-    AgentTracer.SpanBuilder builder = tracer.buildSpan("test", "operation");
+    SpanBuilder builder = this.tracer.buildSpan("test", "operation");
 
     AgentSpan span = builder.withLink(null).start();
     span.addLink(null);
     span.finish();
-    writer.waitForTraces(1);
-    String spanLinksTag = (String) writer.get(0).get(0).getTag(SPAN_LINKS);
+    this.writer.waitForTraces(1);
 
+    assertEquals(1, this.writer.get(0).size());
+    String spanLinksTag = (String) this.writer.get(0).get(0).getTag(SPAN_LINKS);
     assertNull(spanLinksTag);
   }
 
@@ -215,9 +216,10 @@ private void assertLink(SpanLink expected, SpanLinkAsTag actual) {
   }
 
   static List<SpanLinkAsTag> deserializeSpanLinks(String json) throws IOException {
-    return JSON_MAPPER.readValue(
-        json,
-        JSON_MAPPER.getTypeFactory().constructCollectionType(List.class, SpanLinkAsTag.class));
+    if (json == null) {
+      return emptyList();
+    }
+    return JSON_MAPPER.readValue(json, SPAN_LINK_LIST_TYPE);
   }
 
   static class SpanLinkAsTag {

dd-trace-core/src/test/java/datadog/trace/core/propagation/HttpCodecTestHelper.java

@@ -10,7 +10,7 @@ public class HttpCodecTestHelper {
   public static final String TRACE_PARENT_KEY = W3CHttpCodec.TRACE_PARENT_KEY;
   public static final String TRACE_STATE_KEY = W3CHttpCodec.TRACE_STATE_KEY;
 
-  public static HttpCodec.Extractor W3CHttpCodecNewExtractor(
+  public static HttpCodec.Extractor newW3cHttpCodecExtractor(
       Config config, Supplier<TraceConfig> traceConfigSupplier) {
     return W3CHttpCodec.newExtractor(config, traceConfigSupplier);
   }

settings.gradle.kts

@@ -159,6 +159,7 @@ include(
   ":dd-java-agent:testing",
   ":utils:config-utils",
   ":utils:container-utils",
+  ":utils:junit-utils",
   ":utils:filesystem-utils",
   ":utils:flare-utils",
   ":utils:logging-utils",

utils/junit-utils/build.gradle.kts

@@ -0,0 +1,15 @@
+plugins {
+  `java-library`
+}
+
+apply(from = "$rootDir/gradle/java.gradle")
+
+dependencies {
+  api(libs.bytebuddy)
+  api(libs.bytebuddyagent)
+  api(libs.forbiddenapis)
+  api(project(":components:environment"))
+
+  compileOnly(libs.junit.jupiter)
+  compileOnly(libs.tabletest)
+}

utils/junit-utils/src/main/java/datadog/trace/junit/utils/config/WithConfig.java

@@ -0,0 +1,49 @@
+package datadog.trace.junit.utils.config;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Repeatable;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+/**
+ * Declares a configuration override for a test. Can be placed on a test class (applies to all
+ * tests) or on individual test methods.
+ *
+ * <p>By default, injects a system property with the {@code dd.} prefix. Use {@code env = true} for
+ * environment variables (prefix {@code DD_}).
+ *
+ * <p>Examples:
+ *
+ * <pre>{@code
+ * @WithConfig(key = "service", value = "my_service")
+ * @WithConfig(key = "trace.resolver.enabled", value = "false")
+ * class MyTest extends DDJavaSpecification {
+ *
+ *   @Test
+ *   @WithConfig(key = "AGENT_HOST", value = "localhost", env = true)
+ *   void testWithEnv() { ... }
+ * }
+ * }</pre>
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.TYPE, ElementType.METHOD})
+@Repeatable(WithConfigs.class)
+@ExtendWith(WithConfigExtension.class)
+public @interface WithConfig {
+  /**
+   * Config key (e.g. {@code "trace.resolver.enabled"}). The {@code dd.}/{@code DD_} prefix is
+   * auto-added unless {@link #addPrefix()} is {@code false}.
+   */
+  String key();
+
+  /** Config value. */
+  String value();
+
+  /** If {@code true}, sets an environment variable instead of a system property. */
+  boolean env() default false;
+
+  /** If {@code false}, the key is used as-is without adding the {@code dd.}/{@code DD_} prefix. */
+  boolean addPrefix() default true;
+}