From fda3ba7e3c4669442b3ed10284f8cabccb55e7e7 Mon Sep 17 00:00:00 2001
From: Anna Kozlova <anna.kozlova@jetbrains.com>
Date: Fri, 8 Jan 2016 15:55:41 +0100
Subject: [PATCH] test discovery: fork per module when tests were found in
 different modules

---
 .../JavaTestFrameworkRunnableState.java       |  5 +-
 .../TestDiscoveryConfiguration.java           | 12 +++-
 .../TestDiscoveryConfigurationProducer.java   |  6 ++
 .../execution/junit/JUnitDebuggerRunner.java  |  4 +-
 .../intellij/execution/junit/TestObject.java  |  9 ++-
 .../JUnitTestDiscoveryConfiguration.java      | 46 +++++++++++--
 .../configuration/TestNGDebuggerRunner.java   |  6 +-
 .../configuration/TestNGRunnableState.java    | 65 ++++++++++---------
 .../TestNGTestDiscoveryConfiguration.java     | 18 ++++-
 9 files changed, 126 insertions(+), 45 deletions(-)

diff --git a/java/execution/impl/src/com/intellij/execution/JavaTestFrameworkRunnableState.java b/java/execution/impl/src/com/intellij/execution/JavaTestFrameworkRunnableState.java
index 4bb082b1d29c..f5ce0f3e0442 100644
--- a/java/execution/impl/src/com/intellij/execution/JavaTestFrameworkRunnableState.java
+++ b/java/execution/impl/src/com/intellij/execution/JavaTestFrameworkRunnableState.java
@@ -322,8 +322,7 @@ public abstract class JavaTestFrameworkRunnableState<T extends ModuleBasedConfig
     }
   }
 
-  private boolean spansMultipleModules() {
-    final String qualifiedName = getConfiguration().getPackage();
+  protected boolean spansMultipleModules(final String qualifiedName) {
     if (qualifiedName != null) {
       final Project project = getConfiguration().getProject();
       final PsiPackage aPackage = JavaPsiFacade.getInstance(project).findPackage(qualifiedName);
@@ -350,7 +349,7 @@ public abstract class JavaTestFrameworkRunnableState<T extends ModuleBasedConfig
     final String workingDirectory = getConfiguration().getWorkingDirectory();
     return getScope() != TestSearchScope.SINGLE_MODULE &&
            ("$" + PathMacroUtil.MODULE_DIR_MACRO_NAME + "$").equals(workingDirectory) &&
-           spansMultipleModules();
+           spansMultipleModules(getConfiguration().getPackage());
   }
 
   protected void createTempFiles(JavaParameters javaParameters) {
diff --git a/java/execution/impl/src/com/intellij/execution/testDiscovery/TestDiscoveryConfiguration.java b/java/execution/impl/src/com/intellij/execution/testDiscovery/TestDiscoveryConfiguration.java
index 3548ffbf46ad..4b485484e5ba 100644
--- a/java/execution/impl/src/com/intellij/execution/testDiscovery/TestDiscoveryConfiguration.java
+++ b/java/execution/impl/src/com/intellij/execution/testDiscovery/TestDiscoveryConfiguration.java
@@ -16,8 +16,14 @@
 package com.intellij.execution.testDiscovery;
 
 import com.intellij.diagnostic.logging.LogConfigurationPanel;
-import com.intellij.execution.*;
-import com.intellij.execution.configurations.*;
+import com.intellij.execution.ExecutionBundle;
+import com.intellij.execution.Executor;
+import com.intellij.execution.JavaRunConfigurationExtensionManager;
+import com.intellij.execution.JavaTestConfigurationBase;
+import com.intellij.execution.configurations.ConfigurationFactory;
+import com.intellij.execution.configurations.JavaRunConfigurationModule;
+import com.intellij.execution.configurations.RunConfiguration;
+import com.intellij.execution.configurations.RuntimeConfigurationException;
 import com.intellij.execution.testframework.sm.runner.SMTRunnerConsoleProperties;
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.module.ModuleManager;
@@ -153,7 +159,7 @@ public abstract class TestDiscoveryConfiguration extends JavaTestConfigurationBa
   @Nullable
   @Override
   public String getPackage() {
-    return null;
+    return "";
   }
 
   @Override
diff --git a/java/execution/impl/src/com/intellij/execution/testDiscovery/TestDiscoveryConfigurationProducer.java b/java/execution/impl/src/com/intellij/execution/testDiscovery/TestDiscoveryConfigurationProducer.java
index 89f799d818a1..939259619165 100644
--- a/java/execution/impl/src/com/intellij/execution/testDiscovery/TestDiscoveryConfigurationProducer.java
+++ b/java/execution/impl/src/com/intellij/execution/testDiscovery/TestDiscoveryConfigurationProducer.java
@@ -22,6 +22,7 @@ import com.intellij.execution.actions.ConfigurationContext;
 import com.intellij.execution.configurations.ConfigurationType;
 import com.intellij.execution.junit.JavaRunConfigurationProducerBase;
 import com.intellij.execution.testframework.TestSearchScope;
+import com.intellij.openapi.module.Module;
 import com.intellij.openapi.util.Condition;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.Ref;
@@ -77,6 +78,11 @@ public abstract class TestDiscoveryConfigurationProducer extends JavaRunConfigur
     return false;
   }
 
+  @Override
+  protected Module findModule(TestDiscoveryConfiguration configuration, Module contextModule) {
+    return null;
+  }
+
   private static Pair<String, String> getPosition(Location location) {
     final PsiElement psiElement = location.getPsiElement();
     final PsiMethod psiMethod = PsiTreeUtil.getParentOfType(psiElement, PsiMethod.class);
diff --git a/plugins/junit/src/com/intellij/execution/junit/JUnitDebuggerRunner.java b/plugins/junit/src/com/intellij/execution/junit/JUnitDebuggerRunner.java
index 2c6e2f613eb6..d0a5b221583f 100644
--- a/plugins/junit/src/com/intellij/execution/junit/JUnitDebuggerRunner.java
+++ b/plugins/junit/src/com/intellij/execution/junit/JUnitDebuggerRunner.java
@@ -17,6 +17,7 @@ package com.intellij.execution.junit;
 
 import com.intellij.execution.JavaTestFrameworkDebuggerRunner;
 import com.intellij.execution.configurations.RunProfile;
+import com.intellij.execution.junit.testDiscovery.JUnitTestDiscoveryConfiguration;
 import org.jetbrains.annotations.NotNull;
 
 /**
@@ -25,7 +26,8 @@ import org.jetbrains.annotations.NotNull;
 public class JUnitDebuggerRunner extends JavaTestFrameworkDebuggerRunner {
   @Override
   protected boolean validForProfile(@NotNull RunProfile profile) {
-    return profile instanceof JUnitConfiguration;
+    return profile instanceof JUnitConfiguration || 
+           profile instanceof JUnitTestDiscoveryConfiguration;
   }
 
   @NotNull
diff --git a/plugins/junit/src/com/intellij/execution/junit/TestObject.java b/plugins/junit/src/com/intellij/execution/junit/TestObject.java
index 452dd932b9a8..a48926d97494 100644
--- a/plugins/junit/src/com/intellij/execution/junit/TestObject.java
+++ b/plugins/junit/src/com/intellij/execution/junit/TestObject.java
@@ -338,8 +338,9 @@ public abstract class TestObject extends JavaTestFrameworkRunnableState<JUnitCon
           continue;
         }
 
-        if (perModule != null && element instanceof PsiElement) {
-          final Module module = ModuleUtilCore.findModuleForPsiElement((PsiElement)element);
+        final PsiElement psiElement = retrievePsiElement(element);
+        if (perModule != null && psiElement != null) {
+          final Module module = ModuleUtilCore.findModuleForPsiElement(psiElement);
           if (module != null) {
             List<String> list = perModule.get(module);
             if (list == null) {
@@ -374,6 +375,10 @@ public abstract class TestObject extends JavaTestFrameworkRunnableState<JUnitCon
     }
   }
 
+  protected PsiElement retrievePsiElement(Object element) {
+    return element instanceof PsiElement ? (PsiElement)element : null;
+  }
+
   @Override
   protected void deleteTempFiles() {
     super.deleteTempFiles();
diff --git a/plugins/junit/src/com/intellij/execution/junit/testDiscovery/JUnitTestDiscoveryConfiguration.java b/plugins/junit/src/com/intellij/execution/junit/testDiscovery/JUnitTestDiscoveryConfiguration.java
index 4605b9959ea5..70460e4ca74c 100644
--- a/plugins/junit/src/com/intellij/execution/junit/testDiscovery/JUnitTestDiscoveryConfiguration.java
+++ b/plugins/junit/src/com/intellij/execution/junit/testDiscovery/JUnitTestDiscoveryConfiguration.java
@@ -25,8 +25,12 @@ import com.intellij.execution.runners.ExecutionEnvironment;
 import com.intellij.execution.testDiscovery.TestDiscoveryConfiguration;
 import com.intellij.execution.testDiscovery.TestDiscoverySearchHelper;
 import com.intellij.execution.testframework.SearchForTestsTask;
+import com.intellij.execution.testframework.SourceScope;
+import com.intellij.execution.testframework.TestSearchScope;
 import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.*;
+import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.refactoring.listeners.RefactoringElementListener;
 import com.intellij.util.FunctionUtil;
 import org.jetbrains.annotations.NotNull;
@@ -71,18 +75,52 @@ public class JUnitTestDiscoveryConfiguration extends TestDiscoveryConfiguration
     public JUnitTestDiscoveryRunnableState(ExecutionEnvironment environment) {
       super(((JUnitConfiguration)myDelegate), environment);
     }
-    
+
+    @Override
+    protected TestSearchScope getScope() {
+      return getConfigurationModule().getModule() != null ? TestSearchScope.MODULE_WITH_DEPENDENCIES : TestSearchScope.WHOLE_PROJECT;
+    }
+
+    @Override
+    protected boolean forkPerModule() {
+      return spansMultipleModules("");
+    }
+
+    @Override
+    protected PsiElement retrievePsiElement(Object pattern) {
+      if (pattern instanceof String) {
+        final String className = StringUtil.getPackageName((String)pattern, ',');
+        if (!pattern.equals(className)) {
+          final JavaPsiFacade facade = JavaPsiFacade.getInstance(getProject());
+          final SourceScope sourceScope = getSourceScope();
+          final GlobalSearchScope globalSearchScope = sourceScope != null ? sourceScope.getGlobalSearchScope()
+                                                                          : GlobalSearchScope.projectScope(getProject());
+          return facade.findClass(className, globalSearchScope);
+        }
+      }
+      return null;
+    }
+
     @Override
     public SearchForTestsTask createSearchingForTestsTask() {
       return new SearchForTestsTask(getProject(), myServerSocket) {
+
+        private Set<String> myPatterns;
+
         @Override
         protected void search() throws ExecutionException {
-          final Set<String> patterns = TestDiscoverySearchHelper.search(getProject(), getPosition(), getChangeList(), getFrameworkPrefix());
-          addClassesListToJavaParameters(patterns, FunctionUtil.<String>id(), "", false, getJavaParameters());
+          myPatterns = TestDiscoverySearchHelper.search(getProject(), getPosition(), getChangeList(), getFrameworkPrefix());
         }
 
         @Override
-        protected void onFound() {}
+        protected void onFound() {
+          if (myPatterns != null) {
+            try {
+              addClassesListToJavaParameters(myPatterns, FunctionUtil.<String>id(), "", false, getJavaParameters());
+            }
+            catch (ExecutionException ignored) {}
+          }
+        }
       };
     }
 
diff --git a/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGDebuggerRunner.java b/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGDebuggerRunner.java
index 175f12a2ba3a..bdd14ce60176 100644
--- a/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGDebuggerRunner.java
+++ b/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGDebuggerRunner.java
@@ -15,15 +15,17 @@
  */
 package com.theoryinpractice.testng.configuration;
 
-import com.intellij.execution.configurations.RunProfile;
 import com.intellij.execution.JavaTestFrameworkDebuggerRunner;
+import com.intellij.execution.configurations.RunProfile;
+import com.theoryinpractice.testng.configuration.testDiscovery.TestNGTestDiscoveryConfiguration;
 import org.jetbrains.annotations.NotNull;
 
 
 public class TestNGDebuggerRunner extends JavaTestFrameworkDebuggerRunner {
   @Override
   protected boolean validForProfile(@NotNull RunProfile profile) {
-    return profile instanceof TestNGConfiguration;
+    return profile instanceof TestNGConfiguration || 
+           profile instanceof TestNGTestDiscoveryConfiguration;
   }
 
   @NotNull
diff --git a/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGRunnableState.java b/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGRunnableState.java
index 0bc127bf2344..7562061f400a 100644
--- a/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGRunnableState.java
+++ b/plugins/testng/src/com/theoryinpractice/testng/configuration/TestNGRunnableState.java
@@ -23,7 +23,10 @@ import com.intellij.execution.configurations.ParametersList;
 import com.intellij.execution.process.*;
 import com.intellij.execution.runners.ExecutionEnvironment;
 import com.intellij.execution.runners.ProgramRunner;
-import com.intellij.execution.testframework.*;
+import com.intellij.execution.testframework.Printable;
+import com.intellij.execution.testframework.Printer;
+import com.intellij.execution.testframework.TestFrameworkRunningModel;
+import com.intellij.execution.testframework.TestSearchScope;
 import com.intellij.execution.ui.ConsoleViewContentType;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.module.Module;
@@ -36,6 +39,7 @@ import com.intellij.openapi.util.registry.Registry;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.JavaPsiFacade;
 import com.intellij.psi.PsiClass;
+import com.intellij.psi.PsiMethod;
 import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.rt.execution.testFrameworks.ForkedDebuggerHelper;
 import com.intellij.util.PathUtil;
@@ -239,38 +243,41 @@ public class TestNGRunnableState extends JavaTestFrameworkRunnableState<TestNGCo
       @Override
       protected void onFound() {
         super.onFound();
-
-        if (forkPerModule()) {
-          final Map<Module, List<String>> perModule = new TreeMap<Module, List<String>>(new Comparator<Module>() {
-            @Override
-            public int compare(Module o1, Module o2) {
-              return StringUtil.compare(o1.getName(), o2.getName(), true);
-            }
-          });
-
-          for (final PsiClass psiClass : myClasses.keySet()) {
-            final Module module = ModuleUtilCore.findModuleForPsiElement(psiClass);
-            if (module != null) {
-              List<String> list = perModule.get(module);
-              if (list == null) {
-                list = new ArrayList<String>();
-                perModule.put(module, list);
-              }
-              list.add(psiClass.getQualifiedName());
-            }
-          }
-
-          try {
-            writeClassesPerModule(getConfiguration().getPackage(), getJavaParameters(), perModule);
-          }
-          catch (Exception e) {
-            LOG.error(e);
-          }
-        }
+        writeClassesPerModule(myClasses);
       }
     };
   }
 
+  protected void writeClassesPerModule(Map<PsiClass, Map<PsiMethod, List<String>>> classes) {
+    if (forkPerModule()) {
+      final Map<Module, List<String>> perModule = new TreeMap<Module, List<String>>(new Comparator<Module>() {
+        @Override
+        public int compare(Module o1, Module o2) {
+          return StringUtil.compare(o1.getName(), o2.getName(), true);
+        }
+      });
+
+      for (final PsiClass psiClass : classes.keySet()) {
+        final Module module = ModuleUtilCore.findModuleForPsiElement(psiClass);
+        if (module != null) {
+          List<String> list = perModule.get(module);
+          if (list == null) {
+            list = new ArrayList<String>();
+            perModule.put(module, list);
+          }
+          list.add(psiClass.getQualifiedName());
+        }
+      }
+
+      try {
+        writeClassesPerModule(getConfiguration().getPackage(), getJavaParameters(), perModule);
+      }
+      catch (Exception e) {
+        LOG.error(e);
+      }
+    }
+  }
+
   public static boolean supportSerializationProtocol(TestNGConfiguration config) {
     final Project project = config.getProject();
     final GlobalSearchScope scopeToDetermineTestngIn;
diff --git a/plugins/testng/src/com/theoryinpractice/testng/configuration/testDiscovery/TestNGTestDiscoveryConfiguration.java b/plugins/testng/src/com/theoryinpractice/testng/configuration/testDiscovery/TestNGTestDiscoveryConfiguration.java
index 5406f280093f..09efc7ef06aa 100644
--- a/plugins/testng/src/com/theoryinpractice/testng/configuration/testDiscovery/TestNGTestDiscoveryConfiguration.java
+++ b/plugins/testng/src/com/theoryinpractice/testng/configuration/testDiscovery/TestNGTestDiscoveryConfiguration.java
@@ -75,7 +75,17 @@ public class TestNGTestDiscoveryConfiguration extends TestDiscoveryConfiguration
     public TestNGTestDiscoveryRunnableState(ExecutionEnvironment environment) {
       super(environment, ((TestNGConfiguration)myDelegate));
     }
-    
+
+    @Override
+    protected TestSearchScope getScope() {
+      return TestSearchScope.MODULE_WITH_DEPENDENCIES;
+    }
+
+    @Override
+    protected boolean forkPerModule() {
+      return spansMultipleModules("");
+    }
+
     @Override
     public SearchingForTestsTask createSearchingForTestsTask() {
       return new SearchingForTestsTask(myServerSocket, getConfiguration(), myTempFile, client) {
@@ -89,6 +99,12 @@ public class TestNGTestDiscoveryConfiguration extends TestDiscoveryConfiguration
           TestNGTestPattern.fillTestObjects(myClasses, patterns, TestSearchScope.MODULE_WITH_DEPENDENCIES, 
                                             TestNGTestDiscoveryConfiguration.this, searchScope);
         }
+
+        @Override
+        protected void onFound() {
+          super.onFound();
+          writeClassesPerModule(myClasses);
+        }
       };
     }
   }