From 531ac706b2705ab41d7c9cbe870087d86884894c Mon Sep 17 00:00:00 2001
From: Sergey Ignatov <sergey.ignatov@jetbrains.com>
Date: Tue, 29 May 2018 16:15:59 +0300
Subject: [PATCH] fix test discovery's tree

---
 .../actions/DiscoveredTestsTree.java          | 25 +++++++++++++------
 .../actions/DiscoveredTestsTreeModel.java     |  4 +--
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/java/execution/impl/src/com/intellij/execution/testDiscovery/actions/DiscoveredTestsTree.java b/java/execution/impl/src/com/intellij/execution/testDiscovery/actions/DiscoveredTestsTree.java
index ca4c0f54af1c..8cbffef6b2cc 100644
--- a/java/execution/impl/src/com/intellij/execution/testDiscovery/actions/DiscoveredTestsTree.java
+++ b/java/execution/impl/src/com/intellij/execution/testDiscovery/actions/DiscoveredTestsTree.java
@@ -30,6 +30,7 @@ import javax.swing.tree.TreeModel;
 import javax.swing.tree.TreePath;
 import javax.swing.tree.TreeSelectionModel;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -39,7 +40,7 @@ class DiscoveredTestsTree extends Tree implements DataProvider {
   public DiscoveredTestsTree(String title) {
     myModel = new DiscoveredTestsTreeModel();
     setModel(new AsyncTreeModel(myModel));
-    HintUpdateSupply.installSimpleHintUpdateSupply(this);
+    HintUpdateSupply.installHintUpdateSupply(this, DiscoveredTestsTree::obj2psi);
     TreeUIHelper.getInstance().installTreeSpeedSearch(this, o -> {
       Object component = o.getLastPathComponent();
       return component instanceof PsiMember ? ((PsiMember)component).getName() : null;
@@ -116,9 +117,18 @@ class DiscoveredTestsTree extends Tree implements DataProvider {
     return myModel.getTestMethods();
   }
 
+  @Nullable
   public PsiElement getSelectedElement() {
     TreePath path = getSelectionModel().getSelectionPath();
-    return ObjectUtils.tryCast(path == null ? null : path.getLastPathComponent(), PsiElement.class);
+    return obj2psi(path == null ? null : path.getLastPathComponent());
+  }
+
+  @Nullable
+  private static PsiElement obj2psi(@Nullable Object obj) {
+    return Optional.ofNullable(ObjectUtils.tryCast(obj, DiscoveredTestsTreeModel.Node.class))
+                   .map(n -> n.getPointer())
+                   .map(p -> p.getElement())
+                   .orElse(null);
   }
 
   public int getTestCount() {
@@ -133,18 +143,19 @@ class DiscoveredTestsTree extends Tree implements DataProvider {
       List<PsiElement> result = ContainerUtil.newSmartList();
       TreeModel model = getModel();
       for (TreePath p : paths) {
-        Object e = p.getLastPathComponent();
+        Object o = p.getLastPathComponent();
+        PsiElement e = obj2psi(o);
         if (e instanceof PsiMethod) {
-          result.add((PsiMethod)e);
+          result.add(e);
         }
         else {
           int count = model.getChildCount(e);
-          if (count == 0 && e instanceof PsiElement) {
-            result.add((PsiElement)e);
+          if (count == 0 && e != null) {
+            result.add(e);
           }
           else {
             for (int i = 0; i < count; i++) {
-              ContainerUtil.addIfNotNull(result, ObjectUtils.tryCast(model.getChild(e, i), PsiMethod.class));
+              ContainerUtil.addIfNotNull(result, ObjectUtils.tryCast(obj2psi(model.getChild(e, i)), PsiMethod.class));
             }
           }
         }
diff --git a/java/execution/impl/src/com/intellij/execution/testDiscovery/actions/DiscoveredTestsTreeModel.java b/java/execution/impl/src/com/intellij/execution/testDiscovery/actions/DiscoveredTestsTreeModel.java
index 484cfa98b718..09fc877100ac 100644
--- a/java/execution/impl/src/com/intellij/execution/testDiscovery/actions/DiscoveredTestsTreeModel.java
+++ b/java/execution/impl/src/com/intellij/execution/testDiscovery/actions/DiscoveredTestsTreeModel.java
@@ -98,7 +98,7 @@ class DiscoveredTestsTreeModel extends BaseTreeModel<Object> {
     return myTests.values().stream().mapToInt(ms -> ms.size()).sum();
   }
 
-  static abstract class Node<Psi extends PsiMember> {
+  public static abstract class Node<Psi extends PsiMember> {
     @NotNull
     private final SmartPsiElementPointer<Psi> myPointer;
     private final String myName;
@@ -140,7 +140,7 @@ class DiscoveredTestsTreeModel extends BaseTreeModel<Object> {
     }
 
     @NotNull
-    SmartPsiElementPointer<Psi> getPointer() {
+    public SmartPsiElementPointer<Psi> getPointer() {
       return myPointer;
     }