feat(test): Add instrumentation test for JUnit

Author: PerfectSlayer

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

@@ -1,9 +1,12 @@
 package datadog.trace.agent.test;
 
+import static java.util.function.Function.identity;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import datadog.instrument.classinject.ClassInjector;
+import datadog.trace.agent.test.assertions.TraceAssertions;
+import datadog.trace.agent.test.assertions.TraceMatcher;
 import datadog.trace.agent.tooling.AgentInstaller;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.agent.tooling.TracerInstaller;
@@ -22,11 +25,19 @@
 import java.util.ServiceLoader;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.function.Function;
 import net.bytebuddy.agent.ByteBuddyAgent;
 import org.junit.jupiter.api.AfterEach;
 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}.
+ */
 @ExtendWith(TestClassShadowingExtension.class)
 public abstract class AbstractInstrumentationTest {
   static final Instrumentation INSTRUMENTATION = ByteBuddyAgent.getInstrumentation();
@@ -100,6 +111,33 @@ public void tearDown() {
     this.transformerLister = null;
   }
 
+  /**
+   * Checks the structure of the traces captured from the test agent.
+   *
+   * @param matchers The matchers to verify the trace collection, one matcher by expected trace.
+   */
+  protected void assertTraces(TraceMatcher... matchers) {
+    assertTraces(identity(), matchers);
+  }
+
+  /**
+   * Checks the structure of the traces captured from the test agent.
+   *
+   * @param options The {@link TraceAssertions.Options} to configure the checks.
+   * @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) {
+    int expectedTraceCount = matchers.length;
+    try {
+      this.writer.waitForTraces(expectedTraceCount);
+    } catch (InterruptedException | TimeoutException e) {
+      throw new AssertionFailedError("Timeout while waiting for traces", e);
+    }
+    TraceAssertions.assertTraces(this.writer, options, matchers);
+  }
+
   protected void blockUntilChildSpansFinished(final int numberOfSpans) {
     blockUntilChildSpansFinished(this.tracer.activeSpan(), numberOfSpans);
   }

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

@@ -0,0 +1,26 @@
+package datadog.trace.agent.test.assertions;
+
+import java.util.Optional;
+
+public class Is<T> implements Matcher<T> {
+  private final T expected;
+
+  Is(T expected) {
+    this.expected = expected;
+  }
+
+  @Override
+  public Optional<T> expected() {
+    return Optional.of(this.expected);
+  }
+
+  @Override
+  public String message() {
+    return "Unexpected value";
+  }
+
+  @Override
+  public boolean test(T t) {
+    return this.expected.equals(t);
+  }
+}

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

@@ -0,0 +1,10 @@
+package datadog.trace.agent.test.assertions;
+
+import java.util.Optional;
+import java.util.function.Predicate;
+
+public interface Matcher<T> extends Predicate<T> {
+  Optional<T> expected();
+
+  String message();
+}

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

@@ -0,0 +1,26 @@
+package datadog.trace.agent.test.assertions;
+
+import java.util.function.Predicate;
+import java.util.regex.Pattern;
+
+public class Matchers {
+  public static <T> Matcher<T> is(T expected) {
+    return new Is<>(expected);
+  }
+
+  public static <T> Matcher<T> nonNull() {
+    return new NonNull<>();
+  }
+
+  public static Matcher<CharSequence> matches(String regex) {
+    return new Matches(Pattern.compile(regex));
+  }
+
+  public static Matcher<CharSequence> matches(Pattern pattern) {
+    return new Matches(pattern);
+  }
+
+  public static <T> Matcher<T> validates(Predicate<T> validator) {
+    return new Validates<>(validator);
+  }
+}

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

@@ -0,0 +1,27 @@
+package datadog.trace.agent.test.assertions;
+
+import java.util.Optional;
+import java.util.regex.Pattern;
+
+public class Matches implements Matcher<CharSequence> {
+  private final Pattern pattern;
+
+  Matches(Pattern pattern) {
+    this.pattern = pattern;
+  }
+
+  @Override
+  public Optional<CharSequence> expected() {
+    return Optional.empty();
+  }
+
+  @Override
+  public String message() {
+    return "Non matching value";
+  }
+
+  @Override
+  public boolean test(CharSequence s) {
+    return this.pattern.matcher(s).matches();
+  }
+}

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

@@ -0,0 +1,20 @@
+package datadog.trace.agent.test.assertions;
+
+import java.util.Optional;
+
+public class NonNull<T> implements Matcher<T> {
+  @Override
+  public Optional<T> expected() {
+    return Optional.empty();
+  }
+
+  @Override
+  public String message() {
+    return "Non-null value expected";
+  }
+
+  @Override
+  public boolean test(T t) {
+    return t != null;
+  }
+}

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

@@ -0,0 +1,127 @@
+package datadog.trace.agent.test.assertions;
+
+import static datadog.trace.agent.test.assertions.Matchers.is;
+import static datadog.trace.agent.test.assertions.Matchers.matches;
+import static datadog.trace.agent.test.assertions.Matchers.nonNull;
+import static datadog.trace.agent.test.assertions.Matchers.validates;
+import static java.time.temporal.ChronoUnit.NANOS;
+
+import datadog.trace.core.DDSpan;
+import java.time.Duration;
+import java.util.Optional;
+import java.util.function.Predicate;
+import java.util.regex.Pattern;
+import org.opentest4j.AssertionFailedError;
+
+public final class SpanMatcher {
+  private Matcher<Long> idMatcher;
+  private Matcher<Long> parentIdMatcher;
+  private Matcher<String> serviceNameMatcher;
+  private Matcher<CharSequence> operationNameMatcher;
+  private Matcher<CharSequence> resourceNameMatcher;
+  private Matcher<Duration> durationMatcher;
+
+  private static final Matcher<Long> CHILD_OF_PREVIOUS_MATCHER = is(0L);
+
+  private SpanMatcher() {}
+
+  public static SpanMatcher span() {
+    return new SpanMatcher();
+  }
+
+  public SpanMatcher withId(long id) {
+    this.idMatcher = is(id);
+    return this;
+  }
+
+  public SpanMatcher isRoot() {
+    return childOf(0L);
+  }
+
+  public SpanMatcher childOf(long parentId) {
+    this.parentIdMatcher = is(parentId);
+    return this;
+  }
+
+  public SpanMatcher childOfPrevious() {
+    this.parentIdMatcher = CHILD_OF_PREVIOUS_MATCHER;
+    return this;
+  }
+
+  public SpanMatcher hasServiceName() {
+    this.serviceNameMatcher = nonNull();
+    return this;
+  }
+
+  public SpanMatcher withServiceName(String serviceName) {
+    this.serviceNameMatcher = is(serviceName);
+    return this;
+  }
+
+  public SpanMatcher withOperationName(String operationName) {
+    this.operationNameMatcher = is(operationName);
+    return this;
+  }
+
+  public SpanMatcher operationNameMatching(Pattern pattern) {
+    this.operationNameMatcher = matches(pattern);
+    return this;
+  }
+
+  public SpanMatcher withResourceName(String resourceName) {
+    this.resourceNameMatcher = is(resourceName);
+    return this;
+  }
+
+  public SpanMatcher resourceNameMatching(Pattern pattern) {
+    this.resourceNameMatcher = matches(pattern);
+    return this;
+  }
+
+  public SpanMatcher resourceNameMatching(Predicate<CharSequence> validator) {
+    this.resourceNameMatcher = validates(validator);
+    return this;
+  }
+
+  public SpanMatcher durationShorterThan(Duration duration) {
+    this.durationMatcher = validates(d -> d.compareTo(duration) < 0);
+    return this;
+  }
+
+  public SpanMatcher durationLongerThan(Duration duration) {
+    this.durationMatcher = validates(d -> d.compareTo(duration) > 0);
+    return this;
+  }
+
+  public SpanMatcher durationMatching(Predicate<Duration> validator) {
+    this.durationMatcher = validates(validator);
+    return this;
+  }
+
+  public void assertSpan(DDSpan span, DDSpan previousSpan) {
+    // Apply parent id matcher from the previous span
+    if (this.parentIdMatcher == CHILD_OF_PREVIOUS_MATCHER) {
+      this.parentIdMatcher = is(previousSpan.getSpanId());
+    }
+    // Assert span values
+    assertValue(this.idMatcher, span.getSpanId(), "Expected identifier");
+    assertValue(this.parentIdMatcher, span.getParentId(), "Expected parent identifier");
+    assertValue(this.serviceNameMatcher, span.getServiceName(), "Expected service name");
+    assertValue(this.operationNameMatcher, span.getOperationName(), "Expected operation name");
+    assertValue(this.resourceNameMatcher, span.getResourceName(), "Expected resource name");
+    assertValue(
+        this.durationMatcher, Duration.of(span.getDurationNano(), NANOS), "Expected duration");
+    // TODO Add more values to test (tags, links, ...)
+  }
+
+  private <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.message(), expected.get(), value);
+      } else {
+        throw new AssertionFailedError(message + ": " + value + ". " + matcher.message());
+      }
+    }
+  }
+}