Trace setup and teardown operations in JUnit 5 (#7714)

Author: nikita-tkachenko-datadog

dd-java-agent/agent-ci-visibility/src/testFixtures/groovy/datadog/trace/civisibility/CiVisibilityInstrumentationTest.groovy

@@ -199,6 +199,11 @@ abstract class CiVisibilityInstrumentationTest extends AgentTestRunner {
     }
   }
 
+  @Override
+  protected String idGenerationStrategyName() {
+    "RANDOM"
+  }
+
   @Override
   void setup() {
     skippableTests.clear()

dd-java-agent/agent-ci-visibility/src/testFixtures/groovy/datadog/trace/civisibility/CiVisibilityTestUtils.groovy

@@ -41,6 +41,7 @@ abstract class CiVisibilityTestUtils {
     // Regardless, the values of these fields should be treated as different
     path("content.start", false),
     path("content.duration", false),
+    path("content.metrics.['_dd.host.vcpu_count']", false),
     path("content.meta.['_dd.p.tid']", false),
     path("content.meta.['error.stack']", false),
   ]
@@ -83,15 +84,15 @@ abstract class CiVisibilityTestUtils {
       replacementMap.put(labelGenerator.forKey(e.key), "\"$e.value\"")
     }
 
-    def expectedEvents = getFreemarkerTemplate(baseTemplatesPath + "/events.ftl", replacementMap)
+    def expectedEvents = getFreemarkerTemplate(baseTemplatesPath + "/events.ftl", replacementMap, events)
     def actualEvents = JSON_MAPPER.writeValueAsString(events)
     try {
       JSONAssert.assertEquals(expectedEvents, actualEvents, JSONCompareMode.LENIENT)
     } catch (AssertionError e) {
       throw new org.opentest4j.AssertionFailedError("Events mismatch", expectedEvents, actualEvents, e)
     }
 
-    def expectedCoverages = getFreemarkerTemplate(baseTemplatesPath + "/coverages.ftl", replacementMap)
+    def expectedCoverages = getFreemarkerTemplate(baseTemplatesPath + "/coverages.ftl", replacementMap, coverages)
     def actualCoverages = JSON_MAPPER.writeValueAsString(coverages)
     try {
       JSONAssert.assertEquals(expectedCoverages, actualCoverages, JSONCompareMode.LENIENT)
@@ -122,11 +123,15 @@ abstract class CiVisibilityTestUtils {
     }
   }
 
-  private static String getFreemarkerTemplate(String templatePath, Map<String, Object> replacements) {
-    Template coveragesTemplate = FREEMARKER.getTemplate(templatePath)
-    StringWriter coveragesOut = new StringWriter()
-    coveragesTemplate.process(replacements, coveragesOut)
-    return coveragesOut.toString()
+  private static String getFreemarkerTemplate(String templatePath, Map<String, Object> replacements, List<Map<?, ?>> replacementsSource) {
+    try {
+      Template coveragesTemplate = FREEMARKER.getTemplate(templatePath)
+      StringWriter coveragesOut = new StringWriter()
+      coveragesTemplate.process(replacements, coveragesOut)
+      return coveragesOut.toString()
+    } catch (Exception e) {
+      throw new RuntimeException("Could not get Freemarker template " + templatePath + "; replacements map: " + replacements + "; replacements source: " + replacementsSource, e)
+    }
   }
 
   private static final class TemplateGenerator {

dd-java-agent/instrumentation/junit-5.3/junit-5.8/build.gradle

@@ -0,0 +1,54 @@
+apply from: "$rootDir/gradle/java.gradle"
+
+def jupiterVersion = '5.8.0'
+def platformVersion = '1.8.0'
+
+muzzle {
+  pass {
+    group = 'org.junit.platform'
+    module = 'junit-platform-launcher'
+    versions = "[$platformVersion,)"
+    extraDependency "org.junit.jupiter:junit-jupiter-engine:$jupiterVersion"
+  }
+  pass {
+    group = 'org.junit.jupiter'
+    module = 'junit-jupiter-engine'
+    versions = "[$jupiterVersion,)"
+  }
+  pass {
+    group = 'org.junit.jupiter'
+    module = 'junit-jupiter-api'
+    versions = "[$jupiterVersion,)"
+    extraDependency "org.junit.jupiter:junit-jupiter-engine:$jupiterVersion"
+  }
+}
+
+addTestSuiteForDir('latestDepTest', 'test')
+
+dependencies {
+  implementation project(':dd-java-agent:instrumentation:junit-5.3')
+
+  compileOnly group: 'org.junit.platform', name: 'junit-platform-launcher', version: "$platformVersion"
+  compileOnly group: 'org.junit.jupiter', name: 'junit-jupiter-api', version: "$jupiterVersion"
+  compileOnly group: 'org.junit.jupiter', name: 'junit-jupiter-engine', version: "$jupiterVersion"
+
+  testImplementation testFixtures(project(':dd-java-agent:agent-ci-visibility'))
+
+  // versions used below are not the minimum ones that we support,
+  // but the tests need to use them in order to be compliant with Spock 2.x
+  testImplementation group: 'org.junit.platform', name: 'junit-platform-launcher', version: '1.8.2'
+  testImplementation group: 'org.junit.jupiter', name: 'junit-jupiter-engine', version: '5.8.2'
+  testImplementation group: 'org.junit.jupiter', name: 'junit-jupiter-params', version: '5.8.2'
+
+  latestDepTestImplementation group: 'org.junit.platform', name: 'junit-platform-launcher', version: '+'
+  latestDepTestImplementation group: 'org.junit.jupiter', name: 'junit-jupiter-engine', version: '+'
+  latestDepTestImplementation group: 'org.junit.jupiter', name: 'junit-jupiter-params', version: '+'
+}
+
+configurations.matching({ it.name.startsWith('test') }).each({
+  it.resolutionStrategy {
+    force group: 'org.junit.platform', name: 'junit-platform-launcher', version: '1.8.2'
+    force group: 'org.junit.jupiter', name: 'junit-jupiter-engine', version: '5.8.2'
+    force group: 'org.junit.jupiter', name: 'junit-jupiter-params', version: '5.8.2'
+  }
+})

dd-java-agent/instrumentation/junit-5.3/junit-5.8/src/main/java/datadog/trace/instrumentation/junit5/BeforeAfterOperationsTracer.java

@@ -0,0 +1,60 @@
+package datadog.trace.instrumentation.junit5;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import java.lang.reflect.Method;
+import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.api.extension.InvocationInterceptor;
+import org.junit.jupiter.api.extension.ReflectiveInvocationContext;
+
+public class BeforeAfterOperationsTracer implements InvocationInterceptor {
+
+  @Override
+  public void interceptBeforeAllMethod(
+      Invocation<Void> invocation,
+      ReflectiveInvocationContext<Method> invocationContext,
+      ExtensionContext extensionContext)
+      throws Throwable {
+    traceInvocation(invocation, invocationContext.getExecutable(), "BeforeAll");
+  }
+
+  @Override
+  public void interceptBeforeEachMethod(
+      Invocation<Void> invocation,
+      ReflectiveInvocationContext<Method> invocationContext,
+      ExtensionContext extensionContext)
+      throws Throwable {
+    traceInvocation(invocation, invocationContext.getExecutable(), "BeforeEach");
+  }
+
+  @Override
+  public void interceptAfterEachMethod(
+      Invocation<Void> invocation,
+      ReflectiveInvocationContext<Method> invocationContext,
+      ExtensionContext extensionContext)
+      throws Throwable {
+    traceInvocation(invocation, invocationContext.getExecutable(), "AfterEach");
+  }
+
+  @Override
+  public void interceptAfterAllMethod(
+      Invocation<Void> invocation,
+      ReflectiveInvocationContext<Method> invocationContext,
+      ExtensionContext extensionContext)
+      throws Throwable {
+    traceInvocation(invocation, invocationContext.getExecutable(), "AfterAll");
+  }
+
+  private static void traceInvocation(
+      Invocation<Void> invocation, Method executable, String operationName) throws Throwable {
+    AgentSpan agentSpan = AgentTracer.startSpan("junit", executable.getName());
+    agentSpan.setTag(Tags.TEST_CALLBACK, operationName);
+    try (AgentScope agentScope = AgentTracer.activateSpan(agentSpan)) {
+      invocation.proceed();
+    } finally {
+      agentSpan.finish();
+    }
+  }
+}

dd-java-agent/instrumentation/junit-5.3/junit-5.8/src/main/java/datadog/trace/instrumentation/junit5/JUnit5BeforeAfterInstrumentation.java

@@ -0,0 +1,53 @@
+package datadog.trace.instrumentation.junit5;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.junit.jupiter.engine.extension.ExtensionRegistrar;
+
+@AutoService(InstrumenterModule.class)
+public class JUnit5BeforeAfterInstrumentation extends InstrumenterModule.CiVisibility
+    implements Instrumenter.ForTypeHierarchy {
+
+  public JUnit5BeforeAfterInstrumentation() {
+    super("ci-visibility", "junit-5", "setup-teardown");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "org.junit.jupiter.engine.extension.ExtensionRegistrar";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named(hierarchyMarkerType()));
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".BeforeAfterOperationsTracer",
+    };
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isConstructor(),
+        JUnit5BeforeAfterInstrumentation.class.getName() + "$RegisterExtensionAdvice");
+  }
+
+  public static class RegisterExtensionAdvice {
+    @Advice.OnMethodExit
+    public static void registerExtension(@Advice.This ExtensionRegistrar extensionRegistrar) {
+      extensionRegistrar.registerExtension(BeforeAfterOperationsTracer.class);
+    }
+  }
+}

dd-java-agent/instrumentation/junit-5.3/junit-5.8/src/test/groovy/JUnit58Test.groovy

@@ -0,0 +1,63 @@
+import datadog.trace.api.DisableTestTrace
+import datadog.trace.civisibility.CiVisibilityInstrumentationTest
+import datadog.trace.instrumentation.junit5.TestEventsHandlerHolder
+import org.example.*
+import org.junit.jupiter.engine.JupiterTestEngine
+import org.junit.platform.engine.DiscoverySelector
+import org.junit.platform.launcher.core.LauncherConfig
+import org.junit.platform.launcher.core.LauncherDiscoveryRequestBuilder
+import org.junit.platform.launcher.core.LauncherFactory
+
+import static org.junit.platform.engine.discovery.DiscoverySelectors.selectClass
+
+@DisableTestTrace(reason = "avoid self-tracing")
+class JUnit58Test extends CiVisibilityInstrumentationTest {
+
+  def "test #testcaseName"() {
+    runTests(tests)
+
+    assertSpansData(testcaseName, expectedTracesCount)
+
+    where:
+    testcaseName                  | tests                            | expectedTracesCount
+    "test-before-each-after-each" | [TestSucceedBeforeEachAfterEach] | 2
+    "test-before-all-after-all"   | [TestSucceedBeforeAllAfterAll]   | 2
+  }
+
+  private static void runTests(List<Class<?>> tests) {
+    TestEventsHandlerHolder.startForcefully()
+
+    DiscoverySelector[] selectors = new DiscoverySelector[tests.size()]
+    for (i in 0..<tests.size()) {
+      selectors[i] = selectClass(tests[i])
+    }
+
+    def launcherReq = LauncherDiscoveryRequestBuilder.request()
+      .selectors(selectors)
+      .build()
+
+    def launcherConfig = LauncherConfig
+      .builder()
+      .enableTestEngineAutoRegistration(false)
+      .addTestEngines(new JupiterTestEngine())
+      .build()
+
+    def launcher = LauncherFactory.create(launcherConfig)
+    try {
+      launcher.execute(launcherReq)
+    } catch (Throwable ignored) {
+    }
+
+    TestEventsHandlerHolder.stop()
+  }
+
+  @Override
+  String instrumentedLibraryName() {
+    return "junit5"
+  }
+
+  @Override
+  String instrumentedLibraryVersion() {
+    return JupiterTestEngine.getPackage().getImplementationVersion()
+  }
+}

dd-java-agent/instrumentation/junit-5.3/junit-5.8/src/test/java/org/example/TestSucceedBeforeAllAfterAll.java

@@ -0,0 +1,24 @@
+package org.example;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import org.junit.jupiter.api.*;
+
+public class TestSucceedBeforeAllAfterAll {
+
+  @BeforeAll
+  public static void setUp() {}
+
+  @AfterAll
+  public static void tearDown() {}
+
+  @Test
+  public void test_succeed() {
+    assertTrue(true);
+  }
+
+  @Test
+  public void another_test_succeed() {
+    assertTrue(true);
+  }
+}

dd-java-agent/instrumentation/junit-5.3/junit-5.8/src/test/java/org/example/TestSucceedBeforeEachAfterEach.java

@@ -0,0 +1,26 @@
+package org.example;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class TestSucceedBeforeEachAfterEach {
+
+  @BeforeEach
+  public void setUp() {}
+
+  @AfterEach
+  public void tearDown() {}
+
+  @Test
+  public void test_succeed() {
+    assertTrue(true);
+  }
+
+  @Test
+  public void another_test_succeed() {
+    assertTrue(true);
+  }
+}

dd-java-agent/instrumentation/junit-5.3/junit-5.8/src/test/resources/test-before-all-after-all/coverages.ftl

@@ -0,0 +1 @@
+[ ]