From d7f75583f97a15a77e6d784116e103f4142b3034 Mon Sep 17 00:00:00 2001
From: Anna Kozlova <anna.kozlova@jetbrains.com>
Date: Tue, 1 Sep 2015 13:26:01 +0300
Subject: [PATCH] junit: make it possible to start pattern configuration with
 parameter from @Parameterized (IDEA-144608)

---
 ...bstractPatternBasedConfigurationProducer.java |  5 ++++-
 .../junit/PatternConfigurationProducer.java      | 10 ++++++++++
 .../intellij/junit4/JUnit4TestRunnerUtil.java    | 12 ++++++++++--
 .../TestNGConfigurationProducer.java             |  9 +++++++++
 .../TestNGInClassConfigurationProducer.java      |  8 +-------
 .../TestNGPatternConfigurationProducer.java      |  3 +++
 .../ui/actions/RerunFailedTestsAction.java       | 16 ++++++++++------
 7 files changed, 47 insertions(+), 16 deletions(-)

diff --git a/java/execution/impl/src/com/intellij/execution/testframework/AbstractPatternBasedConfigurationProducer.java b/java/execution/impl/src/com/intellij/execution/testframework/AbstractPatternBasedConfigurationProducer.java
index b1186ab0efc3..db9b42a863eb 100644
--- a/java/execution/impl/src/com/intellij/execution/testframework/AbstractPatternBasedConfigurationProducer.java
+++ b/java/execution/impl/src/com/intellij/execution/testframework/AbstractPatternBasedConfigurationProducer.java
@@ -21,6 +21,7 @@ import com.intellij.execution.actions.ConfigurationContext;
 import com.intellij.execution.configurations.ConfigurationType;
 import com.intellij.execution.configurations.ModuleBasedConfiguration;
 import com.intellij.execution.junit.JavaRunConfigurationProducerBase;
+import com.intellij.execution.junit2.PsiMemberParameterizedLocation;
 import com.intellij.execution.junit2.info.MethodLocation;
 import com.intellij.openapi.actionSystem.CommonDataKeys;
 import com.intellij.openapi.actionSystem.DataContext;
@@ -193,7 +194,9 @@ public abstract class AbstractPatternBasedConfigurationProducer<T extends Module
     }
     else if (psiMember instanceof PsiMember) {
       final PsiClass containingClass = location instanceof MethodLocation
-                                       ? ((MethodLocation)location).getContainingClass(): ((PsiMember)psiMember).getContainingClass();
+                                       ? ((MethodLocation)location).getContainingClass()
+                                       : location instanceof PsiMemberParameterizedLocation ? ((PsiMemberParameterizedLocation)location).getContainingClass() 
+                                                                                            : ((PsiMember)psiMember).getContainingClass();
       assert containingClass != null;
       return ClassUtil.getJVMClassName(containingClass) + "," + ((PsiMember)psiMember).getName();
     } else if (psiMember instanceof PsiPackage) {
diff --git a/plugins/junit/src/com/intellij/execution/junit/PatternConfigurationProducer.java b/plugins/junit/src/com/intellij/execution/junit/PatternConfigurationProducer.java
index bf528c10d0bb..b9fd05a5787b 100644
--- a/plugins/junit/src/com/intellij/execution/junit/PatternConfigurationProducer.java
+++ b/plugins/junit/src/com/intellij/execution/junit/PatternConfigurationProducer.java
@@ -16,6 +16,9 @@
 
 package com.intellij.execution.junit;
 
+import com.intellij.execution.Location;
+import com.intellij.execution.junit2.PsiMemberParameterizedLocation;
+import com.intellij.execution.junit2.info.MethodLocation;
 import com.intellij.execution.testframework.AbstractPatternBasedConfigurationProducer;
 import com.intellij.execution.actions.ConfigurationContext;
 import com.intellij.execution.actions.ConfigurationFromContext;
@@ -61,6 +64,13 @@ public class PatternConfigurationProducer extends AbstractPatternBasedConfigurat
     data.TEST_OBJECT = JUnitConfiguration.TEST_PATTERN;
     data.setScope(setupPackageConfiguration(context, configuration, data.getScope()));
     configuration.setGeneratedName();
+    final Location contextLocation = context.getLocation();
+    if (contextLocation instanceof PsiMemberParameterizedLocation) {
+      final String paramSetName = ((PsiMemberParameterizedLocation)contextLocation).getParamSetName();
+      if (paramSetName != null) {
+        configuration.setProgramParameters(paramSetName);
+      }
+    }
     return true;
   }
 
diff --git a/plugins/junit_rt/src/com/intellij/junit4/JUnit4TestRunnerUtil.java b/plugins/junit_rt/src/com/intellij/junit4/JUnit4TestRunnerUtil.java
index 2fde52e4f095..a2b0387c8440 100644
--- a/plugins/junit_rt/src/com/intellij/junit4/JUnit4TestRunnerUtil.java
+++ b/plugins/junit_rt/src/com/intellij/junit4/JUnit4TestRunnerUtil.java
@@ -44,7 +44,7 @@ public class JUnit4TestRunnerUtil {
    */
   private static final ResourceBundle ourBundle = ResourceBundle.getBundle("RuntimeBundle");
 
-  public static Request buildRequest(String[] suiteClassNames, String name, boolean notForked) {
+  public static Request buildRequest(String[] suiteClassNames, final String name, boolean notForked) {
     if (suiteClassNames.length == 0) {
       return null;
     }
@@ -101,7 +101,15 @@ public class JUnit4TestRunnerUtil {
               public boolean shouldRun(Description description) {
                 if (description.isTest()) {
                   final Set methods = (Set)classMethods.get(JUnit4ReflectionUtil.getClassName(description));
-                  return methods == null || methods.contains(JUnit4ReflectionUtil.getMethodName(description));
+                  if (methods == null) {
+                    return true;
+                  }
+                  String methodName = JUnit4ReflectionUtil.getMethodName(description);
+                  if (methods.contains(methodName)) {
+                    return true;
+                  }
+                  return name != null && methodName.endsWith(name) && 
+                         methods.contains(methodName.substring(0, methodName.length() - name.length()));
                 }
                 return true;
               }
diff --git a/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGConfigurationProducer.java b/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGConfigurationProducer.java
index 0914e2b476dc..dfb4074e15c7 100644
--- a/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGConfigurationProducer.java
+++ b/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGConfigurationProducer.java
@@ -71,6 +71,15 @@ public abstract class TestNGConfigurationProducer extends JavaRunConfigurationPr
     }
     return false;
   }
+
+  public static void setupConfigurationParamName(TestNGConfiguration configuration, Location contextLocation) {
+    if (contextLocation instanceof PsiMemberParameterizedLocation) {
+      final String paramSetName = ((PsiMemberParameterizedLocation)contextLocation).getParamSetName();
+      if (paramSetName != null) {
+        configuration.setProgramParameters(getInvocationNumber(paramSetName));
+      }
+    }
+  }
   
   public static String getInvocationNumber(String str) {
     return StringUtil.trimEnd(StringUtil.trimStart(str, "["), "]");
diff --git a/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGInClassConfigurationProducer.java b/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGInClassConfigurationProducer.java
index d30386d78f0c..186e94c5cd0d 100644
--- a/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGInClassConfigurationProducer.java
+++ b/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGInClassConfigurationProducer.java
@@ -27,7 +27,6 @@ import com.intellij.execution.actions.ConfigurationContext;
 import com.intellij.execution.actions.ConfigurationFromContext;
 import com.intellij.execution.actions.RunConfigurationProducer;
 import com.intellij.execution.junit.InheritorChooser;
-import com.intellij.execution.junit2.PsiMemberParameterizedLocation;
 import com.intellij.execution.junit2.info.MethodLocation;
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.project.Project;
@@ -103,12 +102,7 @@ public class TestNGInClassConfigurationProducer extends TestNGConfigurationProdu
     }
 
     final Location contextLocation = context.getLocation();
-    if (contextLocation instanceof PsiMemberParameterizedLocation) {
-      final String paramSetName = getInvocationNumber(((PsiMemberParameterizedLocation)contextLocation).getParamSetName());
-      if (paramSetName != null) {
-        configuration.setProgramParameters(paramSetName);
-      }
-    }
+    setupConfigurationParamName(configuration, contextLocation);
 
     PsiClass psiClass = null;
     PsiElement element = context.getPsiLocation();
diff --git a/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGPatternConfigurationProducer.java b/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGPatternConfigurationProducer.java
index 1f2e0b023b05..b40168da310d 100644
--- a/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGPatternConfigurationProducer.java
+++ b/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGPatternConfigurationProducer.java
@@ -20,6 +20,8 @@
  */
 package com.theoryinpractice.testng.configuration;
 
+import com.intellij.execution.Location;
+import com.intellij.execution.junit2.PsiMemberParameterizedLocation;
 import com.intellij.execution.testframework.AbstractPatternBasedConfigurationProducer;
 import com.intellij.execution.actions.ConfigurationContext;
 import com.intellij.openapi.module.Module;
@@ -65,6 +67,7 @@ public class TestNGPatternConfigurationProducer extends AbstractPatternBasedConf
     data.TEST_OBJECT = TestType.PATTERN.getType();
     data.setScope(setupPackageConfiguration(context, configuration, data.getScope()));
     configuration.setGeneratedName();
+    TestNGConfigurationProducer.setupConfigurationParamName(configuration, context.getLocation());
     return true;
   }
 
diff --git a/plugins/testng/src/com/theoryinpractice/testng/ui/actions/RerunFailedTestsAction.java b/plugins/testng/src/com/theoryinpractice/testng/ui/actions/RerunFailedTestsAction.java
index abb194ce6180..4cbee48eddd9 100644
--- a/plugins/testng/src/com/theoryinpractice/testng/ui/actions/RerunFailedTestsAction.java
+++ b/plugins/testng/src/com/theoryinpractice/testng/ui/actions/RerunFailedTestsAction.java
@@ -113,14 +113,18 @@ public class RerunFailedTestsAction extends JavaRerunFailedTestsAction {
         if (strings == null || strings.isEmpty()) {
           strings = new ArrayList<String>();
         }
-        if (location instanceof PsiMemberParameterizedLocation) {
-          final String paramSetName = ((PsiMemberParameterizedLocation)location).getParamSetName();
-          if (paramSetName != null) {
-            strings.add(TestNGConfigurationProducer.getInvocationNumber(paramSetName));
-          }
-        }
+        setupParameterName(location, strings);
         psiMethods.put(psiMethod, strings);
       }
     }
   }
+
+  private static void setupParameterName(Location location, List<String> strings) {
+    if (location instanceof PsiMemberParameterizedLocation) {
+      final String paramSetName = ((PsiMemberParameterizedLocation)location).getParamSetName();
+      if (paramSetName != null) {
+        strings.add(TestNGConfigurationProducer.getInvocationNumber(paramSetName));
+      }
+    }
+  }
 }