Implement flaky test retries for Test NG

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/DDTestFrameworkModule.java

@@ -2,6 +2,7 @@
 
 import datadog.trace.api.civisibility.config.TestIdentifier;
 import datadog.trace.api.civisibility.retry.TestRetryPolicy;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
 /** Test module abstraction that is used by test framework instrumentations (e.g. JUnit, TestNG) */
@@ -29,6 +30,7 @@ DDTestSuiteImpl testSuiteStart(
    */
   boolean skip(TestIdentifier test);
 
+  @Nonnull
   TestRetryPolicy retryPolicy(TestIdentifier test);
 
   void end(Long startTime);

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/DDTestFrameworkModuleImpl.java

@@ -18,6 +18,7 @@
 import java.util.HashSet;
 import java.util.concurrent.atomic.LongAdder;
 import java.util.function.Consumer;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
 /**
@@ -81,6 +82,7 @@ public boolean skip(TestIdentifier test) {
   }
 
   @Override
+  @Nonnull
   public TestRetryPolicy retryPolicy(TestIdentifier test) {
     return test != null && flakyTests.contains(test)
         ? new RetryIfFailed(config.getCiVisibilityFlakyRetryCount())

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/DDTestFrameworkModuleProxy.java

@@ -22,6 +22,7 @@
 import java.util.TreeSet;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.LongAdder;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -95,6 +96,7 @@ public boolean skip(TestIdentifier test) {
   }
 
   @Override
+  @Nonnull
   public TestRetryPolicy retryPolicy(TestIdentifier test) {
     return test != null && flakyTests.contains(test)
         ? new RetryIfFailed(config.getCiVisibilityFlakyRetryCount())

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/events/TestEventsHandlerImpl.java

@@ -17,6 +17,7 @@
 import java.util.Collections;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import org.objectweb.asm.Type;
 import org.slf4j.Logger;
@@ -289,6 +290,7 @@ public boolean skip(TestIdentifier test) {
   }
 
   @Override
+  @Nonnull
   public TestRetryPolicy retryPolicy(TestIdentifier test) {
     return testModule.retryPolicy(test);
   }

dd-java-agent/instrumentation/testng/src/main/java/datadog/trace/instrumentation/testng/TestNGInstrumentation.java

@@ -1,10 +1,11 @@
 package datadog.trace.instrumentation.testng;
 
-import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
-
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.api.Config;
+import datadog.trace.instrumentation.testng.retry.RetryAnnotationTransformer;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
 import org.testng.ITestListener;
 import org.testng.ITestNGListener;
 import org.testng.TestNG;
@@ -25,8 +26,7 @@ public String instrumentedType() {
   @Override
   public void adviceTransformations(AdviceTransformation transformation) {
     transformation.applyAdvice(
-        named("initializeDefaultListeners"),
-        TestNGInstrumentation.class.getName() + "$TestNGAdvice");
+        MethodDescription::isConstructor, TestNGInstrumentation.class.getName() + "$TestNGAdvice");
   }
 
   @Override
@@ -36,7 +36,9 @@ public String[] helperClassNames() {
       packageName + ".TestNGSuiteListener",
       packageName + ".TestNGClassListener",
       packageName + ".TestEventsHandlerHolder",
-      packageName + ".TracingListener"
+      packageName + ".TracingListener",
+      packageName + ".retry.RetryAnalyzer",
+      packageName + ".retry.RetryAnnotationTransformer",
     };
   }
 
@@ -54,6 +56,12 @@ public static void addTracingListener(@Advice.This final TestNG testNG) {
 
       TestNGSuiteListener suiteListener = new TestNGSuiteListener(tracingListener);
       testNG.addListener((ITestNGListener) suiteListener);
+
+      if (Config.get().isCiVisibilityFlakyRetryEnabled()) {
+        final RetryAnnotationTransformer transformer =
+            new RetryAnnotationTransformer(testNG.getAnnotationTransformer());
+        testNG.addListener(transformer);
+      }
     }
 
     // TestNG 6.4 and above

dd-java-agent/instrumentation/testng/src/main/java/datadog/trace/instrumentation/testng/retry/RetryAnalyzer.java

@@ -0,0 +1,31 @@
+package datadog.trace.instrumentation.testng.retry;
+
+import datadog.trace.api.civisibility.config.TestIdentifier;
+import datadog.trace.api.civisibility.retry.TestRetryPolicy;
+import datadog.trace.instrumentation.testng.TestEventsHandlerHolder;
+import org.testng.IRetryAnalyzer;
+import org.testng.ITestResult;
+
+public class RetryAnalyzer implements IRetryAnalyzer {
+
+  private volatile TestRetryPolicy retryPolicy;
+
+  @Override
+  public boolean retry(ITestResult result) {
+    if (TestEventsHandlerHolder.TEST_EVENTS_HANDLER == null) {
+      return false;
+    }
+    if (retryPolicy == null) {
+      synchronized (this) {
+        if (retryPolicy == null) {
+          String testSuiteName = result.getInstanceName();
+          String testName =
+              (result.getName() != null) ? result.getName() : result.getMethod().getMethodName();
+          TestIdentifier testIdentifier = new TestIdentifier(testSuiteName, testName, null, null);
+          retryPolicy = TestEventsHandlerHolder.TEST_EVENTS_HANDLER.retryPolicy(testIdentifier);
+        }
+      }
+    }
+    return retryPolicy.retry(result.isSuccess());
+  }
+}

dd-java-agent/instrumentation/testng/src/main/java/datadog/trace/instrumentation/testng/retry/RetryAnnotationTransformer.java

@@ -0,0 +1,23 @@
+package datadog.trace.instrumentation.testng.retry;
+
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Method;
+import org.testng.IAnnotationTransformer;
+import org.testng.annotations.ITestAnnotation;
+
+public class RetryAnnotationTransformer implements IAnnotationTransformer {
+  private final IAnnotationTransformer delegate;
+
+  public RetryAnnotationTransformer(IAnnotationTransformer delegate) {
+    this.delegate = delegate;
+  }
+
+  @Override
+  public void transform(
+      ITestAnnotation annotation, Class testClass, Constructor testConstructor, Method testMethod) {
+    annotation.setRetryAnalyzer(RetryAnalyzer.class);
+    if (delegate != null) {
+      delegate.transform(annotation, testClass, testConstructor, testMethod);
+    }
+  }
+}

dd-java-agent/instrumentation/testng/src/testFixtures/groovy/datadog/trace/instrumentation/testng/TestNGTest.groovy

@@ -1,13 +1,14 @@
 package datadog.trace.instrumentation.testng
 
-
 import datadog.trace.api.civisibility.config.TestIdentifier
 import datadog.trace.civisibility.CiVisibilityInstrumentationTest
 import org.example.TestError
 import org.example.TestFailed
 import org.example.TestFailedAndSucceed
+import org.example.TestFailedParameterized
 import org.example.TestFailedSuiteSetup
 import org.example.TestFailedSuiteTearDown
+import org.example.TestFailedThenSucceed
 import org.example.TestFailedWithSuccessPercentage
 import org.example.TestInheritance
 import org.example.TestParameterized
@@ -30,57 +31,93 @@ abstract class TestNGTest extends CiVisibilityInstrumentationTest {
   static testOutputDir = "build/tmp/test-output"
 
   def "test #testcaseName"() {
-    setup:
-    givenSkippableTests(skippedTests)
+    runTests(tests, null)
+
+    assertSpansData(testcaseName, expectedTracesCount)
+
+    where:
+    testcaseName                                          | tests                                              | expectedTracesCount
+    "test-succeed"                                        | [TestSucceed]                                      | 2
+    "test-inheritance"                                    | [TestInheritance]                                  | 2
+    "test-failed-${version()}"                            | [TestFailed]                                       | 2
+    "test-failed-with-success-percentage-${version()}"    | [TestFailedWithSuccessPercentage]                  | 6
+    "test-error"                                          | [TestError]                                        | 2
+    "test-skipped"                                        | [TestSkipped]                                      | 2
+    "test-parameterized"                                  | [TestParameterized]                                | 3
+    "test-success-with-groups"                            | [TestSucceedGroups]                                | 2
+    "test-class-skipped"                                  | [TestSkippedClass]                                 | 3
+    "test-success-and-skipped"                            | [TestSucceedAndSkipped]                            | 3
+    "test-success-and-failure-${version()}"               | [TestFailedAndSucceed]                             | 4
+    "test-suite-teardown-failure"                         | [TestFailedSuiteTearDown]                          | 3
+    "test-suite-setup-failure"                            | [TestFailedSuiteSetup]                             | 3
+    "test-multiple-successful-suites"                     | [TestSucceed, TestSucceedAndSkipped]               | 4
+    "test-successful-suite-and-failed-suite-${version()}" | [TestSucceed, TestFailedAndSucceed]                | 5
+    "test-nested-successful-suites"                       | [TestSucceedNested, TestSucceedNested.NestedSuite] | 3
+    "test-nested-skipped-suites-${version()}"             | [TestSkippedNested]                                | 3
+    "test-factory-data-provider"                          | [TestSucceedDataProvider]                          | 2
+  }
+
+  def "test parallel execution #testcaseName"() {
+    runTests(tests, parallelMode)
+
+    assertSpansData(testcaseName, expectedTracesCount)
+
+    where:
+    testcaseName                                | tests                 | expectedTracesCount | parallelMode
+    "test-successful-test-cases-in-parallel"    | [TestSucceedMultiple] | 3                   | "methods"
+    "test-parameterized-test-cases-in-parallel" | [TestParameterized]   | 3                   | "methods"
+  }
 
-    if (!tests.isEmpty()) {
-      runTests(tests, parallelMode)
-    } else {
-      def xmlSuite = null
-      TestNGTest.classLoader.getResourceAsStream(testcaseName + "/suite.xml").withCloseable {
-        xmlSuite = new SuiteXmlParser().parse("testng.xml", it, true)
-      }
-      runXmlSuites([xmlSuite], parallelMode)
+  def "test XML suites #testcaseName"() {
+    def xmlSuite = null
+    TestNGTest.classLoader.getResourceAsStream(testcaseName + "/suite.xml").withCloseable {
+      xmlSuite = new SuiteXmlParser().parse("testng.xml", it, true)
     }
+    runXmlSuites([xmlSuite], parallelMode)
 
-    expect:
     assertSpansData(testcaseName, expectedTracesCount)
 
     where:
-    testcaseName                                                                      | tests                                              | expectedTracesCount | parallelMode | skippedTests
-    "test-succeed"                                                                    | [TestSucceed]                                      | 2                   | null         | []
-    "test-inheritance"                                                                | [TestInheritance]                                  | 2                   | null         | []
-    "test-failed-${version()}"                                                        | [TestFailed]                                       | 2                   | null         | []
-    "test-failed-with-success-percentage-${version()}"                                | [TestFailedWithSuccessPercentage]                  | 6                   | null         | []
-    "test-error"                                                                      | [TestError]                                        | 2                   | null         | []
-    "test-skipped"                                                                    | [TestSkipped]                                      | 2                   | null         | []
-    "test-parameterized"                                                              | [TestParameterized]                                | 3                   | null         | []
-    "test-success-with-groups"                                                        | [TestSucceedGroups]                                | 2                   | null         | []
-    "test-class-skipped"                                                              | [TestSkippedClass]                                 | 3                   | null         | []
-    "test-success-and-skipped"                                                        | [TestSucceedAndSkipped]                            | 3                   | null         | []
-    "test-success-and-failure-${version()}"                                           | [TestFailedAndSucceed]                             | 4                   | null         | []
-    "test-suite-teardown-failure"                                                     | [TestFailedSuiteTearDown]                          | 3                   | null         | []
-    "test-suite-setup-failure"                                                        | [TestFailedSuiteSetup]                             | 3                   | null         | []
-    "test-multiple-successful-suites"                                                 | [TestSucceed, TestSucceedAndSkipped]               | 4                   | null         | []
-    "test-successful-suite-and-failed-suite-${version()}"                             | [TestSucceed, TestFailedAndSucceed]                | 5                   | null         | []
-    "test-nested-successful-suites"                                                   | [TestSucceedNested, TestSucceedNested.NestedSuite] | 3                   | null         | []
-    "test-nested-skipped-suites-${version()}"                                         | [TestSkippedNested]                                | 3                   | null         | []
-    "test-factory-data-provider"                                                      | [TestSucceedDataProvider]                          | 2                   | null         | []
-    "test-successful-test-cases-in-parallel"                                          | [TestSucceedMultiple]                              | 3                   | "methods"    | []
-    "test-parameterized-test-cases-in-parallel"                                       | [TestParameterized]                                | 3                   | "methods"    | []
-    "test-itr-skipping"                                                               | [TestFailedAndSucceed]                             | 4                   | null         | [
+    testcaseName                                                                      | expectedTracesCount | parallelMode
+    "test-successful-test-cases-in-TESTS-parallel-mode"                               | 3                   | "tests"
+    "test-successful-test-cases-in-TESTS-parallel-mode-not-all-test-methods-included" | 3                   | "tests"
+    "test-successful-test-cases-in-TESTS-parallel-mode-same-test-case"                | 4                   | "tests"
+  }
+
+  def "test ITR #testcaseName"() {
+    givenSkippableTests(skippedTests)
+    runTests(tests, null)
+
+    assertSpansData(testcaseName, expectedTracesCount)
+
+    where:
+    testcaseName                              | tests                     | expectedTracesCount | skippedTests
+    "test-itr-skipping"                       | [TestFailedAndSucceed]    | 4                   | [
       new TestIdentifier("org.example.TestFailedAndSucceed", "test_another_succeed", null, null),
       new TestIdentifier("org.example.TestFailedAndSucceed", "test_failed", null, null)
     ]
-    "test-itr-skipping-parameterized"                                                 | [TestParameterized]                                | 3                   | null         | [
+    "test-itr-skipping-parameterized"         | [TestParameterized]       | 3                   | [
       new TestIdentifier("org.example.TestParameterized", "parameterized_test_succeed", '{"arguments":{"0":"hello","1":"true"}}', null)
     ]
-    "test-itr-skipping-factory-data-provider"                                         | [TestSucceedDataProvider]                          | 2                   | null         | [new TestIdentifier("org.example.TestSucceedDataProvider", "testMethod", null, null)]
-    "test-itr-unskippable"                                                            | [TestSucceedUnskippable]                           | 2                   | null         | [new TestIdentifier("org.example.TestSucceedUnskippable", "test_succeed", null, null)]
-    "test-itr-unskippable-not-skipped"                                                | [TestSucceedUnskippable]                           | 2                   | null         | []
-    "test-successful-test-cases-in-TESTS-parallel-mode"                               | []                                                 | 3                   | "tests"      | []
-    "test-successful-test-cases-in-TESTS-parallel-mode-not-all-test-methods-included" | []                                                 | 3                   | "tests"      | []
-    "test-successful-test-cases-in-TESTS-parallel-mode-same-test-case"                | []                                                 | 4                   | "tests"      | []
+    "test-itr-skipping-factory-data-provider" | [TestSucceedDataProvider] | 2                   | [new TestIdentifier("org.example.TestSucceedDataProvider", "testMethod", null, null)]
+    "test-itr-unskippable"                    | [TestSucceedUnskippable]  | 2                   | [new TestIdentifier("org.example.TestSucceedUnskippable", "test_succeed", null, null)]
+    "test-itr-unskippable-not-skipped"        | [TestSucceedUnskippable]  | 2                   | []
+  }
+
+  def "test flaky retries #testcaseName"() {
+    givenFlakyTests(retriedTests)
+    runTests(tests, null)
+
+    assertSpansData(testcaseName, expectedTracesCount)
+
+    where:
+    testcaseName                            | tests                     | expectedTracesCount | retriedTests
+    "test-failed-${version()}"              | [TestFailed]              | 2                   | []
+    "test-skipped"                          | [TestSkipped]             | 2                   | [new TestIdentifier("org.example.TestSkipped", "test_skipped", null, null)]
+    "test-retry-failed-${version()}"        | [TestFailed]              | 6                   | [new TestIdentifier("org.example.TestFailed", "test_failed", null, null)]
+    "test-retry-error"                      | [TestError]               | 6                   | [new TestIdentifier("org.example.TestError", "test_error", null, null)]
+    "test-retry-parameterized"              | [TestFailedParameterized] | 7                   | [new TestIdentifier("org.example.TestFailedParameterized", "parameterized_test_succeed", null, null)]
+    "test-failed-then-succeed-${version()}" | [TestFailedThenSucceed]   | 4                   | [new TestIdentifier("org.example.TestFailedThenSucceed", "test_failed", null, null)]
   }
 
   protected void runTests(List<Class> testClasses, String parallelMode = null) {

dd-java-agent/instrumentation/testng/src/testFixtures/java/org/example/TestFailedParameterized.java

@@ -0,0 +1,22 @@
+package org.example;
+
+import static org.testng.AssertJUnit.assertTrue;
+
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+public class TestFailedParameterized {
+
+  @DataProvider(name = "dataProvider")
+  public static Object[][] data() {
+    return new Object[][] {
+      {"hello", true},
+      {"\"goodbye\"", false}
+    };
+  }
+
+  @Test(dataProvider = "dataProvider")
+  public void parameterized_test_succeed(final String str, final boolean booleanValue) {
+    assertTrue(booleanValue);
+  }
+}

dd-java-agent/instrumentation/testng/src/testFixtures/java/org/example/TestFailedThenSucceed.java

@@ -0,0 +1,15 @@
+package org.example;
+
+import static org.testng.Assert.assertTrue;
+
+import org.testng.annotations.Test;
+
+public class TestFailedThenSucceed {
+
+  private int retry;
+
+  @Test
+  public void test_failed() {
+    assertTrue(++retry >= 3);
+  }
+}