Fix ScalaTest suites run through JUnit 4

dd-java-agent/instrumentation/junit/junit-4/junit-4.10/src/main/java/datadog/trace/instrumentation/junit4/JUnit4Instrumentation.java

@@ -99,34 +99,48 @@ public static void onRunnerCreation(@Advice.This final Runner runner) {
 
   public static class JUnit4Advice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void addTracingListener(
+    public static JUnit4TracingListener addTracingListener(
         @Advice.This final Runner runner, @Advice.Argument(0) final RunNotifier runNotifier) {
       if (!JUnit4Utils.runnerToFramework(runner).equals(TestFrameworkInstrumentation.JUNIT4)) {
         // checking class names in hierarchyMatcher alone is not enough:
         // for example, Karate calls #run method of its super class,
         // that was transformed
-        return;
+        return null;
       }
 
       // No public accessor to get already installed listeners.
       // The installed RunListeners list are obtained using reflection.
       final List<RunListener> runListeners = JUnit4Utils.runListenersFromRunNotifier(runNotifier);
       if (runListeners == null) {
-        return;
+        return null;
       }
 
       for (final RunListener listener : runListeners) {
         RunListener tracingListener = JUnit4Utils.toTracingListener(listener);
         if (tracingListener != null) {
           // prevents installing TracingListener multiple times
-          return;
+          return null;
         }
       }
 
-      final TracingListener tracingListener =
+      final JUnit4TracingListener tracingListener =
           new JUnit4TracingListener(
               InstrumentationContext.get(Description.class, TestExecutionTracker.class));
       runNotifier.addListener(tracingListener);
+      return tracingListener;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void finishTestSuites(
+        @Advice.Enter final JUnit4TracingListener tracingListener,
+        @Advice.Argument(0) final RunNotifier runNotifier) {
+      if (tracingListener != null) {
+        try {
+          tracingListener.finishTestSuites();
+        } finally {
+          runNotifier.removeListener(tracingListener);
+        }
+      }
     }
 
     // JUnit 4.10 and above

dd-java-agent/instrumentation/junit/junit-4/junit-4.10/src/main/java/datadog/trace/instrumentation/junit4/JUnit4TracingListener.java

@@ -8,6 +8,8 @@
 import datadog.trace.bootstrap.ContextStore;
 import java.lang.reflect.Method;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import org.junit.Ignore;
 import org.junit.runner.Description;
 import org.junit.runner.notification.Failure;
@@ -18,6 +20,7 @@ public class JUnit4TracingListener extends TracingListener {
   private static final String FRAMEWORK_VERSION = JUnit4Utils.getVersion();
 
   private final ContextStore<Description, TestExecutionTracker> executionTrackers;
+  private final Map<TestSuiteDescriptor, Boolean> startedSuites = new ConcurrentHashMap<>();
 
   public JUnit4TracingListener(ContextStore<Description, TestExecutionTracker> executionTrackers) {
     this.executionTrackers = executionTrackers;
@@ -32,6 +35,21 @@ public void testSuiteStarted(final Description description) {
     }
 
     TestSuiteDescriptor suiteDescriptor = JUnit4Utils.toSuiteDescriptor(description);
+    startTestSuite(suiteDescriptor, description, false);
+  }
+
+  private void startTestSuite(
+      TestSuiteDescriptor suiteDescriptor, Description description, boolean parallelized) {
+    startedSuites.computeIfAbsent(
+        suiteDescriptor,
+        key -> {
+          onTestSuiteStart(key, description, parallelized);
+          return parallelized;
+        });
+  }
+
+  private void onTestSuiteStart(
+      TestSuiteDescriptor suiteDescriptor, Description description, boolean parallelized) {
     Class<?> testClass = description.getTestClass();
     String testSuiteName = JUnit4Utils.getSuiteName(testClass, description);
     List<String> categories = JUnit4Utils.getCategories(testClass, null);
@@ -44,7 +62,7 @@ public void testSuiteStarted(final Description description) {
             FRAMEWORK_VERSION,
             testClass,
             categories,
-            false,
+            parallelized,
             TestFrameworkInstrumentation.JUNIT4,
             null);
   }
@@ -58,18 +76,36 @@ public void testSuiteFinished(final Description description) {
     }
 
     TestSuiteDescriptor suiteDescriptor = JUnit4Utils.toSuiteDescriptor(description);
+    finishTestSuite(suiteDescriptor);
+  }
+
+  private void finishTestSuite(TestSuiteDescriptor suiteDescriptor) {
+    if (startedSuites.remove(suiteDescriptor) == null) {
+      return;
+    }
     TestEventsHandlerHolder.HANDLERS
         .get(TestFrameworkInstrumentation.JUNIT4)
         .onTestSuiteFinish(suiteDescriptor, null);
   }
 
+  public void finishTestSuites() {
+    // Fallback suites have no thread-bound scope: their callbacks may run on an executor.
+    // Suites with native lifecycle events are finished by their own callbacks.
+    for (Map.Entry<TestSuiteDescriptor, Boolean> suite : startedSuites.entrySet()) {
+      if (suite.getValue()) {
+        finishTestSuite(suite.getKey());
+      }
+    }
+  }
+
   @Override
   public void testStarted(final Description description) {
     if (JUnit4Utils.isJUnitPlatformRunnerTest(description)) {
       return;
     }
 
     TestSuiteDescriptor suiteDescriptor = JUnit4Utils.toSuiteDescriptor(description);
+    startTestSuite(suiteDescriptor, description, true);
     TestDescriptor testDescriptor = JUnit4Utils.toTestDescriptor(description);
     TestSourceData testSourceData = JUnit4Utils.toTestSourceData(description);
 
@@ -179,21 +215,7 @@ public void testIgnored(final Description description) {
 
       TestSuiteDescriptor suiteDescriptor = JUnit4Utils.toSuiteDescriptor(description);
       Class<?> testClass = description.getTestClass();
-      String testSuiteName = JUnit4Utils.getSuiteName(testClass, description);
-      List<String> categories = JUnit4Utils.getCategories(testClass, null);
-
-      TestEventsHandlerHolder.HANDLERS
-          .get(TestFrameworkInstrumentation.JUNIT4)
-          .onTestSuiteStart(
-              suiteDescriptor,
-              testSuiteName,
-              FRAMEWORK_NAME,
-              FRAMEWORK_VERSION,
-              testClass,
-              categories,
-              false,
-              TestFrameworkInstrumentation.JUNIT4,
-              null);
+      startTestSuite(suiteDescriptor, description, false);
       TestEventsHandlerHolder.HANDLERS
           .get(TestFrameworkInstrumentation.JUNIT4)
           .onTestSuiteSkip(suiteDescriptor, reason);
@@ -203,9 +225,7 @@ public void testIgnored(final Description description) {
         testIgnored(description, testMethod, reason);
       }
 
-      TestEventsHandlerHolder.HANDLERS
-          .get(TestFrameworkInstrumentation.JUNIT4)
-          .onTestSuiteFinish(suiteDescriptor, null);
+      finishTestSuite(suiteDescriptor);
     }
   }
 
@@ -215,6 +235,7 @@ private void testIgnored(Description description, Method testMethod, String reas
     }
 
     TestSuiteDescriptor suiteDescriptor = JUnit4Utils.toSuiteDescriptor(description);
+    startTestSuite(suiteDescriptor, description, true);
     TestDescriptor testDescriptor = JUnit4Utils.toTestDescriptor(description);
 
     Class<?> testClass = description.getTestClass();

dd-java-agent/instrumentation/junit/junit-4/junit-4.10/src/test/groovy/JUnit4Test.groovy

@@ -2,19 +2,121 @@ import datadog.trace.api.DisableTestTrace
 import datadog.trace.api.civisibility.config.TestFQN
 import datadog.trace.api.civisibility.config.TestIdentifier
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer
 import datadog.trace.civisibility.CiVisibilityInstrumentationTest
 import datadog.trace.civisibility.diff.LineDiff
 import datadog.trace.instrumentation.junit4.JUnit4Utils
 import datadog.trace.instrumentation.junit4.TestEventsHandlerHolder
 import org.example.*
 import org.junit.jupiter.api.Assumptions
+import org.junit.experimental.ParallelComputer
+import org.junit.runner.Computer
 import org.junit.runner.JUnitCore
 
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.Executors
+
 @DisableTestTrace(reason = "avoid self-tracing")
 class JUnit4Test extends CiVisibilityInstrumentationTest {
 
   def runner = new JUnitCore()
 
+  def "test runner without suite lifecycle events"() {
+    when:
+    TestEventsHandlerHolder.start(TestFrameworkInstrumentation.JUNIT4, JUnit4Utils.capabilities(false))
+    def result
+    try {
+      result = runner.run(TestSucceedCustomRunner)
+    } finally {
+      TestEventsHandlerHolder.stop(TestFrameworkInstrumentation.JUNIT4)
+    }
+
+    assert result.runCount == 2
+    assert result.ignoreCount == 1
+    assert result.failures*.description*.methodName == ["existing failure"]
+    assert result.failures*.exception*.class == [AssertionError]
+    assertSpansData("test-custom-runner")
+
+    then:
+    noExceptionThrown()
+  }
+
+  def "test multiple runners without suite lifecycle events #mode"() {
+    when:
+    def initialSpan = AgentTracer.activeSpan()
+    TestEventsHandlerHolder.start(TestFrameworkInstrumentation.JUNIT4, JUnit4Utils.capabilities(false))
+    def result
+    try {
+      result = runner.run(computer, tests as Class[])
+    } finally {
+      TestEventsHandlerHolder.stop(TestFrameworkInstrumentation.JUNIT4)
+    }
+
+    assert result.runCount == 4
+    assert result.ignoreCount == 2
+    assert result.failures*.description*.methodName == ["existing failure", "existing failure"]
+    assert result.failures*.exception*.class == [AssertionError, AssertionError]
+
+    def events = getEventsAsJson(TEST_WRITER.toList())
+    def suites = events.findAll { it.type == "test_suite_end" }
+    assert suites*.content*.meta*.get("test.suite").sort() == [TestSucceedCustomRunner.name, TestSucceedCustomRunner.AnotherSuite.name].sort()
+    suites.each { suite ->
+      def suiteTests = events.findAll { it.type == "test" && it.content.test_suite_id == suite.content.test_suite_id }
+      assert suiteTests.size() == 3
+      assert suiteTests.collectEntries { [(it.content.meta["test.name"]): it.content.meta["test.status"]] } == ["adds": "pass", "existing failure": "fail", "ignored": "skip"]
+    }
+    assert AgentTracer.activeSpan() == initialSpan
+
+    then:
+    noExceptionThrown()
+
+    where:
+    mode                  | computer                          | tests
+    "parallel"            | new ParallelComputer(true, false) | [TestSucceedCustomRunner, TestSucceedCustomRunner.AnotherSuite]
+    "sequential"          | new Computer()                    | [TestSucceedCustomRunner, TestSucceedCustomRunner.AnotherSuite]
+    "sequential reversed" | new Computer()                    | [TestSucceedCustomRunner.AnotherSuite, TestSucceedCustomRunner]
+  }
+
+  def "test executor callbacks without suite lifecycle events"() {
+    setup:
+    def executor = Executors.newSingleThreadExecutor()
+    TestSucceedExecutorRunner.executor = executor
+
+    when:
+    def initialSpan = AgentTracer.activeSpan()
+    def initialWorkerSpan = CompletableFuture.supplyAsync({ AgentTracer.activeSpan() }, executor).join()
+    TestEventsHandlerHolder.start(TestFrameworkInstrumentation.JUNIT4, JUnit4Utils.capabilities(false))
+    def result
+    def workerSpan
+    try {
+      result = runner.run(TestSucceedExecutorRunner)
+      workerSpan = CompletableFuture.supplyAsync({ AgentTracer.activeSpan() }, executor).join()
+    } finally {
+      TestEventsHandlerHolder.stop(TestFrameworkInstrumentation.JUNIT4)
+    }
+
+    assert workerSpan == initialWorkerSpan
+    assert AgentTracer.activeSpan() == initialSpan
+    assert result.runCount == 2
+    assert result.ignoreCount == 1
+    assert result.failures*.description*.methodName == ["existing failure"]
+    assert result.failures*.exception*.class == [AssertionError]
+
+    def events = getEventsAsJson(TEST_WRITER.toList())
+    def suites = events.findAll { it.type == "test_suite_end" }
+    assert suites*.content*.meta*.get("test.suite") == [TestSucceedExecutorRunner.name]
+    def suiteTests = events.findAll { it.type == "test" && it.content.test_suite_id == suites[0].content.test_suite_id }
+    assert suiteTests.size() == 3
+    assert suiteTests.collectEntries { [(it.content.meta["test.name"]): it.content.meta["test.status"]] } == ["adds": "pass", "existing failure": "fail", "ignored": "skip"]
+
+    then:
+    noExceptionThrown()
+
+    cleanup:
+    executor.shutdownNow()
+    TestSucceedExecutorRunner.executor = null
+  }
+
   def "test #testcaseName"() {
     runTests(tests, success)
 

dd-java-agent/instrumentation/junit/junit-4/junit-4.10/src/test/java/org/example/TestSucceedCustomRunner.java

@@ -0,0 +1,51 @@
+package org.example;
+
+import org.junit.runner.Description;
+import org.junit.runner.RunWith;
+import org.junit.runner.Runner;
+import org.junit.runner.notification.Failure;
+import org.junit.runner.notification.RunNotifier;
+
+@RunWith(TestSucceedCustomRunner.CustomRunner.class)
+public class TestSucceedCustomRunner {
+
+  public static class AnotherSuite extends TestSucceedCustomRunner {}
+
+  public static class CustomRunner extends Runner {
+    private final Class<?> testClass;
+
+    public CustomRunner(Class<?> testClass) {
+      this.testClass = testClass;
+    }
+
+    @Override
+    public Description getDescription() {
+      Description suite = Description.createSuiteDescription(testClass);
+      suite.addChild(Description.createTestDescription(testClass, "adds"));
+      suite.addChild(Description.createTestDescription(testClass, "existing failure"));
+      suite.addChild(Description.createTestDescription(testClass, "ignored"));
+      return suite;
+    }
+
+    @Override
+    public void run(RunNotifier notifier) {
+      emitTestEvents(notifier);
+    }
+
+    protected void emitTestEvents(RunNotifier notifier) {
+      // Like ScalaTest's JUnitRunner, emit test callbacks without suite lifecycle callbacks.
+      notifier.fireTestIgnored(testDescription("ignored"));
+      Description passing = testDescription("adds");
+      notifier.fireTestStarted(passing);
+      notifier.fireTestFinished(passing);
+      Description failing = testDescription("existing failure");
+      notifier.fireTestStarted(failing);
+      notifier.fireTestFailure(new Failure(failing, new AssertionError("existing failure")));
+      notifier.fireTestFinished(failing);
+    }
+
+    private Description testDescription(String name) {
+      return Description.createSuiteDescription(name + "(" + testClass.getName() + ")");
+    }
+  }
+}

dd-java-agent/instrumentation/junit/junit-4/junit-4.10/src/test/java/org/example/TestSucceedExecutorRunner.java

@@ -0,0 +1,23 @@
+package org.example;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import org.junit.runner.RunWith;
+import org.junit.runner.notification.RunNotifier;
+
+@RunWith(TestSucceedExecutorRunner.ExecutorRunner.class)
+public class TestSucceedExecutorRunner {
+
+  public static ExecutorService executor;
+
+  public static class ExecutorRunner extends TestSucceedCustomRunner.CustomRunner {
+    public ExecutorRunner(Class<?> testClass) {
+      super(testClass);
+    }
+
+    @Override
+    public void run(RunNotifier notifier) {
+      CompletableFuture.runAsync(() -> emitTestEvents(notifier), executor).join();
+    }
+  }
+}

dd-java-agent/instrumentation/junit/junit-4/junit-4.10/src/test/resources/test-custom-runner/coverages.ftl

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