From 6594af0dc1ca896da4ce90b625128756e82a6fcb Mon Sep 17 00:00:00 2001
From: Anna Kozlova <anna.kozlova@jetbrains.com>
Date: Wed, 16 Dec 2015 13:41:53 +0100
Subject: [PATCH] tests: allow tests selection from file with multiple test
 classes (IDEA-148273)

---
 .../AbstractPatternBasedConfigurationProducer.java  | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

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 719d8d20763f..1275cf8d5af8 100644
--- a/java/execution/impl/src/com/intellij/execution/testframework/AbstractPatternBasedConfigurationProducer.java
+++ b/java/execution/impl/src/com/intellij/execution/testframework/AbstractPatternBasedConfigurationProducer.java
@@ -34,6 +34,7 @@ import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.*;
 import com.intellij.psi.search.PsiElementProcessor;
 import com.intellij.psi.util.ClassUtil;
+import com.intellij.psi.util.PsiTreeUtil;
 
 import java.util.ArrayList;
 import java.util.LinkedHashSet;
@@ -163,6 +164,7 @@ public abstract class AbstractPatternBasedConfigurationProducer<T extends Module
       return true;
     }
     else {
+      final PsiElement element = CommonDataKeys.PSI_ELEMENT.getData(dataContext);
       final VirtualFile[] files = CommonDataKeys.VIRTUAL_FILE_ARRAY.getData(dataContext);
       if (files != null) {
         Project project = CommonDataKeys.PROJECT.getData(dataContext);
@@ -171,7 +173,16 @@ public abstract class AbstractPatternBasedConfigurationProducer<T extends Module
           for (VirtualFile file : files) {
             final PsiFile psiFile = psiManager.findFile(file);
             if (psiFile instanceof PsiClassOwner) {
-              collectTestMembers(((PsiClassOwner)psiFile).getClasses(), checkAbstract, checkIsTest, processor);
+              PsiClass[] psiClasses = ((PsiClassOwner)psiFile).getClasses();
+              if (element != null && psiClasses.length > 0) {
+                for (PsiClass aClass : psiClasses) {
+                  if (PsiTreeUtil.isAncestor(aClass, element, false)) {
+                    psiClasses = new PsiClass[] {aClass};
+                    break;
+                  }
+                }
+              }
+              collectTestMembers(psiClasses, checkAbstract, checkIsTest, processor);
               for (PsiElement psiMember : processor.getCollection()) {
                 classes.add(((PsiClass)psiMember).getQualifiedName());
               }