Check return type of factory methods

3 files changed, 32 insertions, 3 deletions

diff --git a/src/main/java/org/testng/internal/ClassHelper.java b/src/main/java/org/testng/internal/ClassHelper.java
index 5b854a04..a348da34 100644
--- a/src/main/java/org/testng/internal/ClassHelper.java
+++ b/src/main/java/org/testng/internal/ClassHelper.java
@@ -166,9 +166,6 @@ public final class ClassHelper {
 //      }
 //    }
 
-    // Found the method, verify that it returns an array of objects
-    // TBD
-
     return result;
   }
 
diff --git a/src/main/java/org/testng/internal/FactoryMethod.java b/src/main/java/org/testng/internal/FactoryMethod.java
index fb802871..4ad59c05 100644
--- a/src/main/java/org/testng/internal/FactoryMethod.java
+++ b/src/main/java/org/testng/internal/FactoryMethod.java
@@ -5,6 +5,7 @@ import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 
+import org.testng.IInstanceInfo;
 import org.testng.ITestContext;
 import org.testng.ITestNGMethod;
 import org.testng.TestNGException;
@@ -35,6 +36,7 @@ public class FactoryMethod extends BaseTestMethod {
   {
     super(com.getName(), com, annotationFinder, instance);
     Utils.checkInstanceOrStatic(instance, com.getMethod());
+    Utils.checkReturnType(com.getMethod(), Object[].class, IInstanceInfo[].class);
     Class<?> declaringClass = com.getDeclaringClass();
     if (instance != null && ! declaringClass.isAssignableFrom(instance.getClass())) {
       throw new TestNGException("Mismatch between instance/method classes:"
diff --git a/src/main/java/org/testng/internal/Utils.java b/src/main/java/org/testng/internal/Utils.java
index aa399b0c..b98163c8 100644
--- a/src/main/java/org/testng/internal/Utils.java
+++ b/src/main/java/org/testng/internal/Utils.java
@@ -802,6 +802,36 @@ public final class Utils {
     }
   }
 
+  public static void checkReturnType(Method method, Class<?>... returnTypes) {
+    if (method == null) {
+      return;
+    }
+    for (Class<?> returnType : returnTypes) {
+      if (method.getReturnType() == returnType) {
+        return;
+      }
+    }
+    throw new TestNGException(method.getDeclaringClass().getName() + "."
+              + method.getName() + " MUST return " + toString(returnTypes) + " but returns " + method.getReturnType().getName());
+  }
+
+  private static String toString(Class<?>[] classes) {
+    StringBuilder sb = new StringBuilder("[ ");
+    for (int i=0; i<classes.length;) {
+      Class<?> clazz = classes[i];
+      if (clazz.isArray()) {
+        sb.append(clazz.getComponentType().getName()).append("[]");
+      } else {
+        sb.append(clazz.getName());
+      }
+      if (++i < classes.length) { // increment and compare
+        sb.append(" or ");
+      }
+    }
+    sb.append(" ]");
+    return sb.toString();
+  }
+
   /**
    * Returns the string representation of the specified object, transparently
    * handling null references and arrays.