junit: track current test; still expand if user selects one

plugins/junit/src/com/intellij/execution/junit2/ui/actions/RunningTestTracker.java

@@ -16,17 +16,26 @@
 
 package com.intellij.execution.junit2.ui.actions;
 
+import com.intellij.execution.junit2.TestProxy;
 import com.intellij.execution.junit2.events.StateChangedEvent;
 import com.intellij.execution.junit2.events.TestEvent;
-import com.intellij.execution.junit2.TestProxy;
 import com.intellij.execution.junit2.ui.model.JUnitAdapter;
 import com.intellij.execution.junit2.ui.model.JUnitRunningModel;
 import com.intellij.execution.junit2.ui.properties.JUnitConsoleProperties;
 import com.intellij.execution.testframework.TestFrameworkPropertyListener;
 import com.intellij.execution.testframework.actions.TestFrameworkActions;
+import com.intellij.ide.util.treeView.NodeDescriptor;
+import com.intellij.openapi.Disposable;
 import com.intellij.openapi.diagnostic.Logger;
+import com.intellij.openapi.util.Disposer;
 import com.intellij.rt.execution.junit.states.PoolOfTestStates;
 
+import javax.swing.*;
+import javax.swing.tree.DefaultMutableTreeNode;
+import javax.swing.tree.TreePath;
+import java.awt.event.MouseAdapter;
+import java.awt.event.MouseEvent;
+
 class RunningTestTracker extends JUnitAdapter implements TestFrameworkPropertyListener<Boolean> {
   private static final Logger LOG = Logger.getInstance("#com.intellij.execution.junit2.ui.actions.TrackRunningTestAction");
 
@@ -37,6 +46,29 @@ class RunningTestTracker extends JUnitAdapter implements TestFrameworkPropertyLi
 
   private RunningTestTracker(final JUnitRunningModel model) {
     myModel = model;
+    final JTree tree = myModel.getTree();
+    final MouseAdapter userSelectionListener = new MouseAdapter() {
+      @Override
+      public void mousePressed(MouseEvent e) {
+        final TreePath treePath = tree.getPathForLocation(e.getX(), e.getY());
+        if (treePath != null) {
+          final Object component = treePath.getLastPathComponent();
+          if (component instanceof DefaultMutableTreeNode) {
+            final Object userObject = ((DefaultMutableTreeNode)component).getUserObject();
+            if (userObject instanceof NodeDescriptor) {
+              myLastSelected = (TestProxy)((NodeDescriptor)userObject).getElement();
+            }
+          }
+        }
+      }
+    };
+    tree.addMouseListener(userSelectionListener);
+    Disposer.register(myModel, new Disposable() {
+      @Override
+      public void dispose() {
+        tree.removeMouseListener(userSelectionListener);
+      }
+    });
     choosePolicy();
   }
 
@@ -45,13 +77,6 @@ class RunningTestTracker extends JUnitAdapter implements TestFrameworkPropertyLi
     myTrackingPolicy.apply();
   }
 
-  @Override
-  public void onTestSelected(TestProxy test) {
-    if (test != null && test.isLeaf()) {
-      myLastSelected = test;
-    }
-  }
-
   public void onTestChanged(final TestEvent event) {
     if (event instanceof StateChangedEvent) {
       final TestProxy proxy = event.getSource();
@@ -105,8 +130,14 @@ class RunningTestTracker extends JUnitAdapter implements TestFrameworkPropertyLi
     }
 
     private void selectLastTest() {
-      if (myLastRan != null && isRunningState(myLastRan) && myLastSelected == null)
-        myModel.selectTest(myLastRan);
+      if (myLastRan != null && isRunningState(myLastRan)) {
+        if (myLastSelected == null) {
+          myModel.selectTest(myLastRan);
+        }
+        else {
+          myModel.expandTest(myLastRan);
+        }
+      }
     }
   };
 }

plugins/junit/src/com/intellij/execution/junit2/ui/model/JUnitRunningModel.java

@@ -79,6 +79,10 @@ public class JUnitRunningModel implements TestFrameworkRunningModel {
     selectTest(getRoot());
     myTreeListener.install();
   }
+  
+  public JTree getTree() {
+    return myTreeView;
+  }
 
   public void dispose() {
   }
@@ -116,6 +120,12 @@ public class JUnitRunningModel implements TestFrameworkRunningModel {
     myTreeBuilder.select(test, null);
   }
 
+  public void expandTest(final TestProxy test) {
+    if (test == null) return;
+
+    myTreeBuilder.expand(test, null);
+  }
+
 
   public void collapse(final TestProxy test) {
     if (test == getRoot())