DataProvider: possibility to dynamically load dataprovider class

With current Data providers implementation, it's code will stick around in method area (JVM spec $2.5.4) for entire test run. By specifying dataprovider class with it's full qualified name, and by using new custom classloader to load it, when needed, JVM gets a chance to unload dataprovider class, when we're done and deleted all links to it Test dataprovider class unloaded by analysing memory dumps.
testng-team · Mar 12, 2022 · a0b7229 · a0b7229
1 parent 331bfbe
commit a0b7229
Show file tree

Hide file tree

Showing 20 changed files with 305 additions and 6 deletions.
diff --git a/testng-core-api/src/main/java/org/testng/annotations/ITestAnnotation.java b/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();

diff --git a/testng-core-api/src/main/java/org/testng/annotations/Test.java b/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.

diff --git a/testng-core-api/src/main/java/org/testng/internal/annotations/IDataProvidable.java b/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);
 }
diff --git a/testng-core/src/main/java/org/testng/internal/DataProviderLoader.java b/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;
+  }
+}
diff --git a/testng-core/src/main/java/org/testng/internal/DataProviderMethod.java b/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) {

diff --git a/testng-core/src/main/java/org/testng/internal/DataProviderMethodRemovable.java b/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;
+  }
+}
diff --git a/testng-core/src/main/java/org/testng/internal/Parameters.java b/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()) {

diff --git a/testng-core/src/main/java/org/testng/internal/annotations/FactoryAnnotation.java b/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;

diff --git a/testng-core/src/main/java/org/testng/internal/annotations/JDK15AnnotationFinder.java b/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);

diff --git a/testng-core/src/main/java/org/testng/internal/annotations/JDK15TagFactory.java b/testng-core/src/main/java/org/testng/internal/annotations/JDK15TagFactory.java
@@ -481,6 +481,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));

diff --git a/testng-core/src/main/java/org/testng/internal/annotations/TestAnnotation.java b/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;

diff --git a/testng-core/src/test/kotlin/org/testng/dataprovider/DynamicDataProviderLoadingTest.kt b/testng-core/src/test/kotlin/org/testng/dataprovider/DynamicDataProviderLoadingTest.kt
@@ -0,0 +1,85 @@
+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.annotations.Test
+import org.testng.dataprovider.sample.issue2724.TestClass
+import org.testng.dataprovider.sample.issue2724.TestDPUnloaded
+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(TestClass::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")
+        System.setProperty("memdump.path", dumpPath)
+
+        run(TestDPUnloaded::class.java)
+
+        val heapDumpFile = File(dumpPath)
+        assertThat(heapDumpFile).exists()
+        val 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()
+        }
+    }
+
+    fun getGCPath(instance: Instance) : String {
+        var result = ""
+        if (!instance.isGCRoot) {
+            result += getGCPath(instance.nearestGCRootPointer)
+        }
+        return result + "${instance.javaClass.name}\n"
+    }
+
+    companion object {
+        var classLoadCount = 0
+    }
+}
diff --git a/testng-core/src/test/kotlin/org/testng/dataprovider/sample/issue2724/ClassLoadedCounter.kt b/testng-core/src/test/kotlin/org/testng/dataprovider/sample/issue2724/ClassLoadedCounter.kt
@@ -0,0 +1,11 @@
+package org.testng.dataprovider.sample.issue2724
+
+import org.testng.dataprovider.DynamicDataProviderLoadingTest
+
+class ClassLoadedCounter {
+    companion object {
+        init {
+            DynamicDataProviderLoadingTest.classLoadCount++
+        }
+    }
+}
diff --git a/testng-core/src/test/kotlin/org/testng/dataprovider/sample/issue2724/DataProviders.kt b/testng-core/src/test/kotlin/org/testng/dataprovider/sample/issue2724/DataProviders.kt
@@ -0,0 +1,14 @@
+package org.testng.dataprovider.sample.issue2724
+
+import org.testng.annotations.DataProvider
+
+class DataProviders {
+    @DataProvider
+    fun data() : Array<Array<Any>> {
+        return arrayOf(
+            arrayOf("Mike", 34, "student"),
+            arrayOf("Mike", 23, "driver"),
+            arrayOf("Paul", 20, "director")
+        )
+    }
+}
diff --git a/testng-core/src/test/kotlin/org/testng/dataprovider/sample/issue2724/TestClass.kt b/testng-core/src/test/kotlin/org/testng/dataprovider/sample/issue2724/TestClass.kt
@@ -0,0 +1,15 @@
+package org.testng.dataprovider.sample.issue2724
+
+import org.testng.annotations.Test
+
+class TestClass {
+
+    @Suppress("UNUSED")
+    @Test(
+        dataProviderDynamicClass = "org.testng.dataprovider.sample.issue2724.DataProviders",
+        dataProvider = "data"
+    )
+    fun testDynamicDataProvider(name: String, age: Int, status: String) {
+
+    }
+}