feat(test): Improve trace assert API

PerfectSlayer · PerfectSlayer · commit 5e0f9d55f94c · 2026-04-27T13:26:03.000+02:00
diff --git a/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/AbstractInstrumentationTest.java b/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/AbstractInstrumentationTest.java
@@ -1,6 +1,6 @@
 package datadog.trace.agent.test;
 
-import static java.util.function.Function.identity;
+import static java.util.function.UnaryOperator.identity;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -27,8 +27,8 @@
 import java.util.ServiceLoader;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
-import java.util.function.Function;
 import java.util.function.Predicate;
+import java.util.function.UnaryOperator;
 import net.bytebuddy.agent.ByteBuddyAgent;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
@@ -148,8 +148,7 @@ protected void assertTraces(TraceMatcher... matchers) {
    * @param matchers The matchers to verify the trace collection, one matcher by expected trace.
    */
   protected void assertTraces(
-      Function<TraceAssertions.Options, TraceAssertions.Options> options,
-      TraceMatcher... matchers) {
+      UnaryOperator<TraceAssertions.Options> options, TraceMatcher... matchers) {
     int expectedTraceCount = matchers.length;
     try {
       writer.waitForTraces(expectedTraceCount);
diff --git a/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/SpanMatcher.java b/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/SpanMatcher.java
@@ -12,6 +12,7 @@
 import static java.time.Duration.ofNanos;
 import static org.junit.jupiter.api.AssertionFailureBuilder.assertionFailure;
 
+import datadog.trace.api.DDTraceId;
 import datadog.trace.api.TagMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
 import datadog.trace.core.DDSpan;
@@ -49,6 +50,7 @@
  * </ul>
  */
 public final class SpanMatcher {
+  private Matcher<DDTraceId> traceIdMatcher;
   private Matcher<Long> idMatcher;
   private Matcher<Long> parentIdMatcher;
   private Matcher<String> serviceNameMatcher;
@@ -77,6 +79,18 @@ public static SpanMatcher span() {
     return new SpanMatcher();
   }
 
+  /**
+   * Checks the trace identifier matches the given value.
+   *
+   * @param traceId The trace identifier to match against.
+   * @return The current {@link SpanMatcher} instance with the specified trace identifier constraint
+   *     applied.
+   */
+  public SpanMatcher traceId(DDTraceId traceId) {
+    this.traceIdMatcher = Matchers.is(traceId);
+    return this;
+  }
+
   /**
    * Checks the span identifier matches the given value.
    *
@@ -293,6 +307,7 @@ void assertSpan(DDSpan span, DDSpan previousSpan) {
       this.parentIdMatcher = is(previousSpan.getSpanId());
     }
     // Assert span values
+    assertValue(this.traceIdMatcher, span.getTraceId(), "Expected trace identifier");
     assertValue(this.idMatcher, span.getSpanId(), "Expected identifier");
     assertValue(this.parentIdMatcher, span.getParentId(), "Expected parent identifier");
     assertValue(this.serviceNameMatcher, span.getServiceName(), "Expected service name");
diff --git a/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/TraceAssertions.java b/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/TraceAssertions.java
@@ -1,20 +1,20 @@
 package datadog.trace.agent.test.assertions;
 
-import static java.util.function.Function.identity;
+import static java.util.function.UnaryOperator.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 java.util.function.UnaryOperator;
 
 /**
- * This class is a helper class to verify traces structure.
+ * This class is a helper class to verify trace structure.
  *
- * <p>To check for traces structure, use the static factory methods: {@link #assertTraces(List,
+ * <p>To check for trace structure, use the static factory methods: {@link #assertTraces(List,
  * TraceMatcher...)} with the expected {@link TraceMatcher}s (one per trace), or {@link
- * #assertTraces(List, Function, TraceMatcher...)} to configure the checks with a {@link Options}
- * object.
+ * #assertTraces(List, UnaryOperator, TraceMatcher...)} to configure the checks with a {@link
+ * Options} object.
  *
  * <p>The following predefined configurations:
  *
@@ -40,15 +40,15 @@ public final class TraceAssertions {
    * Trace assertions options.
    */
   /** Ignores addition traces. If there are more traces than expected, do not fail. */
-  public static final Function<Options, Options> IGNORE_ADDITIONAL_TRACES =
+  public static final UnaryOperator<Options> IGNORE_ADDITIONAL_TRACES =
       Options::ignoredAdditionalTraces;
 
   /** Sorts traces by start time. */
-  public static final Function<Options, Options> SORT_BY_START_TIME =
+  public static final UnaryOperator<Options> SORT_BY_START_TIME =
       options -> options.sorter(TRACE_START_TIME_COMPARATOR);
 
   /** Sorts traces by their root span identifier. */
-  public static final Function<Options, Options> SORT_BY_ROOT_SPAN_ID =
+  public static final UnaryOperator<Options> SORT_BY_ROOT_SPAN_ID =
       options -> options.sorter(TRACE_ROOT_SPAN_ID_COMPARATOR);
 
   private TraceAssertions() {}
@@ -81,7 +81,7 @@ public static void assertTraces(List<List<DDSpan>> traces, TraceMatcher... match
    * @param matchers The matchers to verify the trace collection, one matcher by expected trace.
    */
   public static void assertTraces(
-      List<List<DDSpan>> traces, Function<Options, Options> options, TraceMatcher... matchers) {
+      List<List<DDSpan>> traces, UnaryOperator<Options> options, TraceMatcher... matchers) {
     Options opts = options.apply(new Options());
     int expectedTraceCount = matchers.length;
     int traceCount = traces.size();
diff --git a/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/TraceMatcher.java b/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/TraceMatcher.java
@@ -1,16 +1,18 @@
 package datadog.trace.agent.test.assertions;
 
+import static java.util.Comparator.comparingLong;
+
 import datadog.trace.core.DDSpan;
 import java.util.Comparator;
 import java.util.List;
-import java.util.function.Function;
+import java.util.function.UnaryOperator;
 import org.opentest4j.AssertionFailedError;
 
 /**
  * This class is a helper class to verify a trace structure.
  *
  * <p>To get a {@code TraceMatcher}, use the static factory methods: {@link #trace(SpanMatcher...)}
- * with the expected {@link SpanMatcher}s (one per expected span), or {@link #trace(Function,
+ * with the expected {@link SpanMatcher}s (one per expected span), or {@link #trace(UnaryOperator,
  * SpanMatcher...)} to configure the checks with a {@link Options} object.
  *
  * <p>{@link #SORT_BY_START_TIME} can be used as predefined configuration to sort spans by start
@@ -21,8 +23,8 @@
  */
 public final class TraceMatcher {
   public static final Comparator<DDSpan> START_TIME_COMPARATOR =
-      Comparator.comparingLong(DDSpan::getStartTime);
-  public static Function<Options, Options> SORT_BY_START_TIME =
+      comparingLong(DDSpan::getStartTime);
+  public static UnaryOperator<Options> SORT_BY_START_TIME =
       options -> options.sorter(START_TIME_COMPARATOR);
 
   private final Options options;
@@ -51,7 +53,7 @@ public static TraceMatcher trace(SpanMatcher... matchers) {
    * @param options The {@link TraceAssertions.Options} to configure the checks.
    * @param matchers The matchers to verify the trace structure.
    */
-  public static TraceMatcher trace(Function<Options, Options> options, SpanMatcher... matchers) {
+  public static TraceMatcher trace(UnaryOperator<Options> options, SpanMatcher... matchers) {
     return new TraceMatcher(options.apply(new Options()), matchers);
   }
 
