DataProvider: possibility to unload dataprovider class, when done with it

CHANGES.txt

@@ -1,5 +1,6 @@
 Current
 7.6.0
+New: GITHUB-2724: DataProvider: possibility to unload dataprovider class, when done with it (Dzmitry Sankouski)
 Fixed: GITHUB-217: Configure TestNG to fail when there's a failure in data provider (Krishnan Mahadevan)
 Fixed: GITHUB-2743: SuiteRunner could not be initial by default Configuration (Nan Liang)
 Fixed: GITHUB-2729: beforeConfiguration() listener method should be invoked for skipped configurations as well(Nan Liang)

testng-core-api/src/main/java/org/testng/annotations/ITestAnnotation.java

@@ -72,6 +72,10 @@ public interface ITestAnnotation extends ITestOrConfiguration, IDataProvidable {
 
   void setDataProviderClass(Class<?> v);
 
+  String getDataProviderDynamicClass();
+
+  void setDataProviderDynamicClass(String v);
+
   void setRetryAnalyzer(Class<? extends IRetryAnalyzer> c);
 
   Class<? extends IRetryAnalyzer> getRetryAnalyzerClass();

testng-core-api/src/main/java/org/testng/annotations/Test.java

@@ -104,6 +104,8 @@
    */
   Class<?> dataProviderClass() default Object.class;
 
+  String dataProviderDynamicClass() default "";
+
   /**
    * If set to true, this test method will always be run even if it depends on a method that failed.
    * This attribute will be ignored if this test doesn't depend on any method or group.

testng-core-api/src/main/java/org/testng/internal/annotations/IDataProvidable.java

@@ -9,4 +9,8 @@ public interface IDataProvidable {
   Class<?> getDataProviderClass();
 
   void setDataProviderClass(Class<?> v);
+
+  String getDataProviderDynamicClass();
+
+  void setDataProviderDynamicClass(String v);
 }

testng-core/src/main/java/org/testng/internal/DataProviderLoader.java

@@ -0,0 +1,43 @@
+package org.testng.internal;
+
+import java.io.IOException;
+import java.io.InputStream;
+import org.testng.log4testng.Logger;
+
+public class DataProviderLoader extends ClassLoader {
+  private static final int BUFFER_SIZE = 1 << 20;
+  private static final Logger log = Logger.getLogger(DataProviderLoader.class);
+
+  public Class loadClazz(String path) throws ClassNotFoundException {
+    Class clazz = findLoadedClass(path);
+    if (clazz == null) {
+      byte[] bt = loadClassData(path);
+      clazz = defineClass(path, bt, 0, bt.length);
+    }
+
+    return clazz;
+  }
+
+  private byte[] loadClassData(String className) throws ClassNotFoundException {
+    InputStream in =
+        this.getClass()
+            .getClassLoader()
+            .getResourceAsStream(className.replace(".", "/") + ".class");
+    if (in == null) {
+      throw new ClassNotFoundException("Cannot load resource input stream: " + className);
+    }
+
+    byte[] classBytes;
+    try {
+      classBytes = in.readAllBytes();
+    } catch (IOException e) {
+      throw new ClassNotFoundException("ERROR reading class file" + e);
+    }
+
+    if (classBytes == null) {
+      throw new ClassNotFoundException("Cannot load class: " + className);
+    }
+
+    return classBytes;
+  }
+}

testng-core/src/main/java/org/testng/internal/DataProviderMethod.java

@@ -8,8 +8,8 @@
 /** Represents an @{@link org.testng.annotations.DataProvider} annotated method. */
 class DataProviderMethod implements IDataProviderMethod {
 
-  private final Object instance;
-  private final Method method;
+  protected Object instance;
+  protected Method method;
   private final IDataProviderAnnotation annotation;
 
   DataProviderMethod(Object instance, Method method, IDataProviderAnnotation annotation) {

testng-core/src/main/java/org/testng/internal/DataProviderMethodRemovable.java

@@ -0,0 +1,20 @@
+package org.testng.internal;
+
+import java.lang.reflect.Method;
+import org.testng.annotations.IDataProviderAnnotation;
+
+/** Represents an @{@link org.testng.annotations.DataProvider} annotated method. */
+class DataProviderMethodRemovable extends DataProviderMethod {
+
+  DataProviderMethodRemovable(Object instance, Method method, IDataProviderAnnotation annotation) {
+    super(instance, method, annotation);
+  }
+
+  public void setInstance(Object instance) {
+    this.instance = instance;
+  }
+
+  public void setMethod(Method method) {
+    this.method = method;
+  }
+}

testng-core/src/main/java/org/testng/internal/Parameters.java

@@ -495,6 +495,15 @@ private static IDataProviderMethod findDataProvider(
     if (dp != null) {
       String dataProviderName = dp.getDataProvider();
       Class<?> dataProviderClass = dp.getDataProviderClass();
+      boolean isDynamicDataProvider =
+          dataProviderClass == null && !dp.getDataProviderDynamicClass().isEmpty();
+      if (isDynamicDataProvider) {
+        try {
+          dataProviderClass = new DataProviderLoader().loadClazz(dp.getDataProviderDynamicClass());
+        } catch (ClassNotFoundException e) {
+          throw new TestNGException("Dynamic data provider class %s not found", e);
+        }
+      }
 
       if (!Utils.isStringEmpty(dataProviderName)) {
         result =
@@ -505,6 +514,7 @@ private static IDataProviderMethod findDataProvider(
                 finder,
                 dataProviderName,
                 dataProviderClass,
+                isDynamicDataProvider,
                 context);
 
         if (null == result) {
@@ -566,6 +576,10 @@ private static IDataProvidable merge(ITestAnnotation methodLevel, ITestAnnotatio
     if (isDataProviderClassEmpty(methodLevel) && !isDataProviderClassEmpty(classLevel)) {
       methodLevel.setDataProviderClass(classLevel.getDataProviderClass());
     }
+    if (isDynamicDataProviderClassEmpty(methodLevel)
+        && !isDynamicDataProviderClassEmpty(classLevel)) {
+      methodLevel.setDataProviderDynamicClass(classLevel.getDataProviderDynamicClass());
+    }
     return methodLevel;
   }
 
@@ -574,6 +588,10 @@ private static boolean isDataProviderClassEmpty(ITestAnnotation annotation) {
         || Object.class.equals(annotation.getDataProviderClass());
   }
 
+  private static boolean isDynamicDataProviderClassEmpty(ITestAnnotation annotation) {
+    return annotation.getDataProviderDynamicClass().isEmpty();
+  }
+
   private static boolean isDataProviderNameEmpty(ITestAnnotation annotation) {
     return Strings.isNullOrEmpty(annotation.getDataProvider());
   }
@@ -586,6 +604,7 @@ private static IDataProviderMethod findDataProvider(
       IAnnotationFinder finder,
       String name,
       Class<?> dataProviderClass,
+      boolean isDynamicDataProvider,
       ITestContext context) {
     IDataProviderMethod result = null;
 
@@ -620,7 +639,12 @@ private static IDataProviderMethod findDataProvider(
         if (result != null) {
           throw new TestNGException("Found two providers called '" + name + "' on " + cls);
         }
-        result = new DataProviderMethod(instanceToUse, m, dp);
+
+        if (isDynamicDataProvider) {
+          result = new DataProviderMethodRemovable(instanceToUse, m, dp);
+        } else {
+          result = new DataProviderMethod(instanceToUse, m, dp);
+        }
       }
     }
 
@@ -839,6 +863,14 @@ public void remove() {
                 filteredParameters, dataProviderMethod, testMethod, methodParams.context);
       }
 
+      if (dataProviderMethod instanceof DataProviderMethodRemovable) {
+        ((DataProviderMethodRemovable) dataProviderMethod).setMethod(null);
+        ((DataProviderMethodRemovable) dataProviderMethod).setInstance(null);
+        if (testMethod instanceof TestNGMethod) {
+          ((TestNGMethod) testMethod).setDataProviderMethod(null);
+        }
+      }
+
       return new ParameterHolder(
           filteredParameters, ParameterOrigin.ORIGIN_DATA_PROVIDER, dataProviderMethod);
     } else if (methodParams.xmlParameters.isEmpty()) {

testng-core/src/main/java/org/testng/internal/annotations/FactoryAnnotation.java

@@ -8,6 +8,7 @@ public class FactoryAnnotation extends BaseAnnotation implements IFactoryAnnotat
 
   private String m_dataProvider = null;
   private Class<?> m_dataProviderClass;
+  private String m_dataProviderDynamicClass;
   private boolean m_enabled = true;
   private List<Integer> m_indices;
 
@@ -30,6 +31,16 @@ public Class<?> getDataProviderClass() {
     return m_dataProviderClass;
   }
 
+  @Override
+  public String getDataProviderDynamicClass() {
+    return m_dataProviderDynamicClass;
+  }
+
+  @Override
+  public void setDataProviderDynamicClass(String v) {
+    m_dataProviderDynamicClass = v;
+  }
+
   @Override
   public boolean getEnabled() {
     return m_enabled;

testng-core/src/main/java/org/testng/internal/annotations/JDK15AnnotationFinder.java

@@ -48,7 +48,7 @@ public class JDK15AnnotationFinder implements IAnnotationFinder {
   private final JDK15TagFactory m_tagFactory = new JDK15TagFactory();
   private final Map<Class<? extends IAnnotation>, Class<? extends Annotation>> m_annotationMap =
       new ConcurrentHashMap<>();
-  private final Map<Pair<Annotation, ?>, IAnnotation> m_annotations = new ConcurrentHashMap<>();
+  private final Map<String, IAnnotation> m_annotations = new ConcurrentHashMap<>();
 
   private final IAnnotationTransformer m_transformer;
 
@@ -273,7 +273,7 @@ private <A extends IAnnotation> A findAnnotation(
 
     IAnnotation result =
         m_annotations.computeIfAbsent(
-            p,
+            p.toString(),
             key -> {
               IAnnotation obj = m_tagFactory.createTag(cls, testMethod, a, annotationClass);
               transform(obj, testClass, testConstructor, testMethod, whichClass);

testng-core/src/main/java/org/testng/internal/annotations/JDK15TagFactory.java

@@ -484,6 +484,7 @@ private IAnnotation createTestTag(Class<?> cls, Annotation a) {
     result.setDataProviderClass(
         findInherited(
             test.dataProviderClass(), cls, Test.class, "dataProviderClass", DEFAULT_CLASS));
+    result.setDataProviderDynamicClass(test.dataProviderDynamicClass());
     result.setAlwaysRun(test.alwaysRun());
     result.setDescription(
         findInherited(test.description(), cls, Test.class, "description", DEFAULT_STRING));

testng-core/src/main/java/org/testng/internal/annotations/TestAnnotation.java

@@ -19,6 +19,7 @@ public class TestAnnotation extends TestOrConfiguration implements ITestAnnotati
   private String m_testName = "";
   private boolean m_singleThreaded = false;
   private Class<?> m_dataProviderClass = null;
+  private String m_dataProviderDynamicClass = null;
   private Class<? extends IRetryAnalyzer> m_retryAnalyzerClass = null;
   private boolean m_skipFailedInvocations = false;
   private boolean m_ignoreMissingDependencies = false;
@@ -66,6 +67,16 @@ public void setDataProviderClass(Class<?> dataProviderClass) {
     m_dataProviderClass = dataProviderClass;
   }
 
+  @Override
+  public String getDataProviderDynamicClass() {
+    return m_dataProviderDynamicClass;
+  }
+
+  @Override
+  public void setDataProviderDynamicClass(String v) {
+    m_dataProviderDynamicClass = v;
+  }
+
   @Override
   public void setInvocationCount(int invocationCount) {
     m_invocationCount = invocationCount;

testng-core/src/test/kotlin/org/testng/dataprovider/DynamicDataProviderLoadingTest.kt

@@ -0,0 +1,126 @@
+package org.testng.dataprovider
+
+import org.assertj.core.api.Assertions.assertThat
+import org.assertj.core.api.SoftAssertions
+import org.netbeans.lib.profiler.heap.HeapFactory2
+import org.netbeans.lib.profiler.heap.Instance
+import org.netbeans.lib.profiler.heap.JavaClass
+import org.testng.Reporter
+import org.testng.annotations.Test
+import org.testng.dataprovider.sample.issue2724.*
+import test.SimpleBaseTest
+import java.io.File
+import java.nio.file.Files
+
+const val CLASS_NAME_DP = "org.testng.dataprovider.sample.issue2724.DataProviders"
+const val CLASS_NAME_DP_LOADER = "org.testng.internal.DataProviderLoader"
+
+class DynamicDataProviderLoadingTest : SimpleBaseTest() {
+
+    @Test
+    fun testDynamicDataProviderPasses() {
+        val listener = run(SampleDynamicDP::class.java)
+        assertThat(listener.failedMethodNames).isEmpty()
+        assertThat(listener.succeedMethodNames).containsExactly(
+            "testDynamicDataProvider(Mike,34,student)",
+            "testDynamicDataProvider(Mike,23,driver)",
+            "testDynamicDataProvider(Paul,20,director)",
+        )
+        assertThat(listener.skippedMethodNames).isEmpty()
+    }
+
+    @Test
+    fun testDynamicDataProviderUnloaded() {
+        val tempDirectory = Files.createTempDirectory("temp-testng-")
+        val dumpPath = "%s/%s".format(tempDirectory.toAbsolutePath().toString(), "dump.hprof")
+        val dumpPathBeforeSample =
+            "%s/%s".format(tempDirectory.toAbsolutePath().toString(), "dump-before-sample.hprof")
+        System.setProperty("memdump.path", dumpPath)
+
+        saveMemDump(dumpPathBeforeSample)
+        val heapDumpBeforeSampleFile = File(dumpPathBeforeSample)
+        assertThat(heapDumpBeforeSampleFile).exists()
+        var heap = HeapFactory2.createHeap(heapDumpBeforeSampleFile, null)
+        val beforeSampleDPClassDump: JavaClass? = heap.getJavaClassByName(CLASS_NAME_DP)
+        assertThat(beforeSampleDPClassDump)
+            .describedAs(
+                "Class $CLASS_NAME_DP shouldn't be loaded, before test sample started. "
+            )
+            .isNull()
+
+        run(SampleDPUnloaded::class.java)
+
+        val heapDumpFile = File(dumpPath)
+        assertThat(heapDumpFile).exists()
+        heap = HeapFactory2.createHeap(heapDumpFile, null)
+
+        with(SoftAssertions()) {
+            val dpLoaderClassDump: JavaClass? = heap.getJavaClassByName(CLASS_NAME_DP_LOADER)
+            val dpClassDump: JavaClass? = heap.getJavaClassByName(CLASS_NAME_DP)
+            val dpLoaderMessage = dpLoaderClassDump?.instances?.joinToString("\n") {
+                getGCPath(it)
+            }
+            val dpMessage = dpLoaderClassDump?.instances?.joinToString("\n") {
+                getGCPath(it)
+            }
+
+            this.assertThat(dpLoaderClassDump?.instances)
+                .describedAs(
+                    """
+                    All instances of class $CLASS_NAME_DP_LOADER should be garbage collected, but was not.
+                    Path to GC root is:
+                    $dpLoaderMessage
+                    """.trimIndent()
+                )
+                .isEmpty()
+            this.assertThat(dpClassDump)
+                .describedAs(
+                    """
+                    Class $CLASS_NAME_DP shouldn't be loaded, but it was.
+                    Path to GC root is:
+                    $dpMessage
+                    """.trimIndent()
+                )
+                .isNull()
+            this.assertAll()
+        }
+    }
+
+    @Test
+    fun comparePerformanceAgainstCsvFiles() {
+        val simpleDPSuite = create().apply {
+            setTestClasses(arrayOf(SampleSimpleDP::class.java))
+            setListenerClasses(listOf(TestTimeListener::class.java))
+        }
+        val csvSuite = create().apply {
+            setTestClasses(arrayOf(SampleWithCSVData::class.java))
+            setListenerClasses(listOf(TestTimeListener::class.java))
+        }
+        val dataAsCodeSuite = create().apply {
+            setTestClasses(arrayOf(SampleDynamicDP::class.java))
+            setListenerClasses(listOf(TestTimeListener::class.java))
+        }
+
+        Reporter.log("Test execution time:\n")
+        for (suite in listOf(
+            Pair("simple dataprovider", simpleDPSuite),
+            Pair("dataprovider as code", dataAsCodeSuite),
+            Pair("csv dataprovider", csvSuite),
+        )) {
+            run(false, suite.second)
+            Reporter.log(
+                "${suite.first} execution times: %d milliseconds."
+                    .format(TestTimeListener.testRunTime),
+                true
+            )
+        }
+    }
+
+    fun getGCPath(instance: Instance): String {
+        var result = ""
+        if (!instance.isGCRoot) {
+            result += getGCPath(instance.nearestGCRootPointer)
+        }
+        return result + "${instance.javaClass.name}\n"
+    }
+}