From 9bd5077e94450cd0b837429e45a3635bf387cb89 Mon Sep 17 00:00:00 2001
From: "Dmitry.Krasilschikov" <kradima@jetbrains.com>
Date: Sat, 17 Dec 2016 20:14:57 +0300
Subject: [PATCH] RUBY-5449 find class related tests and vise versa for
 appropriate element

---
 .../src/com/intellij/testIntegration/TestFinder.java        | 5 +++++
 .../src/com/intellij/testIntegration/TestFinderHelper.java  | 6 ++++--
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/platform/lang-api/src/com/intellij/testIntegration/TestFinder.java b/platform/lang-api/src/com/intellij/testIntegration/TestFinder.java
index d956d465398e..f62c477a5ab9 100644
--- a/platform/lang-api/src/com/intellij/testIntegration/TestFinder.java
+++ b/platform/lang-api/src/com/intellij/testIntegration/TestFinder.java
@@ -35,4 +35,9 @@ public interface TestFinder {
   Collection<PsiElement> findClassesForTest(@NotNull PsiElement element);
 
   boolean isTest(@NotNull PsiElement element);
+
+  @Nullable
+  default PsiElement findSelectedElement(@NotNull final PsiElement element) {
+    return element;
+  }
 }
diff --git a/platform/lang-impl/src/com/intellij/testIntegration/TestFinderHelper.java b/platform/lang-impl/src/com/intellij/testIntegration/TestFinderHelper.java
index 67318ceb1238..675cce74e813 100644
--- a/platform/lang-impl/src/com/intellij/testIntegration/TestFinderHelper.java
+++ b/platform/lang-impl/src/com/intellij/testIntegration/TestFinderHelper.java
@@ -39,7 +39,8 @@ public class TestFinderHelper {
   public static Collection<PsiElement> findTestsForClass(PsiElement element) {
     Collection<PsiElement> result = new LinkedHashSet<>();
     for (TestFinder each : getFinders()) {
-      result.addAll(each.findTestsForClass(element));
+      final PsiElement selectedElement = each.findSelectedElement(element);
+      if (selectedElement != null) result.addAll(each.findTestsForClass(selectedElement));
     }
     return result;
   }
@@ -47,7 +48,8 @@ public class TestFinderHelper {
   public static Collection<PsiElement> findClassesForTest(PsiElement element) {
     Collection<PsiElement> result = new LinkedHashSet<>();
     for (TestFinder each : getFinders()) {
-      result.addAll(each.findClassesForTest(element));
+      final PsiElement selectedElement = each.findSelectedElement(element);
+      if (selectedElement != null) result.addAll(each.findClassesForTest(selectedElement));
     }
     return result;
   }