Skip to content

Improve trace assert API #11209

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gh-worker-dd-mergequeue-cf854d merged 1 commit into from
Apr 28, 2026
Merged

Improve trace assert API #11209

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 3 additions & 4 deletions ...mentation-testing/src/main/java/datadog/trace/agent/test/AbstractInstrumentationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;

Expand All @@ -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;
Expand Down Expand Up @@ -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);
Expand Down
15 changes: 15 additions & 0 deletions ...nstrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/SpanMatcher.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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.
*
Expand Down Expand Up @@ -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");
Expand Down
20 changes: 10 additions & 10 deletions ...umentation-testing/src/main/java/datadog/trace/agent/test/assertions/TraceAssertions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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:
*
Expand All @@ -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 =
Copy link
Copy Markdown
Contributor

@amarziali amarziali Apr 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nice

Copy link
Copy Markdown
Contributor Author

@PerfectSlayer PerfectSlayer Apr 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I figured out I was using it wrong during a Devoxx talk when I saw howe they implemented StructuredTaskScope Ruby-style builder 😬 Not too late to fit it 😇

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() {}
Expand Down Expand Up @@ -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();
Expand Down
12 changes: 7 additions & 5 deletions ...strumentation-testing/src/main/java/datadog/trace/agent/test/assertions/TraceMatcher.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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;
Expand Down Expand Up @@ -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);
}

Expand Down
Loading