Allow conditional flaky tests in JUnit

Author: sarahchen6

buildSrc/src/main/kotlin/dd-trace-java.configure-tests.gradle.kts

@@ -103,16 +103,12 @@ tasks.named("check") {
 }
 
 tasks.withType<Test>().configureEach {
-  // Flaky tests management for JUnit 5
-  (options as? JUnitPlatformOptions)?.apply {
-    if (skipFlakyTestsProvider.isPresent) {
-      excludeTags("flaky")
-    } else if (runFlakyTestsProvider.isPresent) {
-      includeTags("flaky")
-    }
+  // Keep suites without test-utils out of flaky-only runs. Runtime extensions refine this tag.
+  if (!skipFlakyTestsProvider.isPresent && runFlakyTestsProvider.isPresent) {
+    (options as? JUnitPlatformOptions)?.includeTags("flaky")
   }
 
-  // Set system property flag that is checked from tests to determine if they should be skipped or run
+  // Let the JUnit and Spock extensions evaluate @Flaky conditions before selecting tests.
   if (skipFlakyTestsProvider.isPresent) {
     jvmArgs("-Drun.flaky.tests=false")
   } else if (runFlakyTestsProvider.isPresent) {

dd-smoke-tests/sample-trace/src/test/groovy/datadog/smoketest/SampleTraceSmokeTest.groovy

@@ -0,0 +1,51 @@
+package datadog.smoketest;
+
+import static datadog.smoketest.backend.AgentBackend.testAgent;
+import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import datadog.environment.JavaVirtualMachine;
+import datadog.smoketest.backend.AgentBackend;
+import datadog.trace.test.util.Flaky;
+import java.io.File;
+import java.util.function.Predicate;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+class SampleTraceSmokeTest {
+  private static final AgentBackend BACKEND = testAgent();
+
+  @RegisterExtension
+  static final SmokeCliApp app =
+      SmokeCliApp.named("sample-trace")
+          .jar(System.getProperty("datadog.smoketest.agent.shadowJar.path"))
+          // tests tracer as a jar sending sample traces instead of a javaagent
+          .noAgent()
+          .backend(BACKEND)
+          .placeholder("agent.host", () -> BACKEND.url().getHost())
+          .placeholder("agent.port", () -> Integer.toString(BACKEND.port()))
+          .jvmArgs(
+              "-Ddd.agent.host=${agent.host}",
+              "-Ddd.trace.agent.port=${agent.port}",
+              "-Ddd.test.agent.session.token=" + BACKEND.sessionToken())
+          .args("sampleTrace", "-c", "10", "-i", "0.1")
+          .workingDirectory(new File(System.getProperty("datadog.smoketest.builddir")))
+          .build();
+
+  @Test
+  @DisplayName("sample traces are sent")
+  @Flaky(condition = IbmJvm.class)
+  void sampleTracesAreSent() {
+    app.traces().waitForTraceCount(10);
+    app.assertCompletesWithValue(30, SECONDS, 0);
+    assertEquals(10, app.traces().getTraces().size());
+  }
+
+  static class IbmJvm implements Predicate<String> {
+    @Override
+    public boolean test(String suite) {
+      return JavaVirtualMachine.isIbm();
+    }
+  }
+}

dd-smoke-tests/sample-trace/src/test/java/datadog/smoketest/SampleTraceSmokeTest.java

@@ -41,8 +41,33 @@ This mechanism exists to make sure either java agent state or static data are re
 
 ### Flaky Tests
 
-If a test runs unreliably, or doesn't have a fully deterministic behavior, this will lead to recurrent unexpected errors in continuous integration.
-In order to identify such tests and avoid the continuous integration to fail, they are marked as _flaky_ and must be annotated with the `@Flaky` annotation.
+Mark unreliable test methods or classes with `@Flaky` in both JUnit and Spock.
+
+All tests run by default. Use `-PskipFlakyTests` to skip flaky tests or `-PrunFlakyTests` to run only flaky tests.
+
+If a test is flaky only in certain environments, use `condition`. In Java, supply a predicate class
+with a no-argument constructor. Its `test` method returns `true` when the test is flaky:
+
+```java
+import datadog.environment.JavaVirtualMachine;
+import java.util.function.Predicate;
+
+@Test
+@Flaky(condition = IbmJvm.class)
+void testOnSupportedJvms() {
+  // ...
+}
+
+static class IbmJvm implements Predicate<String> {
+  @Override
+  public boolean test(String suite) {
+    return JavaVirtualMachine.isIbm();
+  }
+}
+```
+
+Use `suites = {"SomeSubclass"}` to limit the annotation to particular test classes, such as subclasses
+that inherit a test. When both `suites` and `condition` are specified, both must match.
 
 > [!TIP]
 > In case your pull request checks failed due to some unexpected flaky tests, you can retry the continuous 

docs/how_to_test.md

@@ -37,6 +37,7 @@ dependencies {
 
   compileOnly(project(":components:annotations"))
   compileOnly(libs.junit.jupiter)
+  compileOnly(libs.junit.platform.launcher)
   compileOnly(libs.logback.core)
   compileOnly(libs.logback.classic)
 

utils/test-utils/build.gradle.kts

@@ -54,7 +54,7 @@ org.junit.jupiter:junit-jupiter-params:5.14.1=compileClasspath,testCompileClassp
 org.junit.jupiter:junit-jupiter:5.14.1=compileClasspath,testCompileClasspath,testRuntimeClasspath
 org.junit.platform:junit-platform-commons:1.14.1=compileClasspath,testCompileClasspath,testRuntimeClasspath
 org.junit.platform:junit-platform-engine:1.14.1=compileClasspath,testCompileClasspath,testRuntimeClasspath
-org.junit.platform:junit-platform-launcher:1.14.1=testRuntimeClasspath
+org.junit.platform:junit-platform-launcher:1.14.1=compileClasspath,testRuntimeClasspath
 org.junit:junit-bom:5.14.1=compileClasspath,testCompileClasspath,testRuntimeClasspath
 org.mockito:mockito-core:4.4.0=testRuntimeClasspath
 org.objenesis:objenesis:3.3=compileClasspath,testCompileClasspath,testRuntimeClasspath

utils/test-utils/gradle.lockfile

@@ -6,6 +6,7 @@
 import java.lang.annotation.Target;
 import java.util.function.Predicate;
 import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.extension.ExtendWith;
 
 /**
  * Use this annotation for suites or test cases that are flaky. When running in CI, these will be
@@ -14,19 +15,21 @@
 @Retention(RetentionPolicy.RUNTIME)
 @Target({ElementType.TYPE, ElementType.METHOD})
 @Tag("flaky")
+@ExtendWith(FlakyJUnitExtension.class)
 public @interface Flaky {
   /** Reason why the test is flaky (optional). */
   String value() default "";
 
   /**
-   * Fully qualified name of the test suite classes where this test is flaky. Only required when the
-   * test is flaky only when run in a subclass.
+   * Names of the test suite classes where this test is flaky, typically subclasses that inherit the
+   * test. Spock uses simple class names; JUnit accepts simple or fully qualified class names.
    */
   String[] suites() default {};
 
   /**
-   * Closure with a predicate to test at runtime if the actual spec is flaky (e.g. check the JVM
-   * vendor), the parameter is the actual name of the spec under test
+   * Predicate class with a no-argument constructor that determines whether the test is flaky (e.g.
+   * check the JVM vendor). JUnit passes the concrete test class's simple name. Spock also supports
+   * Groovy closures.
    */
   Class<? extends Predicate<String>> condition() default True.class;
 

utils/test-utils/src/main/groovy/datadog/trace/test/util/Flaky.java

@@ -0,0 +1,82 @@
+package datadog.trace.test.util;
+
+import static org.junit.platform.commons.support.AnnotationSupport.findAnnotation;
+
+import java.lang.reflect.AnnotatedElement;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Method;
+import java.util.function.Predicate;
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.extension.ConditionEvaluationResult;
+import org.junit.jupiter.api.extension.ExecutionCondition;
+import org.junit.jupiter.api.extension.ExtensionConfigurationException;
+import org.junit.jupiter.api.extension.ExtensionContext;
+
+/** Selects JUnit tests using the same flaky-test modes as {@link FlakySpockExtension}. */
+public final class FlakyJUnitExtension implements ExecutionCondition {
+  private static final String RUN_FLAKY_TESTS = "run.flaky.tests";
+
+  @Override
+  public ConditionEvaluationResult evaluateExecutionCondition(ExtensionContext context) {
+    if (!"false".equals(System.getProperty(RUN_FLAKY_TESTS))
+        || !context.getTestClass().isPresent()) {
+      return ConditionEvaluationResult.enabled("Flaky tests are not skipped");
+    }
+    Flaky flaky = findFlaky(context.getRequiredTestClass(), context.getTestMethod().orElse(null));
+    if (flaky == null) {
+      return ConditionEvaluationResult.enabled("Test is not flaky");
+    }
+    return ConditionEvaluationResult.disabled(
+        flaky.value().isEmpty() ? "Flaky test" : "Flaky test: " + flaky.value());
+  }
+
+  static Flaky findFlaky(Class<?> testClass, Method method) {
+    Flaky flaky = findFlakyClass(testClass);
+    return flaky != null || method == null ? flaky : matchingAnnotation(method, testClass);
+  }
+
+  private static Flaky findFlakyClass(Class<?> testClass) {
+    for (Class<?> current = testClass; current != null; current = current.getSuperclass()) {
+      Flaky flaky = matchingAnnotation(current, testClass);
+      if (flaky != null) {
+        return flaky;
+      }
+    }
+    if (testClass.isAnnotationPresent(Nested.class)) {
+      return findFlakyClass(testClass.getEnclosingClass());
+    }
+    return null;
+  }
+
+  private static Flaky matchingAnnotation(AnnotatedElement element, Class<?> testClass) {
+    Flaky flaky = findAnnotation(element, Flaky.class).orElse(null);
+    if (flaky == null) {
+      return null;
+    }
+    if (flaky.suites().length > 0) {
+      boolean matches = false;
+      for (String suite : flaky.suites()) {
+        if (suite.equals(testClass.getSimpleName()) || suite.equals(testClass.getName())) {
+          matches = true;
+          break;
+        }
+      }
+      if (!matches) {
+        return null;
+      }
+    }
+    if (flaky.condition() == Flaky.True.class) {
+      return flaky;
+    }
+    try {
+      Constructor<? extends Predicate<String>> constructor =
+          flaky.condition().getDeclaredConstructor();
+      constructor.setAccessible(true);
+      return constructor.newInstance().test(testClass.getSimpleName()) ? flaky : null;
+    } catch (ReflectiveOperationException | RuntimeException e) {
+      throw new ExtensionConfigurationException(
+          "Could not evaluate @Flaky condition " + flaky.condition().getName() + " on " + element,
+          e);
+    }
+  }
+}

utils/test-utils/src/main/groovy/datadog/trace/test/util/FlakyJUnitExtension.java

@@ -0,0 +1,29 @@
+package datadog.trace.test.util;
+
+import static datadog.trace.test.util.FlakyJUnitExtension.findFlaky;
+
+import org.junit.platform.engine.FilterResult;
+import org.junit.platform.engine.TestDescriptor;
+import org.junit.platform.engine.TestSource;
+import org.junit.platform.engine.support.descriptor.MethodSource;
+import org.junit.platform.launcher.PostDiscoveryFilter;
+
+/** Excludes non-flaky JUnit tests from flaky-only runs before execution. */
+public final class FlakyJUnitFilter implements PostDiscoveryFilter {
+  @Override
+  public FilterResult apply(TestDescriptor descriptor) {
+    if (!"true".equals(System.getProperty("run.flaky.tests"))
+        || !descriptor.getUniqueId().getEngineId().filter("junit-jupiter"::equals).isPresent()) {
+      return FilterResult.included("Flaky-only filtering does not apply");
+    }
+    TestSource source = descriptor.getSource().orElse(null);
+    if (source instanceof MethodSource) {
+      MethodSource method = (MethodSource) source;
+      return findFlaky(method.getJavaClass(), method.getJavaMethod()) != null
+          ? FilterResult.included("Flaky test")
+          : FilterResult.excluded("Test is not flaky");
+    }
+    // Keep containers until their methods have been filtered; JUnit prunes empty containers.
+    return FilterResult.included("Container may contain flaky tests");
+  }
+}

utils/test-utils/src/main/groovy/datadog/trace/test/util/FlakyJUnitFilter.java

@@ -42,6 +42,11 @@ class FlakySpockExtension extends AbstractGlobalExtension {
         }
       }
     }
+
+    if (shouldRunFlakyTestsOnly()) {
+      // Preserve selected features through the JUnit Platform's flaky tag filter.
+      spec.getAllFeatures().findAll { !it.excluded }.each { it.addTestTag("flaky") }
+    }
   }
 
   private static void skip(final node) {