Implement ITR tests skipping for TestNG (#5453)

Author: nikita-tkachenko-datadog

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

@@ -0,0 +1,37 @@
+package datadog.trace.instrumentation.testng;
+
+import datadog.trace.api.Config;
+import datadog.trace.api.civisibility.config.SkippableTest;
+import java.lang.reflect.Method;
+import java.util.Set;
+
+public class ItrFilter {
+
+  public static final ItrFilter INSTANCE = new ItrFilter();
+
+  private volatile boolean testsSkipped;
+
+  private ItrFilter() {}
+
+  public boolean skip(Method method, Object instance, Object[] parameters) {
+    SkippableTest test = toSkippableTest(method, instance, parameters);
+    Set<SkippableTest> skippableTests = Config.get().getCiVisibilitySkippableTests();
+    if (skippableTests.contains(test)) {
+      testsSkipped = true;
+      return true;
+    } else {
+      return false;
+    }
+  }
+
+  private SkippableTest toSkippableTest(Method method, Object instance, Object[] parameters) {
+    String testSuiteName = instance.getClass().getName();
+    String testName = method.getName();
+    String testParameters = TestNGUtils.getParameters(parameters);
+    return new SkippableTest(testSuiteName, testName, testParameters, null);
+  }
+
+  public boolean testsSkipped() {
+    return testsSkipped;
+  }
+}

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

@@ -35,6 +35,7 @@ public String[] helperClassNames() {
       packageName + ".TestNGUtils",
       packageName + ".TestNGSuiteListener",
       packageName + ".TestNGClassListener",
+      packageName + ".ItrFilter",
       packageName + ".TracingListener"
     };
   }

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

@@ -0,0 +1,72 @@
+package datadog.trace.instrumentation.testng;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.api.Config;
+import java.lang.reflect.Method;
+import java.util.Set;
+import net.bytebuddy.asm.Advice;
+import org.testng.SkipException;
+import org.testng.annotations.DataProvider;
+
+@AutoService(Instrumenter.class)
+public class TestNGItrInstrumentation extends Instrumenter.CiVisibility
+    implements Instrumenter.ForKnownTypes {
+  public TestNGItrInstrumentation() {
+    super("testng", "testng-itr");
+  }
+
+  @Override
+  public boolean isApplicable(Set<TargetSystem> enabledSystems) {
+    return super.isApplicable(enabledSystems) && Config.get().isCiVisibilityItrEnabled();
+  }
+
+  @Override
+  public String[] knownMatchingTypes() {
+    return new String[] {
+      "org.testng.internal.MethodInvocationHelper",
+      "org.testng.internal.invokers.MethodInvocationHelper"
+    };
+  }
+
+  @Override
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
+        named("invokeMethod")
+            .and(takesArguments(3))
+            .and(takesArgument(0, Method.class))
+            .and(takesArgument(1, Object.class))
+            .and(takesArgument(2, Object[].class)),
+        TestNGItrInstrumentation.class.getName() + "$InvokeMethodAdvice");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".TestNGUtils",
+      packageName + ".TestNGClassListener",
+      packageName + ".ItrFilter",
+    };
+  }
+
+  public static class InvokeMethodAdvice {
+    @Advice.OnMethodEnter
+    public static void invokeMethod(
+        @Advice.Argument(0) final Method method,
+        @Advice.Argument(1) final Object instance,
+        @Advice.Argument(2) final Object[] parameters) {
+      if (ItrFilter.INSTANCE.skip(method, instance, parameters)) {
+        throw new SkipException("Skipped by Datadog Intelligent Test Runner");
+      }
+    }
+
+    // TestNG 6.4 and above
+    public static void muzzleCheck(final DataProvider dataProvider) {
+      dataProvider.name();
+    }
+  }
+}

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

@@ -40,20 +40,24 @@ public static Method getTestMethod(final ITestResult result) {
   }
 
   public static String getParameters(final ITestResult result) {
-    if (result.getParameters() == null || result.getParameters().length == 0) {
+    return getParameters(result.getParameters());
+  }
+
+  public static String getParameters(Object[] parameters) {
+    if (parameters == null || parameters.length == 0) {
       return null;
     }
 
     // We build manually the JSON for test.parameters tag.
     // Example: {"arguments":{"0":"param1","1":"param2"}}
     final StringBuilder sb = new StringBuilder("{\"arguments\":{");
-    for (int i = 0; i < result.getParameters().length; i++) {
+    for (int i = 0; i < parameters.length; i++) {
       sb.append("\"")
           .append(i)
           .append("\":\"")
-          .append(Strings.escapeToJson(String.valueOf(result.getParameters()[i])))
+          .append(Strings.escapeToJson(String.valueOf(parameters[i])))
           .append("\"");
-      if (i != result.getParameters().length - 1) {
+      if (i != parameters.length - 1) {
         sb.append(",");
       }
     }

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

@@ -41,7 +41,7 @@ public void onExecutionStart() {
 
   @Override
   public void onExecutionFinish() {
-    testEventsHandler.onTestModuleFinish(false);
+    testEventsHandler.onTestModuleFinish(ItrFilter.INSTANCE.testsSkipped());
   }
 
   @Override

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

@@ -1,6 +1,9 @@
 package datadog.trace.instrumentation.testng
 
 import datadog.trace.agent.test.asserts.ListWriterAssert
+import datadog.trace.api.civisibility.config.SkippableTest
+import datadog.trace.api.civisibility.config.SkippableTestsSerializer
+import datadog.trace.api.config.CiVisibilityConfig
 import datadog.trace.civisibility.CiVisibilityTest
 import datadog.trace.api.civisibility.CIConstants
 import datadog.trace.bootstrap.instrumentation.api.Tags
@@ -659,6 +662,76 @@ abstract class TestNGTest extends CiVisibilityTest {
     })
   }
 
+  def "test ITR skipping"() {
+    setup:
+    injectSysConfig(CiVisibilityConfig.CIVISIBILITY_SKIPPABLE_TESTS, SkippableTestsSerializer.serialize([
+      new SkippableTest("org.example.TestFailedAndSucceed", "test_another_succeed", null, null),
+      new SkippableTest("org.example.TestFailedAndSucceed", "test_failed", null, null),
+    ]))
+
+    def testNG = new TestNG()
+    testNG.setTestClasses(TestFailedAndSucceed)
+    testNG.setOutputDirectory(testOutputDir)
+    testNG.run()
+
+    expect:
+    ListWriterAssert.assertTraces(TEST_WRITER, 4, false, SORT_TRACES_BY_DESC_SIZE_THEN_BY_NAMES, {
+      long testModuleId
+      long testSuiteId
+      trace(2, true) {
+        testModuleId = testModuleSpan(it, 0, CIConstants.TEST_PASS)
+        testSuiteId = testSuiteSpan(it, 1, testModuleId, "org.example.TestFailedAndSucceed", CIConstants.TEST_PASS)
+      }
+      trace(1) {
+        testSpan(it, 0, testModuleId, testSuiteId, "org.example.TestFailedAndSucceed", "test_another_succeed", CIConstants.TEST_SKIP, testTags)
+      }
+      trace(1) {
+        testSpan(it, 0, testModuleId, testSuiteId, "org.example.TestFailedAndSucceed", "test_failed", CIConstants.TEST_SKIP, testTags)
+      }
+      trace(1) {
+        testSpan(it, 0, testModuleId, testSuiteId, "org.example.TestFailedAndSucceed", "test_succeed", CIConstants.TEST_PASS)
+      }
+    })
+
+    where:
+    testTags = [(Tags.TEST_SKIP_REASON): "Skipped by Datadog Intelligent Test Runner"]
+  }
+
+  def "test ITR skipping for parameterized tests"() {
+    setup:
+    injectSysConfig(CiVisibilityConfig.CIVISIBILITY_SKIPPABLE_TESTS, SkippableTestsSerializer.serialize([
+      new SkippableTest("org.example.TestParameterized", "parameterized_test_succeed", testTags_0[Tags.TEST_PARAMETERS], null),
+    ]))
+
+    def testNG = new TestNG()
+    testNG.setTestClasses(TestParameterized)
+    testNG.setOutputDirectory(testOutputDir)
+    testNG.run()
+
+    expect:
+    ListWriterAssert.assertTraces(TEST_WRITER, 3, false, SORT_TRACES_BY_DESC_SIZE_THEN_BY_NAMES, {
+      long testModuleId
+      long testSuiteId
+      trace(2, true) {
+        testModuleId = testModuleSpan(it, 0, CIConstants.TEST_PASS)
+        testSuiteId = testSuiteSpan(it, 1, testModuleId, "org.example.TestParameterized", CIConstants.TEST_PASS)
+      }
+      trace(1) {
+        testSpan(it, 0, testModuleId, testSuiteId, "org.example.TestParameterized", "parameterized_test_succeed", CIConstants.TEST_SKIP, testTags_0)
+      }
+      trace(1) {
+        testSpan(it, 0, testModuleId, testSuiteId, "org.example.TestParameterized", "parameterized_test_succeed", CIConstants.TEST_PASS, testTags_1)
+      }
+    })
+
+    where:
+    testTags_0 = [
+      (Tags.TEST_PARAMETERS): '{"arguments":{"0":"hello","1":"true"}}',
+      (Tags.TEST_SKIP_REASON): "Skipped by Datadog Intelligent Test Runner"
+    ]
+    testTags_1 = [(Tags.TEST_PARAMETERS): '{"arguments":{"0":"\\\"goodbye\\\"","1":"false"}}']
+  }
+
   @Override
   String expectedOperationPrefix() {
     return "testng"