[Discussion] Implementation of statically typed data providers

build.gradle.kts

@@ -68,7 +68,7 @@ dependencies {
     }
 
     listOf("com.beust:jcommander:1.72", "com.google.inject:guice:4.1.0:no_aop",
-            "org.yaml:snakeyaml:1.21").forEach {
+            "org.yaml:snakeyaml:1.21", "net.bytebuddy:byte-buddy:1.10.1").forEach {
         compile(it)
     }
 

src/main/java/org/testng/annotations/ParameterCollector.java

@@ -0,0 +1,16 @@
+package org.testng.annotations;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.ElementType.PARAMETER;
+
+/**
+ * If this annotation is used on a parameter of a data provider, that parameter is the proxy to the
+ * class which uses data provider. Call methods of that proxy to provide parameters for a test.
+ *
+ * <p>This annotation is ignored everywhere else.
+ */
+@Retention(java.lang.annotation.RetentionPolicy.RUNTIME)
+@Target({PARAMETER})
+public @interface ParameterCollector {}

src/main/java/org/testng/internal/MethodInvocationHelper.java

@@ -1,6 +1,10 @@
 package org.testng.internal;
 
 import java.util.concurrent.TimeUnit;
+
+import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.implementation.InvocationHandlerAdapter;
+import net.bytebuddy.matcher.ElementMatchers;
 import org.testng.IConfigurable;
 import org.testng.IConfigureCallBack;
 import org.testng.IHookCallBack;
@@ -31,6 +35,7 @@
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * Collections of helper methods to help deal with invocation of TestNG methods
@@ -141,11 +146,19 @@ protected static Iterator<Object[]> invokeDataProvider(
       ITestContext testContext,
       Object fedInstance,
       IAnnotationFinder annotationFinder) {
+    List<Object[]> resultHolder = new ArrayList<>();
     List<Object> parameters =
-        getParameters(dataProvider, method, testContext, fedInstance, annotationFinder);
+        getParameters(dataProvider, method, testContext, fedInstance, resultHolder, annotationFinder);
     Object result = invokeMethodNoCheckedException(dataProvider, instance, parameters);
     if (result == null) {
+      if (!resultHolder.isEmpty()) {
+        return resultHolder.iterator();
+      }
       throw new TestNGException("Data Provider " + dataProvider + " returned a null value");
+    } else {
+      if (!resultHolder.isEmpty()) {
+        throw new TestNGException("Data Provider " + dataProvider + " must return void if using @ParameterCollector");
+      }
     }
     // If it returns an Object[][] or Object[], convert it to an Iterator<Object[]>
     if (result instanceof Object[][]) {
@@ -180,6 +193,7 @@ private static List<Object> getParameters(
       ITestNGMethod method,
       ITestContext testContext,
       Object fedInstance,
+      List<Object[]> resultHolder,
       IAnnotationFinder annotationFinder) {
     // Go through all the parameters declared on this Data Provider and
     // make sure we have at most one Method and one ITestContext.
@@ -203,8 +217,32 @@ private static List<Object> getParameters(
         parameters.add(com.getDeclaringClass());
       } else {
         boolean isTestInstance = annotationFinder.hasTestInstance(dataProvider, i);
+        boolean isParameterCollector = annotationFinder.hasParameterCollector(dataProvider, i);
         if (isTestInstance) {
           parameters.add(fedInstance);
+        } else if (isParameterCollector) {
+          final String expectedMethodName = method.getMethodName();
+          Class<?> proxyClass = new ByteBuddy()
+                  .subclass(fedInstance.getClass())
+                  .method(ElementMatchers.any())
+                  .intercept(InvocationHandlerAdapter.of((proxy, calledMethod, args) -> {
+                    final String actualMethodName = calledMethod.getName();
+                    if(!Objects.equals(expectedMethodName, actualMethodName)) {
+                      throw new TestNGException(
+                              String.format("Wrong method %s is called, expected %s",
+                                      actualMethodName, expectedMethodName));
+                    }
+                    resultHolder.add(args);
+                    return null;
+                  }))
+                  .make()
+                  .load(fedInstance.getClass().getClassLoader())
+                  .getLoaded();
+          try {
+            parameters.add(proxyClass.newInstance());
+          } catch (InstantiationException | IllegalAccessException e) {
+            throw new TestNGException("Failed to proxy class", e);
+          }
         } else {
           unresolved.add(new Pair<>(i, cls));
         }

src/main/java/org/testng/internal/annotations/IAnnotationFinder.java

@@ -46,6 +46,10 @@ <A extends IAnnotation> A findAnnotation(
   /** @return true if the ith parameter of the given method has the annotation @TestInstance. */
   boolean hasTestInstance(Method method, int i);
 
+  /** @return true if the ith parameter of the given method has the annotation
+   * {@link org.testng.annotations.ParameterCollector }. */
+  boolean hasParameterCollector(Method method, int i);
+
   /**
    * @return the @Optional values of this method's parameters (<code>null</code> if the parameter
    *     isn't optional)

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

@@ -32,6 +32,7 @@
 import org.testng.annotations.ObjectFactory;
 import org.testng.annotations.Optional;
 import org.testng.annotations.Parameters;
+import org.testng.annotations.ParameterCollector;
 import org.testng.annotations.Test;
 import org.testng.annotations.TestInstance;
 import org.testng.internal.ConstructorOrMethod;
@@ -260,6 +261,20 @@ public boolean hasTestInstance(Method method, int i) {
     return false;
   }
 
+  @Override
+  public boolean hasParameterCollector(Method method, int i) {
+    final Annotation[][] annotations = method.getParameterAnnotations();
+    if (annotations.length > 0 && annotations[i].length > 0) {
+      final Annotation[] pa = annotations[i];
+      for (Annotation a : pa) {
+        if (a instanceof ParameterCollector) {
+          return true;
+        }
+      }
+    }
+    return false;
+  }
+
   @Override
   public String[] findOptionalValues(Method method) {
     return optionalValues(method.getParameterAnnotations());

src/test/java/test/dataprovider/typed/TypedDataProviderSample.java

@@ -0,0 +1,30 @@
+package test.dataprovider.typed;
+
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.ParameterCollector;
+import org.testng.annotations.Test;
+
+public class TypedDataProviderSample {
+
+  private boolean dpDone = false;
+
+  @Test(dataProvider = "dp")
+  public void doStuff(String a, int b) {
+    if (!dpDone) throw new RuntimeException("Method should not be actually called by data provider");
+  }
+
+  @DataProvider(name = "dp")
+  public void createData(@ParameterCollector TypedDataProviderSample test) {
+    test.doStuff("test1", 1);
+    test.doStuff("test2", 2);
+    dpDone = true;
+  }
+
+  @Test(dataProvider = "dp2")
+  public void doStuff2(String a, int b) {}
+
+  @DataProvider(name = "dp2")
+  public void createData2(@ParameterCollector TypedDataProviderSample test) {
+    test.doStuff("test1", 1); // calling wrong method, should be doStuff2
+  }
+}

src/test/java/test/dataprovider/typed/TypedDataProviderTest.java

@@ -0,0 +1,25 @@
+package test.dataprovider.typed;
+
+import org.testng.annotations.Test;
+import test.InvokedMethodNameListener;
+import test.SimpleBaseTest;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class TypedDataProviderTest extends SimpleBaseTest {
+
+  @Test
+  public void typedDataProviderWorks() {
+    InvokedMethodNameListener listener = run(true, TypedDataProviderSample.class);
+    assertThat(listener.getSucceedMethodNames())
+        .containsExactly(
+            "doStuff(test1,1)",
+            "doStuff(test2,2)"
+        );
+    final Throwable throwable = listener.getResult("doStuff2").getThrowable();
+    assertThat(throwable).isNotNull();
+    assertThat(throwable.getMessage())
+            .contains("Wrong method doStuff is called, expected doStuff2");
+  }
+
+}

src/test/resources/testng.xml

@@ -576,6 +576,7 @@
       <class name="test.dataprovider.DataProviderTest"/>
       <class name="test.dataprovider.issue1987.IssueTest"/>
       <class name="test.dataprovider.InterceptorTest"/>
+      <class name="test.dataprovider.typed.TypedDataProviderTest" />
     </classes>
   </test>