Improve instrumentation JUnit testing

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

@@ -14,7 +14,7 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.IdGenerationStrategy;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI;
 import datadog.trace.common.writer.ListWriter;
 import datadog.trace.core.CoreTracer;
 import datadog.trace.core.DDSpan;
@@ -30,48 +30,56 @@
 import java.util.function.Function;
 import java.util.function.Predicate;
 import net.bytebuddy.agent.ByteBuddyAgent;
+import org.junit.jupiter.api.AfterAll;
 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;
 
 /**
- * This class is an experimental base to run instrumentation tests using JUnit Jupiter. It is still
- * early development, and the overall API is expected to change to leverage its extension model. The
- * current implementation is inspired and kept close to it Groovy / Spock counterpart, the {@code
- * InstrumentationSpecification}.
+ * Base class for instrumentation tests using JUnit Jupiter.
+ *
+ * <p>It is still early development, and the overall API might change to leverage its extension
+ * model. The current implementation is inspired and kept close to its Groovy / Spock counterpart,
+ * the {@code InstrumentationSpecification}.
+ *
+ * <ul>
+ *   <li>{@code @BeforeAll}: Installs the agent and creates a shared tracer
+ *   <li>{@code @BeforeEach}: Flushes and resets the writer
+ *   <li>{@code @AfterEach}: Flushes the tracer
+ *   <li>{@code @AfterAll}: Closes the tracer and removes the agent transformer
+ * </ul>
  */
 @ExtendWith({TestClassShadowingExtension.class, AllowContextTestingExtension.class})
 public abstract class AbstractInstrumentationTest {
   static final Instrumentation INSTRUMENTATION = ByteBuddyAgent.getInstrumentation();
 
   static final long TIMEOUT_MILLIS = TimeUnit.SECONDS.toMillis(20);
 
-  protected AgentTracer.TracerAPI tracer;
+  protected static final InstrumentationTestConfig testConfig = new InstrumentationTestConfig();
 
-  protected ListWriter writer;
+  protected static TracerAPI tracer;
+  protected static ListWriter writer;
+  private static ClassFileTransformer activeTransformer;
+  private static ClassFileTransformerListener transformerListener;
 
-  protected ClassFileTransformer activeTransformer;
-  protected ClassFileTransformerListener transformerLister;
-
-  @BeforeEach
-  public void init() {
+  @BeforeAll
+  static void initAll() {
     // If this fails, it's likely the result of another test loading Config before it can be
     // injected into the bootstrap classpath.
-    // If one test extends AgentTestRunner in a module, all tests must extend
     assertNull(Config.class.getClassLoader(), "Config must load on the bootstrap classpath.");
 
-    // Initialize test tracer
-    this.writer = new ListWriter();
-    // Initialize test tracer
-    CoreTracer tracer =
+    // Create shared test writer and tracer
+    writer = new ListWriter();
+    CoreTracer coreTracer =
         CoreTracer.builder()
-            .writer(this.writer)
-            .idGenerationStrategy(IdGenerationStrategy.fromName(idGenerationStrategyName()))
-            .strictTraceWrites(useStrictTraceWrites())
+            .writer(writer)
+            .idGenerationStrategy(IdGenerationStrategy.fromName(testConfig.idGenerationStrategy))
+            .strictTraceWrites(testConfig.strictTraceWrites)
             .build();
-    TracerInstaller.forceInstallGlobalTracer(tracer);
-    this.tracer = tracer;
+    TracerInstaller.forceInstallGlobalTracer(coreTracer);
+    tracer = coreTracer;
 
     ClassInjector.enableClassInjection(INSTRUMENTATION);
 
@@ -85,33 +93,43 @@ public void init() {
             .iterator()
             .hasNext(),
         "No instrumentation found");
-    this.transformerLister = new ClassFileTransformerListener();
-    this.activeTransformer =
+    transformerListener = new ClassFileTransformerListener();
+    activeTransformer =
         AgentInstaller.installBytebuddyAgent(
-            INSTRUMENTATION, true, AgentInstaller.getEnabledSystems(), this.transformerLister);
-  }
-
-  protected String idGenerationStrategyName() {
-    return "SEQUENTIAL";
+            INSTRUMENTATION, true, AgentInstaller.getEnabledSystems(), transformerListener);
   }
 
-  private boolean useStrictTraceWrites() {
-    return true;
+  @BeforeEach
+  public void init() {
+    tracer.flush();
+    writer.start();
   }
 
   @AfterEach
   public void tearDown() {
-    this.tracer.close();
-    this.writer.close();
-    if (this.activeTransformer != null) {
-      INSTRUMENTATION.removeTransformer(this.activeTransformer);
-      this.activeTransformer = null;
-    }
+    tracer.flush();
+  }
 
-    // All cleanups should happen before these assertions.
+  @AfterAll
+  static void tearDownAll() {
+    if (tracer != null) {
+      tracer.close();
+      tracer = null;
+    }
+    if (writer != null) {
+      writer.close();
+      writer = null;
+    }
+    if (activeTransformer != null) {
+      INSTRUMENTATION.removeTransformer(activeTransformer);
+      activeTransformer = null;
+    }
+    // All cleanups should happen before this verify call.
     // If not, a failing assertion may prevent cleanup
-    this.transformerLister.verify();
-    this.transformerLister = null;
+    if (transformerListener != null) {
+      transformerListener.verify();
+      transformerListener = null;
+    }
   }
 
   /**
@@ -134,11 +152,11 @@ protected void assertTraces(
       TraceMatcher... matchers) {
     int expectedTraceCount = matchers.length;
     try {
-      this.writer.waitForTraces(expectedTraceCount);
+      writer.waitForTraces(expectedTraceCount);
     } catch (InterruptedException | TimeoutException e) {
       throw new AssertionFailedError("Timeout while waiting for traces", e);
     }
-    TraceAssertions.assertTraces(this.writer, options, matchers);
+    TraceAssertions.assertTraces(writer, options, matchers);
   }
 
   /**
@@ -149,7 +167,7 @@ protected void assertTraces(
    */
   protected void blockUntilTracesMatch(Predicate<List<List<DDSpan>>> predicate) {
     long deadline = System.currentTimeMillis() + TIMEOUT_MILLIS;
-    while (!predicate.test(this.writer)) {
+    while (!predicate.test(writer)) {
       if (System.currentTimeMillis() > deadline) {
         throw new RuntimeException(new TimeoutException("Timed out waiting for traces/spans."));
       }
@@ -161,8 +179,8 @@ protected void blockUntilTracesMatch(Predicate<List<List<DDSpan>>> predicate) {
     }
   }
 
-  protected void blockUntilChildSpansFinished(final int numberOfSpans) {
-    blockUntilChildSpansFinished(this.tracer.activeSpan(), numberOfSpans);
+  protected void blockUntilChildSpansFinished(int numberOfSpans) {
+    blockUntilChildSpansFinished(tracer.activeSpan(), numberOfSpans);
   }
 
   static void blockUntilChildSpansFinished(AgentSpan span, int numberOfSpans) {
@@ -190,4 +208,20 @@ static void blockUntilChildSpansFinished(AgentSpan span, int numberOfSpans) {
       }
     }
   }
+
+  /** Configuration for {@link AbstractInstrumentationTest}. */
+  protected static class InstrumentationTestConfig {
+    private String idGenerationStrategy = "SEQUENTIAL";
+    private boolean strictTraceWrites = true;
+
+    public InstrumentationTestConfig idGenerationStrategy(String strategy) {
+      this.idGenerationStrategy = strategy;
+      return this;
+    }
+
+    public InstrumentationTestConfig strictTraceWrites(boolean strict) {
+      this.strictTraceWrites = strict;
+      return this;
+    }
+  }
 }

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

@@ -1,9 +1,10 @@
 package datadog.trace.agent.test.assertions;
 
+import static org.junit.jupiter.api.AssertionFailureBuilder.assertionFailure;
+
 import java.util.Optional;
 import java.util.function.Predicate;
 import java.util.regex.Pattern;
-import org.opentest4j.AssertionFailedError;
 
 /** This class is a utility class to create generic matchers. */
 public final class Matchers {
@@ -103,12 +104,11 @@ public static <T> Matcher<T> any() {
   static <T> void assertValue(Matcher<T> matcher, T value, String message) {
     if (matcher != null && !matcher.test(value)) {
       Optional<T> expected = matcher.expected();
-      if (expected.isPresent()) {
-        throw new AssertionFailedError(
-            message + ". " + matcher.failureReason(), expected.get(), value);
-      } else {
-        throw new AssertionFailedError(message + ": " + value + ". " + matcher.failureReason());
-      }
+      assertionFailure()
+          .message(message + ". " + matcher.failureReason())
+          .expected(expected.orElse(null))
+          .actual(value)
+          .buildAndThrow();
     }
   }
 }

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

@@ -10,6 +10,7 @@
 import static datadog.trace.agent.test.assertions.Matchers.validates;
 import static datadog.trace.core.DDSpanAccessor.spanLinks;
 import static java.time.Duration.ofNanos;
+import static org.junit.jupiter.api.AssertionFailureBuilder.assertionFailure;
 
 import datadog.trace.api.TagMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
@@ -322,7 +323,7 @@ private void assertSpanTags(TagMap tags) {
           if (matcher == null) {
             uncheckedTagNames.add(key);
           } else {
-            assertValue(matcher, value, "Unexpected " + key + " tag value.");
+            assertValue(matcher, value, "Unexpected " + key + " tag value");
           }
         });
     // Remove matchers that accept missing tags
@@ -344,10 +345,18 @@ private void assertSpanTags(TagMap tags) {
    * It might evolve into partial link collection testing, matching links using TID/SIP.
    */
   private void assertSpanLinks(List<AgentSpanLink> links) {
+    // Check if links should be asserted at all
+    if (this.linkMatchers == null) {
+      return;
+    }
     int linkCount = links == null ? 0 : links.size();
-    int expectedLinkCount = this.linkMatchers == null ? 0 : this.linkMatchers.length;
+    int expectedLinkCount = this.linkMatchers.length;
     if (linkCount != expectedLinkCount) {
-      throw new AssertionFailedError("Unexpected span link count", expectedLinkCount, linkCount);
+      assertionFailure()
+          .message("Unexpected span link count")
+          .expected(expectedLinkCount)
+          .actual(linkCount)
+          .buildAndThrow();
     }
     for (int i = 0; i < expectedLinkCount; i++) {
       SpanLinkMatcher linkMatcher = this.linkMatchers[expectedLinkCount];

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

@@ -3,18 +3,28 @@
 import static datadog.trace.agent.test.assertions.Matchers.any;
 import static datadog.trace.agent.test.assertions.Matchers.is;
 import static datadog.trace.agent.test.assertions.Matchers.isNonNull;
+import static datadog.trace.api.DDTags.BASE_SERVICE;
+import static datadog.trace.api.DDTags.DD_INTEGRATION;
+import static datadog.trace.api.DDTags.DJM_ENABLED;
+import static datadog.trace.api.DDTags.DSM_ENABLED;
 import static datadog.trace.api.DDTags.ERROR_MSG;
 import static datadog.trace.api.DDTags.ERROR_STACK;
 import static datadog.trace.api.DDTags.ERROR_TYPE;
 import static datadog.trace.api.DDTags.LANGUAGE_TAG_KEY;
+import static datadog.trace.api.DDTags.PARENT_ID;
+import static datadog.trace.api.DDTags.PID_TAG;
+import static datadog.trace.api.DDTags.PROFILING_CONTEXT_ENGINE;
+import static datadog.trace.api.DDTags.PROFILING_ENABLED;
 import static datadog.trace.api.DDTags.REQUIRED_CODE_ORIGIN_TAGS;
 import static datadog.trace.api.DDTags.RUNTIME_ID_TAG;
+import static datadog.trace.api.DDTags.SCHEMA_VERSION_TAG_KEY;
+import static datadog.trace.api.DDTags.SPAN_LINKS;
 import static datadog.trace.api.DDTags.THREAD_ID;
 import static datadog.trace.api.DDTags.THREAD_NAME;
+import static datadog.trace.api.DDTags.TRACER_HOST;
 import static datadog.trace.common.sampling.RateByServiceTraceSampler.SAMPLING_AGENT_RATE;
 import static datadog.trace.common.writer.ddagent.TraceMapper.SAMPLING_PRIORITY_KEY;
 
-import datadog.trace.api.DDTags;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -34,15 +44,17 @@ public static TagsMatcher defaultTags() {
     tagMatchers.put(SAMPLING_AGENT_RATE, any());
     tagMatchers.put(SAMPLING_PRIORITY_KEY.toString(), any());
     tagMatchers.put("_sample_rate", any());
-    tagMatchers.put(DDTags.PID_TAG, any());
-    tagMatchers.put(DDTags.SCHEMA_VERSION_TAG_KEY, any());
-    tagMatchers.put(DDTags.PROFILING_ENABLED, any());
-    tagMatchers.put(DDTags.PROFILING_CONTEXT_ENGINE, any());
-    tagMatchers.put(DDTags.BASE_SERVICE, any());
-    tagMatchers.put(DDTags.DSM_ENABLED, any());
-    tagMatchers.put(DDTags.DJM_ENABLED, any());
-    tagMatchers.put(DDTags.PARENT_ID, any());
-    tagMatchers.put(DDTags.SPAN_LINKS, any()); // this is checked by LinksAsserter
+    tagMatchers.put(PID_TAG, any());
+    tagMatchers.put(SCHEMA_VERSION_TAG_KEY, any());
+    tagMatchers.put(PROFILING_ENABLED, any());
+    tagMatchers.put(PROFILING_CONTEXT_ENGINE, any());
+    tagMatchers.put(BASE_SERVICE, any());
+    tagMatchers.put(DSM_ENABLED, any());
+    tagMatchers.put(DJM_ENABLED, any());
+    tagMatchers.put(PARENT_ID, any());
+    tagMatchers.put(SPAN_LINKS, any()); // this is checked by LinksAsserter
+    tagMatchers.put(DD_INTEGRATION, any());
+    tagMatchers.put(TRACER_HOST, any());
 
     for (String tagName : REQUIRED_CODE_ORIGIN_TAGS) {
       tagMatchers.put(tagName, any());

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

@@ -1,12 +1,12 @@
 package datadog.trace.agent.test.assertions;
 
 import static java.util.function.Function.identity;
+import static org.junit.jupiter.api.AssertionFailureBuilder.assertionFailure;
 
 import datadog.trace.core.DDSpan;
 import java.util.Comparator;
 import java.util.List;
 import java.util.function.Function;
-import org.opentest4j.AssertionFailedError;
 
 /**
  * This class is a helper class to verify traces structure.
@@ -87,11 +87,19 @@ public static void assertTraces(
     int traceCount = traces.size();
     if (opts.ignoredAdditionalTraces) {
       if (traceCount < expectedTraceCount) {
-        throw new AssertionFailedError("Not enough of traces", expectedTraceCount, traceCount);
+        assertionFailure()
+            .message("Not enough of traces")
+            .expected(expectedTraceCount)
+            .actual(traceCount)
+            .buildAndThrow();
       }
     } else {
       if (traceCount != expectedTraceCount) {
-        throw new AssertionFailedError("Invalid number of traces", expectedTraceCount, traceCount);
+        assertionFailure()
+            .message("Invalid number of traces")
+            .expected(expectedTraceCount)
+            .actual(traceCount)
+            .buildAndThrow();
       }
     }
     if (opts.sorter != null) {