Add instrumentation tests for JUnit

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 tracer.
+   *
+   * @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 tracer.
+   *
+   * @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/Any.java

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

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/IsFalse.java

@@ -0,0 +1,20 @@
+package datadog.trace.agent.test.assertions;
+
+import java.util.Optional;
+
+public class IsFalse implements Matcher<Boolean> {
+  @Override
+  public Optional<Boolean> expected() {
+    return Optional.of(false);
+  }
+
+  @Override
+  public String message() {
+    return "False expected";
+  }
+
+  @Override
+  public boolean test(Boolean t) {
+    return !t;
+  }
+}

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

@@ -0,0 +1,20 @@
+package datadog.trace.agent.test.assertions;
+
+import java.util.Optional;
+
+public class IsNonNull<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/IsNull.java

@@ -0,0 +1,20 @@
+package datadog.trace.agent.test.assertions;
+
+import java.util.Optional;
+
+public class IsNull<T> implements Matcher<T> {
+  @Override
+  public Optional<T> expected() {
+    return Optional.empty();
+  }
+
+  @Override
+  public String message() {
+    return "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/IsTrue.java

@@ -0,0 +1,20 @@
+package datadog.trace.agent.test.assertions;
+
+import java.util.Optional;
+
+public class IsTrue implements Matcher<Boolean> {
+  @Override
+  public Optional<Boolean> expected() {
+    return Optional.of(true);
+  }
+
+  @Override
+  public String message() {
+    return "True expected";
+  }
+
+  @Override
+  public boolean test(Boolean t) {
+    return 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,63 @@
+package datadog.trace.agent.test.assertions;
+
+import java.util.Optional;
+import java.util.function.Predicate;
+import java.util.regex.Pattern;
+import org.opentest4j.AssertionFailedError;
+
+/*
+ * TODO: Dev notes
+ * - introduce as few as possible matchers
+ * - only have matchers for generic purpose, don't introduce feature / produce / use-case specific matchers
+ * - name "ignores" as "any"?
+ * - think about extensibility? Open matchers for inheritance
+ */
+
+public class Matchers {
+  public static <T> Matcher<T> is(T expected) {
+    return new Is<>(expected);
+  }
+
+  public static <T> Matcher<T> isNull() {
+    return new IsNull<>();
+  }
+
+  public static <T> Matcher<T> isNonNull() {
+    return new IsNonNull<>();
+  }
+
+  public static Matcher<Boolean> isTrue() {
+    return new IsTrue();
+  }
+
+  public static Matcher<Boolean> isFalse() {
+    return new IsFalse();
+  }
+
+  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);
+  }
+
+  public static <T> Matcher<T> any() {
+    return new 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.message(), expected.get(), value);
+      } else {
+        throw new AssertionFailedError(message + ": " + value + ". " + matcher.message());
+      }
+    }
+  }
+}

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/SpanLinkMatcher.java

@@ -0,0 +1,68 @@
+package datadog.trace.agent.test.assertions;
+
+import static datadog.trace.agent.test.assertions.Matchers.assertValue;
+import static datadog.trace.agent.test.assertions.Matchers.is;
+import static datadog.trace.bootstrap.instrumentation.api.AgentSpanLink.DEFAULT_FLAGS;
+import static datadog.trace.bootstrap.instrumentation.api.SpanAttributes.EMPTY;
+
+import datadog.trace.api.DDTraceId;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
+import datadog.trace.bootstrap.instrumentation.api.SpanAttributes;
+import datadog.trace.core.DDSpan;
+
+public final class SpanLinkMatcher {
+  private final Matcher<DDTraceId> traceIdMatcher;
+  private final Matcher<Long> spanIdMatcher;
+  private Matcher<Byte> traceFlagsMatcher;
+  private Matcher<SpanAttributes> spanAttributesMatcher;
+  private Matcher<String> traceStateMatcher;
+
+  private SpanLinkMatcher(Matcher<DDTraceId> traceIdMatcher, Matcher<Long> spanIdMatcher) {
+    this.traceIdMatcher = traceIdMatcher;
+    this.spanIdMatcher = spanIdMatcher;
+    this.traceFlagsMatcher = is(DEFAULT_FLAGS);
+    this.spanAttributesMatcher = is(EMPTY);
+    this.traceStateMatcher = is("");
+  }
+
+  public static SpanLinkMatcher from(DDSpan span) {
+    return from(span.context());
+  }
+
+  public static SpanLinkMatcher from(AgentSpanContext spanContext) {
+    return from(spanContext.getTraceId(), spanContext.getSpanId());
+  }
+
+  public static SpanLinkMatcher from(DDTraceId traceId, long spanId) {
+    return new SpanLinkMatcher(is(traceId), is(spanId));
+  }
+
+  public static SpanLinkMatcher any() {
+    return new SpanLinkMatcher(Matchers.any(), Matchers.any());
+  }
+
+  public SpanLinkMatcher traceFlags(byte traceFlags) {
+    this.traceFlagsMatcher = is(traceFlags);
+    return this;
+  }
+
+  public SpanLinkMatcher attributes(SpanAttributes spanAttributes) {
+    this.spanAttributesMatcher = is(spanAttributes);
+    return this;
+  }
+
+  public SpanLinkMatcher traceState(String traceState) {
+    this.traceStateMatcher = is(traceState);
+    return this;
+  }
+
+  void assertLink(AgentSpanLink link) {
+    // Assert link values
+    assertValue(this.traceIdMatcher, link.traceId(), "Expected trace identifier");
+    assertValue(this.spanIdMatcher, link.spanId(), "Expected span identifier");
+    assertValue(this.traceFlagsMatcher, link.traceFlags(), "Expected trace flags");
+    assertValue(this.spanAttributesMatcher, link.attributes(), "Expected attributes");
+    assertValue(this.traceStateMatcher, link.traceState(), "Expected trace state");
+  }
+}