IDEA-205919 fixed wrong test runner for package tests

plugins/gradle/java/src/execution/test/runner/AllInPackageGradleConfigurationProducer.java

@@ -3,24 +3,33 @@ package org.jetbrains.plugins.gradle.execution.test.runner;
 
 import com.intellij.execution.ExecutionBundle;
 import com.intellij.execution.actions.ConfigurationContext;
+import com.intellij.execution.actions.ConfigurationFromContext;
 import com.intellij.execution.configurations.ConfigurationFactory;
 import com.intellij.execution.junit.JavaRuntimeConfigurationProducerBase;
+import com.intellij.openapi.externalSystem.model.execution.ExternalSystemTaskExecutionSettings;
 import com.intellij.openapi.externalSystem.service.execution.ExternalSystemRunConfiguration;
 import com.intellij.openapi.externalSystem.util.ExternalSystemApiUtil;
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.util.Ref;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.VirtualFile;
-import com.intellij.psi.PsiDirectory;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiPackage;
+import com.intellij.psi.*;
 import com.intellij.psi.search.GlobalSearchScope;
+import com.intellij.util.ArrayUtil;
+import kotlin.jvm.functions.Function1;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.plugins.gradle.service.execution.GradleExternalTaskConfigurationType;
 import org.jetbrains.plugins.gradle.util.GradleConstants;
-import org.jetbrains.plugins.gradle.util.GradleExecutionSettingsUtil;
 import org.jetbrains.plugins.gradle.util.TasksToRun;
 
+import java.util.List;
+import java.util.Map;
+
+import static org.jetbrains.plugins.gradle.execution.test.runner.TestGradleConfigurationProducerUtilKt.applyTestConfiguration;
+import static org.jetbrains.plugins.gradle.execution.test.runner.TestGradleConfigurationProducerUtilKt.getSourceFile;
+import static org.jetbrains.plugins.gradle.util.GradleExecutionSettingsUtil.createTestFilterFrom;
+
 /**
  * @author Vladislav.Soroka
  */
@@ -35,58 +44,123 @@ public final class AllInPackageGradleConfigurationProducer extends GradleTestRun
   protected boolean doSetupConfigurationFromContext(ExternalSystemRunConfiguration configuration,
                                                     ConfigurationContext context,
                                                     Ref<PsiElement> sourceElement) {
-    final PsiPackage psiPackage = JavaRuntimeConfigurationProducerBase.checkPackage(context.getPsiLocation());
-    if (psiPackage == null) return false;
-    sourceElement.set(psiPackage);
+    ConfigurationData configurationData = extractConfigurationElements(context);
+    if (configurationData == null) return false;
 
-    final Module module = context.getModule();
-    if (module == null) return false;
-
-    if (!ExternalSystemApiUtil.isExternalSystemAwareModule(GradleConstants.SYSTEM_ID, module)) return false;
-
-    final String projectPath = ExternalSystemApiUtil.getExternalProjectPath(module);
-    if (projectPath == null) return false;
-
-    PsiDirectory[] sourceDirs = psiPackage.getDirectories(GlobalSearchScope.moduleScope(module));
-    if (sourceDirs.length == 0) return false;
-    VirtualFile source = sourceDirs[0].getVirtualFile();
-    TasksToRun tasksToRun = findTestsTaskToRun(source, context.getProject());
+    TasksToRun tasksToRun = findTestsTaskToRun(configurationData.source, context.getProject());
     if (tasksToRun.isEmpty()) return false;
 
-    configuration.getSettings().setExternalProjectPath(projectPath);
+    sourceElement.set(configurationData.sourceElement);
+
+    configuration.getSettings().setExternalProjectPath(configurationData.projectPath);
     configuration.getSettings().setTaskNames(tasksToRun);
-    String filter = GradleExecutionSettingsUtil.createTestFilterFrom(psiPackage, /*hasSuffix=*/false);
+    String filter = createTestFilterFrom(configurationData.psiPackage, /*hasSuffix=*/false);
     configuration.getSettings().setScriptParameters(filter);
-    configuration.setName(suggestName(psiPackage, module));
+    configuration.setName(suggestName(configurationData.psiPackage, configurationData.module));
     return true;
   }
 
   @Override
   protected boolean doIsConfigurationFromContext(ExternalSystemRunConfiguration configuration, ConfigurationContext context) {
-    final PsiPackage psiPackage = JavaRuntimeConfigurationProducerBase.checkPackage(context.getPsiLocation());
-    if (psiPackage == null) return false;
-
-    if (context.getModule() == null) return false;
+    ConfigurationData configurationData = extractConfigurationElements(context);
+    if (configurationData == null) return false;
 
     if (!StringUtil.equals(
-      ExternalSystemApiUtil.getExternalProjectPath(context.getModule()),
+      configurationData.projectPath,
       configuration.getSettings().getExternalProjectPath())) {
       return false;
     }
-    Module module = context.getModule();
-    PsiDirectory[] sourceDirs = psiPackage.getDirectories(GlobalSearchScope.moduleScope(module));
-    if (sourceDirs.length == 0) return false;
-    VirtualFile source = sourceDirs[0].getVirtualFile();
-    if (!hasTasksInConfiguration(source, context.getProject(), configuration.getSettings())) return false;
+    if (!hasTasksInConfiguration(configurationData.source, context.getProject(), configuration.getSettings())) return false;
 
     final String scriptParameters = configuration.getSettings().getScriptParameters() + ' ';
-    final String filter = GradleExecutionSettingsUtil.createTestFilterFrom(psiPackage, /*hasSuffix=*/true);
+    final String filter = createTestFilterFrom(configurationData.psiPackage, /*hasSuffix=*/true);
     return scriptParameters.contains(filter);
   }
 
+  @Override
+  public void onFirstRun(@NotNull ConfigurationFromContext fromContext,
+                         @NotNull ConfigurationContext context,
+                         @NotNull Runnable performRunnable) {
+    ConfigurationData configurationData = extractConfigurationElements(context);
+    if (configurationData == null) return;
+    TasksChooser tasksChooser = new TasksChooser() {
+      @Override
+      protected void choosesTasks(@NotNull List<? extends Map<String, ? extends List<String>>> tasks) {
+        ExternalSystemRunConfiguration configuration = (ExternalSystemRunConfiguration)fromContext.getConfiguration();
+        ExternalSystemTaskExecutionSettings settings = configuration.getSettings();
+        Function1<PsiElement, String> createFilter = (e) -> createTestFilterFrom(configurationData.psiPackage, /*hasSuffix=*/false);
+        PsiElement[] sourceElements = ArrayUtil.toObjectArray(PsiElement.class, configurationData.sourceElement);
+        if (!applyTestConfiguration(settings, context.getProject(), tasks, sourceElements, createFilter)) return;
+        configuration.setName(suggestName(configurationData.psiPackage, configurationData.module));
+        performRunnable.run();
+      }
+    };
+    tasksChooser.runTaskChoosing(context, configurationData.sourceElement);
+  }
+
+  @Nullable
+  private ConfigurationData extractConfigurationElements(ConfigurationContext context) {
+    PsiElement contextLocation = context.getPsiLocation();
+    if (contextLocation == null) return null;
+    PsiPackage psiPackage = JavaRuntimeConfigurationProducerBase.checkPackage(contextLocation);
+    if (psiPackage == null) return null;
+    Module module = context.getModule();
+    if (module == null) return null;
+    if (!ExternalSystemApiUtil.isExternalSystemAwareModule(GradleConstants.SYSTEM_ID, module)) return null;
+    String projectPath = resolveProjectPath(module);
+    if (projectPath == null) return null;
+    PsiElement sourceElement = getSourceElement(module, contextLocation);
+    if (sourceElement == null) return null;
+    VirtualFile source = getSourceFile(sourceElement);
+    if (source == null) return null;
+    return new ConfigurationData(module, psiPackage, sourceElement, source, projectPath);
+  }
+
+  @Nullable
+  private static PsiElement getSourceElement(@NotNull Module module, @NotNull PsiElement element) {
+    if (element instanceof PsiFileSystemItem) {
+      return element;
+    }
+    PsiFile containingFile = element.getContainingFile();
+    if (containingFile != null) {
+      return element;
+    }
+    if (element instanceof PsiPackage) {
+      return getPackageDirectory(module, (PsiPackage)element);
+    }
+    return null;
+  }
+
+  @Nullable
+  private static PsiDirectory getPackageDirectory(@NotNull Module module, @NotNull PsiPackage element) {
+    PsiDirectory[] sourceDirs = element.getDirectories(GlobalSearchScope.moduleScope(module));
+    if (sourceDirs.length == 0) return null;
+    return sourceDirs[0];
+  }
+
   private static String suggestName(@NotNull PsiPackage aPackage, @NotNull Module module) {
     return aPackage.getQualifiedName().isEmpty()
            ? ExecutionBundle.message("test.in.scope.presentable.text", module.getName())
            : ExecutionBundle.message("test.in.scope.presentable.text", aPackage.getQualifiedName());
   }
+
+  private static class ConfigurationData {
+    public final @NotNull Module module;
+    public final @NotNull PsiPackage psiPackage;
+    public final @NotNull PsiElement sourceElement;
+    public final @NotNull VirtualFile source;
+    public final @NotNull String projectPath;
+
+    private ConfigurationData(@NotNull Module module,
+                              @NotNull PsiPackage psiPackage,
+                              @NotNull PsiElement sourceElement,
+                              @NotNull VirtualFile source,
+                              @NotNull String projectPath) {
+      this.module = module;
+      this.psiPackage = psiPackage;
+      this.sourceElement = sourceElement;
+      this.source = source;
+      this.projectPath = projectPath;
+    }
+  }
 }

plugins/gradle/java/src/execution/test/runner/GradleTestRunConfigurationProducer.java

@@ -15,6 +15,7 @@ import com.intellij.openapi.externalSystem.model.task.TaskData;
 import com.intellij.openapi.externalSystem.service.execution.ExternalSystemRunConfiguration;
 import com.intellij.openapi.externalSystem.util.ExternalSystemApiUtil;
 import com.intellij.openapi.module.Module;
+import com.intellij.openapi.module.ModuleUtilCore;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.roots.ProjectFileIndex;
 import com.intellij.openapi.util.Ref;
@@ -22,7 +23,6 @@ import com.intellij.openapi.util.io.FileUtil;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiFile;
 import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -224,14 +224,8 @@ public abstract class GradleTestRunConfigurationProducer extends RunConfiguratio
   }
 
   private static TestRunner getTestRunner(@NotNull PsiElement sourceElement) {
-    PsiFile containingFile = sourceElement.getContainingFile();
-    if (containingFile != null) {
-      VirtualFile file = containingFile.getVirtualFile();
-      Module module = file == null ? null : ProjectFileIndex.SERVICE.getInstance(sourceElement.getProject()).getModuleForFile(file);
-      if (module != null) {
-        return GradleSettingsService.getTestRunner(module);
-      }
-    }
-    return PLATFORM;
+    Module module = ModuleUtilCore.findModuleForPsiElement(sourceElement);
+    if (module == null) return PLATFORM;
+    return GradleSettingsService.getTestRunner(module);
   }
 }

plugins/gradle/java/src/execution/test/runner/TasksChooser.kt

@@ -22,7 +22,7 @@ abstract class TasksChooser {
   protected abstract fun choosesTasks(tasks: List<Map<String, List<String>>>)
 
   fun runTaskChoosing(context: ConfigurationContext, elements: Iterable<PsiElement>) {
-    val sources = elements.map { it.containingFile?.virtualFile ?: error("Can not find source file for $it") }
+    val sources = elements.map { getSourceFile(it) ?: error("Can not find source file for $it") }
     runTaskChoosing(context.dataContext, sources, context.project)
   }
 

plugins/gradle/java/src/execution/test/runner/TestGradleConfigurationProducerUtil.kt

@@ -6,52 +6,53 @@ import com.intellij.openapi.module.Module
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ProjectFileIndex
 import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFileSystemItem
 import org.jetbrains.plugins.gradle.execution.GradleRunnerUtil
 import org.jetbrains.plugins.gradle.execution.test.runner.GradleTestRunConfigurationProducer.findTestsTaskToRun
 import java.util.*
 
-fun ExternalSystemTaskExecutionSettings.applyTestConfiguration(
+fun <E : PsiElement> ExternalSystemTaskExecutionSettings.applyTestConfiguration(
   project: Project,
   testTasksToRun: List<Map<String, List<String>>>,
-  containingClasses: Iterable<PsiClass>,
-  createFilter: (PsiClass) -> String
+  sourceElements: Iterable<E>,
+  createFilter: (E) -> String
 ): Boolean {
-  return applyTestConfiguration(project, containingClasses, { it }, { it, _ -> createFilter(it) }) { source ->
+  return applyTestConfiguration(project, sourceElements, { it }, { it, _ -> createFilter(it) }) { source ->
     testTasksToRun.mapNotNull { it[source.path] }
   }
 }
 
-fun ExternalSystemTaskExecutionSettings.applyTestConfiguration(
+fun <E : PsiElement> ExternalSystemTaskExecutionSettings.applyTestConfiguration(
   project: Project,
-  containingClasses: Iterable<PsiClass>,
-  createFilter: (PsiClass) -> String
+  sourceElements: Iterable<E>,
+  createFilter: (E) -> String
 ): Boolean {
-  return applyTestConfiguration(project, containingClasses, { it }, { it, _ -> createFilter(it) })
+  return applyTestConfiguration(project, sourceElements, { it }, { it, _ -> createFilter(it) })
 }
 
-fun ExternalSystemTaskExecutionSettings.applyTestConfiguration(
+fun <E : PsiElement> ExternalSystemTaskExecutionSettings.applyTestConfiguration(
   project: Project,
   testTasksToRun: List<Map<String, List<String>>>,
-  vararg containingClasses: PsiClass,
-  createFilter: (PsiClass) -> String
+  vararg sourceElements: E,
+  createFilter: (E) -> String
 ): Boolean {
-  return applyTestConfiguration(project, containingClasses.toList(), { it }, { it, _ -> createFilter(it) }) { source ->
+  return applyTestConfiguration(project, sourceElements.toList(), { it }, { it, _ -> createFilter(it) }) { source ->
     testTasksToRun.mapNotNull { it[source.path] }
   }
 }
 
-fun ExternalSystemTaskExecutionSettings.applyTestConfiguration(
+fun <E : PsiElement> ExternalSystemTaskExecutionSettings.applyTestConfiguration(
   project: Project,
-  vararg containingClasses: PsiClass,
-  createFilter: (PsiClass) -> String
+  vararg sourceElements: E,
+  createFilter: (E) -> String
 ): Boolean {
-  return applyTestConfiguration(project, containingClasses.toList(), { it }) { it, _ ->
+  return applyTestConfiguration(project, sourceElements.toList(), { it }) { it, _ ->
     createFilter(it)
   }
 }
 
-fun <T> ExternalSystemTaskExecutionSettings.applyTestConfiguration(
+fun <E : PsiElement, T> ExternalSystemTaskExecutionSettings.applyTestConfiguration(
   project: Project,
   testTasksToRun: List<Map<String, List<String>>>,
   tests: Iterable<T>,
@@ -65,31 +66,30 @@ fun <T> ExternalSystemTaskExecutionSettings.applyTestConfiguration(
 fun <E : PsiElement, T> ExternalSystemTaskExecutionSettings.applyTestConfiguration(
   project: Project,
   tests: Iterable<T>,
-  findPsiClass: (T) -> PsiClass?,
-  createFilter: (PsiClass, T) -> String): Boolean {
-  return applyTestConfiguration(project, tests, findPsiClass, createFilter) { source ->
+  findSourceElement: (T) -> E?,
+  createFilter: (E, T) -> String): Boolean {
+  return applyTestConfiguration(project, tests, findSourceElement, createFilter) { source ->
     listOf(findTestsTaskToRun(source, project))
   }
 }
 
-fun <T> ExternalSystemTaskExecutionSettings.applyTestConfiguration(
+fun <E : PsiElement, T> ExternalSystemTaskExecutionSettings.applyTestConfiguration(
   project: Project,
   tests: Iterable<T>,
-  findPsiClass: (T) -> PsiClass?,
-  createFilter: (PsiClass, T) -> String,
+  findSourceElement: (T) -> E?,
+  createFilter: (E, T) -> String,
   getTestsTaskToRun: (VirtualFile) -> List<List<String>>
 ): Boolean {
   val projectFileIndex = ProjectFileIndex.SERVICE.getInstance(project)
   val testRunConfigurations = LinkedHashMap<String, Pair<VirtualFile, MutableList<String>>>()
   var module: Module? = null
   for (test in tests) {
-    val psiClass = findPsiClass(test) ?: return false
-    val psiFile = psiClass.containingFile ?: return false
-    val virtualFile = psiFile.virtualFile
-    module = projectFileIndex.getModuleForFile(virtualFile) ?: return false
+    val sourceElement = findSourceElement(test) ?: return false
+    val sourceFile = getSourceFile(sourceElement) ?: return false
+    module = projectFileIndex.getModuleForFile(sourceFile) ?: return false
     if (!GradleRunnerUtil.isGradleModule(module)) return false
-    val (_, arguments) = testRunConfigurations.getOrPut(module.name) { Pair(virtualFile, ArrayList()) }
-    arguments.add(createFilter(psiClass, test))
+    val (_, arguments) = testRunConfigurations.getOrPut(module.name) { Pair(sourceFile, ArrayList()) }
+    arguments.add(createFilter(sourceElement, test))
   }
   if (module == null) return false
   externalProjectPath = GradleRunnerUtil.resolveProjectPath(module) ?: return false
@@ -142,3 +142,13 @@ private fun StringJoiner.addAll(elements: Iterable<String>) = apply {
   }
 }
 
+fun getSourceFile(sourceElement: PsiElement): VirtualFile? {
+  if (sourceElement is PsiFileSystemItem) {
+    return sourceElement.virtualFile
+  }
+  val containingFile = sourceElement.containingFile
+  if (containingFile != null) {
+    return containingFile.virtualFile
+  }
+  return null
+}