junit: Rework support for parallel junit launch.

A previous commit (2c57cbea0d663b) introduced support for
running Junit tests concurrently on all connected devices.
It relied on changing the TestIdentifier to store device
information, and having a single listener that listens to
test events from all devices.

The change to TestIdentifier caused issues with some users
of ddmlib in cases where the device info is not available
all the time.

This patch reverts all the changes in ddmlib, and moves
the knowledge of what device tests are being run to the
listener layer in ADT. So now we have a per device test
listener that knows only about tests run on that device.

Change-Id: Iffedcb38cdf21e349fbe38ecf9a654060b469b04

7 files changed, 146 insertions, 179 deletions

diff --git a/ddms/libs/ddmlib/src/com/android/ddmlib/testrunner/InstrumentationResultParser.java b/ddms/libs/ddmlib/src/com/android/ddmlib/testrunner/InstrumentationResultParser.java
index 2a65699..71f329a 100644
--- a/ddms/libs/ddmlib/src/com/android/ddmlib/testrunner/InstrumentationResultParser.java
+++ b/ddms/libs/ddmlib/src/com/android/ddmlib/testrunner/InstrumentationResultParser.java
@@ -146,9 +146,6 @@ public class InstrumentationResultParser extends MultiLineReceiver {
         }
     }
 
-    /** Device on which this test was run. */
-    private final String mDeviceName;
-
     /** the name to provide to {@link ITestRunListener#testRunStarted(String, int)} */
     private final String mTestRunName;
 
@@ -218,24 +215,10 @@ public class InstrumentationResultParser extends MultiLineReceiver {
      * @param runName the test run name to provide to
      *            {@link ITestRunListener#testRunStarted(String, int)}
      * @param listeners informed of test results as the tests are executing
-     * @param deviceName name of the device on which this test is running, null if unknown
      */
-    public InstrumentationResultParser(String runName, Collection<ITestRunListener> listeners,
-            String deviceName) {
+    public InstrumentationResultParser(String runName, Collection<ITestRunListener> listeners) {
         mTestRunName = runName;
         mTestListeners = new ArrayList<ITestRunListener>(listeners);
-        mDeviceName = deviceName;
-    }
-
-    /**
-     * Creates the InstrumentationResultParser.
-     *
-     * @param runName the test run name to provide to
-     *            {@link ITestRunListener#testRunStarted(String, int)}
-     * @param listeners informed of test results as the tests are executing
-     */
-    public InstrumentationResultParser(String runName, Collection<ITestRunListener> listeners) {
-        this(runName, listeners, null);
     }
 
     /**
@@ -246,7 +229,7 @@ public class InstrumentationResultParser extends MultiLineReceiver {
      * @param listener informed of test results as the tests are executing
      */
     public InstrumentationResultParser(String runName, ITestRunListener listener) {
-        this(runName, Collections.singletonList(listener), null);
+        this(runName, Collections.singletonList(listener));
     }
 
     /**
@@ -459,8 +442,7 @@ public class InstrumentationResultParser extends MultiLineReceiver {
             return;
         }
         reportTestRunStarted(testInfo);
-        TestIdentifier testId = new TestIdentifier(testInfo.mTestClass, testInfo.mTestName,
-                mDeviceName);
+        TestIdentifier testId = new TestIdentifier(testInfo.mTestClass, testInfo.mTestName);
         Map<String, String> metrics;
 
         switch (testInfo.mCode) {
@@ -570,7 +552,7 @@ public class InstrumentationResultParser extends MultiLineReceiver {
             // received test start msg, but not test complete
             // assume test caused this, report as test failure
             TestIdentifier testId = new TestIdentifier(mLastTestResult.mTestClass,
-                    mLastTestResult.mTestName, mDeviceName);
+                    mLastTestResult.mTestName);
             for (ITestRunListener listener : mTestListeners) {
                 listener.testFailed(ITestRunListener.TestFailure.ERROR, testId,
                     String.format("%1$s. Reason: '%2$s'. %3$s", INCOMPLETE_TEST_ERR_MSG_PREFIX,
diff --git a/ddms/libs/ddmlib/src/com/android/ddmlib/testrunner/RemoteAndroidTestRunner.java b/ddms/libs/ddmlib/src/com/android/ddmlib/testrunner/RemoteAndroidTestRunner.java
index 65ed3c0..124df7d 100644
--- a/ddms/libs/ddmlib/src/com/android/ddmlib/testrunner/RemoteAndroidTestRunner.java
+++ b/ddms/libs/ddmlib/src/com/android/ddmlib/testrunner/RemoteAndroidTestRunner.java
@@ -256,7 +256,7 @@ public class RemoteAndroidTestRunner implements IRemoteAndroidTestRunner  {
         Log.i(LOG_TAG, String.format("Running %1$s on %2$s", runCaseCommandStr,
                 mRemoteDevice.getSerialNumber()));
         String runName = mRunName == null ? mPackageName : mRunName;
-        mParser = new InstrumentationResultParser(runName, listeners, mRemoteDevice.getName());
+        mParser = new InstrumentationResultParser(runName, listeners);
 
         try {
             mRemoteDevice.executeShellCommand(runCaseCommandStr, mParser, mMaxTimeToOutputResponse);
diff --git a/ddms/libs/ddmlib/src/com/android/ddmlib/testrunner/TestIdentifier.java b/ddms/libs/ddmlib/src/com/android/ddmlib/testrunner/TestIdentifier.java
index 00295e7..7de5736 100644
--- a/ddms/libs/ddmlib/src/com/android/ddmlib/testrunner/TestIdentifier.java
+++ b/ddms/libs/ddmlib/src/com/android/ddmlib/testrunner/TestIdentifier.java
@@ -23,33 +23,20 @@ public class TestIdentifier {
 
     private final String mClassName;
     private final String mTestName;
-    private final String mDeviceName;
 
     /**
      * Creates a test identifier.
      *
      * @param className fully qualified class name of the test. Cannot be null.
      * @param testName name of the test. Cannot be null.
-     * @param deviceName device on which the test was run.
      */
-    public TestIdentifier(String className, String testName, String deviceName) {
+    public TestIdentifier(String className, String testName) {
         if (className == null || testName == null) {
             throw new IllegalArgumentException("className and testName must " +
                     "be non-null");
         }
         mClassName = className;
         mTestName = testName;
-        mDeviceName = deviceName;
-    }
-
-    /**
-     * Creates a test identifier.
-     *
-     * @param className fully qualified class name of the test. Cannot be null.
-     * @param testName name of the test. Cannot be null.
-     */
-    public TestIdentifier(String className, String testName) {
-        this(className, testName, null);
     }
 
     /**
@@ -66,19 +53,11 @@ public class TestIdentifier {
         return mTestName;
     }
 
-    /**
-     * Returns the name of the device on which the test was run if available, null otherwise.
-     */
-    public String getDeviceName() {
-        return mDeviceName;
-    }
-
     @Override
     public int hashCode() {
         final int prime = 31;
         int result = 1;
         result = prime * result + ((mClassName == null) ? 0 : mClassName.hashCode());
-        result = prime * result + ((mDeviceName == null) ? 0 : mDeviceName.hashCode());
         result = prime * result + ((mTestName == null) ? 0 : mTestName.hashCode());
         return result;
     }
@@ -97,11 +76,6 @@ public class TestIdentifier {
                 return false;
         } else if (!mClassName.equals(other.mClassName))
             return false;
-        if (mDeviceName == null) {
-            if (other.mDeviceName != null)
-                return false;
-        } else if (!mDeviceName.equals(other.mDeviceName))
-            return false;
         if (mTestName == null) {
             if (other.mTestName != null)
                 return false;
@@ -112,12 +86,6 @@ public class TestIdentifier {
 
     @Override
     public String toString() {
-        String deviceName = getDeviceName();
-        String name = String.format("%s#%s", getClassName(), getTestName());
-        if (deviceName != null) {
-            name += String.format(" (%s)", deviceName);
-        }
-
-        return name;
+        return String.format("%s#%s", getClassName(), getTestName());
     }
 }
diff --git a/ddms/libs/ddmlib/tests/src/com/android/ddmlib/testrunner/InstrumentationResultParserTest.java b/ddms/libs/ddmlib/tests/src/com/android/ddmlib/testrunner/InstrumentationResultParserTest.java
index 148f329..478e09e 100644
--- a/ddms/libs/ddmlib/tests/src/com/android/ddmlib/testrunner/InstrumentationResultParserTest.java
+++ b/ddms/libs/ddmlib/tests/src/com/android/ddmlib/testrunner/InstrumentationResultParserTest.java
@@ -40,7 +40,7 @@ public class InstrumentationResultParserTest extends TestCase {
     private static final String CLASS_NAME = "com.test.FooTest";
     private static final String TEST_NAME = "testFoo";
     private static final String STACK_TRACE = "java.lang.AssertionFailedException";
-    private static final TestIdentifier TEST_ID = new TestIdentifier(CLASS_NAME, TEST_NAME, null);
+    private static final TestIdentifier TEST_ID = new TestIdentifier(CLASS_NAME, TEST_NAME);
 
     /**
      * @param name - test name
diff --git a/eclipse/plugins/com.android.ide.eclipse.adt/src/com/android/ide/eclipse/adt/internal/launch/junit/runtime/RemoteAdtTestRunner.java b/eclipse/plugins/com.android.ide.eclipse.adt/src/com/android/ide/eclipse/adt/internal/launch/junit/runtime/RemoteAdtTestRunner.java
index 23b31f0..1205372 100755
--- a/eclipse/plugins/com.android.ide.eclipse.adt/src/com/android/ide/eclipse/adt/internal/launch/junit/runtime/RemoteAdtTestRunner.java
+++ b/eclipse/plugins/com.android.ide.eclipse.adt/src/com/android/ide/eclipse/adt/internal/launch/junit/runtime/RemoteAdtTestRunner.java
@@ -31,6 +31,8 @@ import org.eclipse.core.runtime.IProgressMonitor;
 import org.eclipse.core.runtime.IStatus;
 import org.eclipse.core.runtime.Status;
 import org.eclipse.core.runtime.jobs.Job;
+import org.eclipse.jdt.internal.junit.runner.IListensToTestExecutions;
+import org.eclipse.jdt.internal.junit.runner.ITestReference;
 import org.eclipse.jdt.internal.junit.runner.MessageIds;
 import org.eclipse.jdt.internal.junit.runner.RemoteTestRunner;
 import org.eclipse.jdt.internal.junit.runner.TestExecution;
@@ -90,10 +92,14 @@ public class RemoteAdtTestRunner extends RemoteTestRunner {
      *
      * The tests are run concurrently on all devices. The overall structure is as follows:
      * <ol>
-     *   <li> First, a separate job per device is run to collect test tree data. </li>
-     *   <li> Once all the devices have reported their test tree data, the tree info is
-     *        passed to the Junit UI </li>
-     *   <li> A job per device is again launched to do the actual test run. </li>
+     *   <li> First, a separate job per device is run to collect test tree data. A per device
+     *        {@link TestCollector} records information regarding the tests run on the device.
+     *        </li>
+     *   <li> Once all the devices have finished collecting the test tree data, the tree info is
+     *        collected from all of them and passed to the Junit UI </li>
+     *   <li> A job per device is again launched to do the actual test run. A per device
+     *        {@link TestRunListener} notifies the shared {@link TestResultsNotifier} of test
+     *        status. </li>
      *   <li> As tests complete, the test run listener updates the Junit UI </li>
      * </ol>
      *
@@ -136,13 +142,17 @@ public class RemoteAdtTestRunner extends RemoteTestRunner {
         }
 
         // Launch all test info collector jobs
-        TestCollector collector = new TestCollector();
-        List<TestRunnerJob> collectorJobs = new ArrayList<TestRunnerJob>(devices.size());
+        List<TestTreeCollectorJob> collectorJobs =
+                new ArrayList<TestTreeCollectorJob>(devices.size());
+        List<TestCollector> perDeviceCollectors = new ArrayList<TestCollector>(devices.size());
         for (int i = 0; i < devices.size(); i++) {
             RemoteAndroidTestRunner runner = runners.get(i);
+            String deviceName = devices.get(i).getName();
+            TestCollector collector = new TestCollector(deviceName);
+            perDeviceCollectors.add(collector);
 
             TestTreeCollectorJob job = new TestTreeCollectorJob(
-                    "Test Tree Collector for " + devices.get(i).getName(),
+                    "Test Tree Collector for " + deviceName,
                     runner, mLaunchInfo.isDebugMode(), collector);
             job.setPriority(Job.INTERACTIVE);
             job.schedule();
@@ -151,7 +161,8 @@ public class RemoteAdtTestRunner extends RemoteTestRunner {
         }
 
         // wait for all test info collector jobs to complete
-        for (TestRunnerJob job : collectorJobs) {
+        int totalTests = 0;
+        for (TestTreeCollectorJob job : collectorJobs) {
             try {
                 job.join();
             } catch (InterruptedException e) {
@@ -163,27 +174,34 @@ public class RemoteAdtTestRunner extends RemoteTestRunner {
                 endTestRunWithError(job.getResult().getMessage());
                 return;
             }
-        }
 
-        if (collector.getErrorMessage() != null) {
-            endTestRunWithError(collector.getErrorMessage());
-            return;
+            TestCollector collector = job.getCollector();
+            String err = collector.getErrorMessage();
+            if (err != null) {
+                endTestRunWithError(err);
+                return;
+            }
+
+            totalTests += collector.getTestCaseCount();
         }
 
         AdtPlugin.printToConsole(mLaunchInfo.getProject(), "Sending test information to Eclipse");
-        notifyTestRunStarted(collector.getTestCaseCount() * devices.size());
-        collector.sendTrees(this);
+        notifyTestRunStarted(totalTests);
+        sendTestTrees(perDeviceCollectors);
 
-        TestRunListener testRunListener = new TestRunListener(devices.size());
         List<TestRunnerJob> instrumentationRunnerJobs =
                 new ArrayList<TestRunnerJob>(devices.size());
 
+        TestResultsNotifier notifier = new TestResultsNotifier(mExecution.getListener(),
+                devices.size());
+
         // Spawn all instrumentation runner jobs
         for (int i = 0; i < devices.size(); i++) {
             RemoteAndroidTestRunner runner = runners.get(i);
-
+            String deviceName = devices.get(i).getName();
+            TestRunListener testRunListener = new TestRunListener(deviceName, notifier);
             InstrumentationRunJob job = new InstrumentationRunJob(
-                    "Test Tree Collector for " + devices.get(i).getName(),
+                    "Test Tree Collector for " + deviceName,
                     runner, mLaunchInfo.isDebugMode(), testRunListener);
             job.setPriority(Job.INTERACTIVE);
             job.schedule();
@@ -207,6 +225,14 @@ public class RemoteAdtTestRunner extends RemoteTestRunner {
         }
     }
 
+    /** Sends info about the test tree to be executed (ie the suites and their enclosed tests) */
+    private void sendTestTrees(List<TestCollector> perDeviceCollectors) {
+        for (TestCollector c : perDeviceCollectors) {
+            ITestReference ref = c.getDeviceSuite();
+            ref.sendTree(this);
+        }
+    }
+
     private static abstract class TestRunnerJob extends Job {
         private ITestRunListener mListener;
         private RemoteAndroidTestRunner mRunner;
@@ -258,12 +284,16 @@ public class RemoteAdtTestRunner extends RemoteTestRunner {
             return mIsDebug;
         }
 
+        public ITestRunListener getListener() {
+            return mListener;
+        }
+
         protected abstract void setupRunner();
     }
 
     private static class TestTreeCollectorJob extends TestRunnerJob {
         public TestTreeCollectorJob(String name, RemoteAndroidTestRunner runner, boolean isDebug,
-                ITestRunListener listener) {
+                TestCollector listener) {
             super(name, runner, isDebug, listener);
         }
 
@@ -279,6 +309,10 @@ public class RemoteAdtTestRunner extends RemoteTestRunner {
             // report Binder transaction failures
             runner.addInstrumentationArg(DELAY_MSEC_KEY, COLLECT_TEST_DELAY_MS);
         }
+
+        public TestCollector getCollector() {
+            return (TestCollector) getListener();
+        }
     }
 
     private static class InstrumentationRunJob extends TestRunnerJob {
@@ -346,30 +380,75 @@ public class RemoteAdtTestRunner extends RemoteTestRunner {
     }
 
     /**
-     * TestRunListener that communicates results in real-time back to JDT JUnit
-     */
+     * This class provides the interface to notify the JDT UI regarding the status of tests.
+     * When running tests on multiple devices, there is a {@link TestRunListener} that listens
+     * to results from each device. Rather than all such listeners directly notifying JDT
+     * from different threads, they all notify this class which notifies JDT. In addition,
+     * the {@link #testRunEnded(String, long)} method make sure that JDT is notified that the
+     * test run has completed only when tests on all devices have completed.
+     * */
+    private class TestResultsNotifier {
+        private final IListensToTestExecutions mListener;
+        private final int mDeviceCount;
+
+        private int mCompletedRuns;
+        private long mMaxElapsedTime;
+
+        public TestResultsNotifier(IListensToTestExecutions listener, int nDevices) {
+            mListener = listener;
+            mDeviceCount = nDevices;
+        }
+
+        public synchronized void testEnded(TestCaseReference ref) {
+            mListener.notifyTestEnded(ref);
+        }
+
+        public synchronized void testFailed(TestReferenceFailure ref) {
+            mListener.notifyTestFailed(ref);
+        }
+
+        public synchronized void testRunEnded(String mDeviceName, long elapsedTime) {
+            mCompletedRuns++;
+
+            if (elapsedTime > mMaxElapsedTime) {
+                mMaxElapsedTime = elapsedTime;
+            }
+
+            if (mCompletedRuns == mDeviceCount) {
+                notifyTestRunEnded(mMaxElapsedTime);
+            }
+        }
+
+        public synchronized void testStarted(TestCaseReference testId) {
+            mListener.notifyTestStarted(testId);
+        }
+    }
+
+    /**
+     * TestRunListener that communicates results in real-time back to JDT JUnit via the
+     * {@link TestResultsNotifier}.
+     * */
     private class TestRunListener implements ITestRunListener {
-        private int mRunCount;
+        private final String mDeviceName;
+        private TestResultsNotifier mNotifier;
 
         /**
-         * Constructs a {@link ITestRunListener} that listens for completion of specified
-         * number of runs.
+         * Constructs a {@link ITestRunListener} that listens for test results on given device.
+         * @param deviceName device on which the tests are being run
+         * @param notifier notifier to inform of test status
          */
-        public TestRunListener(int runCount) {
-            mRunCount = runCount;
+        public TestRunListener(String deviceName, TestResultsNotifier notifier) {
+            mDeviceName = deviceName;
+            mNotifier = notifier;
         }
 
         @Override
-        public synchronized void testEnded(TestIdentifier test,
-                Map<String, String> ignoredTestMetrics) {
-            mExecution.getListener().notifyTestEnded(new TestCaseReference(test));
+        public void testEnded(TestIdentifier test, Map<String, String> ignoredTestMetrics) {
+            mNotifier.testEnded(new TestCaseReference(mDeviceName, test));
         }
 
-        /* (non-Javadoc)
-         * @see com.android.ddmlib.testrunner.ITestRunListener#testFailed(com.android.ddmlib.testrunner.ITestRunListener.TestFailure, com.android.ddmlib.testrunner.TestIdentifier, java.lang.String)
-         */
         @Override
-        public synchronized void testFailed(TestFailure status, TestIdentifier test, String trace) {
+        public void testFailed(TestFailure status, TestIdentifier test, String trace) {
             String statusString;
             if (status == TestFailure.ERROR) {
                 statusString = MessageIds.TEST_ERROR;
@@ -377,47 +456,28 @@ public class RemoteAdtTestRunner extends RemoteTestRunner {
                 statusString = MessageIds.TEST_FAILED;
             }
             TestReferenceFailure failure =
-                new TestReferenceFailure(new TestCaseReference(test),
+                new TestReferenceFailure(new TestCaseReference(mDeviceName, test),
                         statusString, trace, null);
-            mExecution.getListener().notifyTestFailed(failure);
+            mNotifier.testFailed(failure);
         }
 
-        /* (non-Javadoc)
-         * @see com.android.ddmlib.testrunner.ITestRunListener#testRunEnded(long, Map<String, String>)
-         */
         @Override
         public synchronized void testRunEnded(long elapsedTime, Map<String, String> runMetrics) {
-            mRunCount--;
-
-            if (mRunCount > 0) {
-                return;
-            }
-
-            // notify that test run has completed only after all runs have been completed
-            notifyTestRunEnded(elapsedTime);
+            mNotifier.testRunEnded(mDeviceName, elapsedTime);
             AdtPlugin.printToConsole(mLaunchInfo.getProject(),
                     LaunchMessages.RemoteAdtTestRunner_RunCompleteMsg);
         }
 
-        /* (non-Javadoc)
-         * @see com.android.ddmlib.testrunner.ITestRunListener#testRunFailed(java.lang.String)
-         */
         @Override
         public synchronized void testRunFailed(String errorMessage) {
             reportError(errorMessage);
         }
 
-        /* (non-Javadoc)
-         * @see com.android.ddmlib.testrunner.ITestRunListener#testRunStarted(int)
-         */
         @Override
         public synchronized void testRunStarted(String runName, int testCount) {
             // ignore
         }
 
-        /* (non-Javadoc)
-         * @see com.android.ddmlib.testrunner.ITestRunListener#testRunStopped(long)
-         */
         @Override
         public synchronized void testRunStopped(long elapsedTime) {
             notifyTestRunStopped(elapsedTime);
@@ -425,19 +485,14 @@ public class RemoteAdtTestRunner extends RemoteTestRunner {
                     LaunchMessages.RemoteAdtTestRunner_RunStoppedMsg);
         }
 
-        /* (non-Javadoc)
-         * @see com.android.ddmlib.testrunner.ITestRunListener#testStarted(com.android.ddmlib.testrunner.TestIdentifier)
-         */
         @Override
         public synchronized void testStarted(TestIdentifier test) {
-            TestCaseReference testId = new TestCaseReference(test);
-            mExecution.getListener().notifyTestStarted(testId);
+            TestCaseReference testId = new TestCaseReference(mDeviceName, test);
+            mNotifier.testStarted(testId);
         }
     }
 
-    /**
-     * Override parent to get extra logs.
-     */
+    /** Override parent to get extra logs. */
     @Override
     protected boolean connect() {
         boolean result = super.connect();
@@ -448,9 +503,7 @@ public class RemoteAdtTestRunner extends RemoteTestRunner {
         return result;
     }
 
-    /**
-     * Override parent to dump error message to console.
-     */
+    /** Override parent to dump error message to console. */
     @Override
     public void runFailed(String message, Exception exception) {
         if (exception != null) {
diff --git a/eclipse/plugins/com.android.ide.eclipse.adt/src/com/android/ide/eclipse/adt/internal/launch/junit/runtime/TestCaseReference.java b/eclipse/plugins/com.android.ide.eclipse.adt/src/com/android/ide/eclipse/adt/internal/launch/junit/runtime/TestCaseReference.java
index a81e266..e05e9b8 100644
--- a/eclipse/plugins/com.android.ide.eclipse.adt/src/com/android/ide/eclipse/adt/internal/launch/junit/runtime/TestCaseReference.java
+++ b/eclipse/plugins/com.android.ide.eclipse.adt/src/com/android/ide/eclipse/adt/internal/launch/junit/runtime/TestCaseReference.java
@@ -25,26 +25,18 @@ import org.eclipse.jdt.internal.junit.runner.IVisitsTestTrees;
  */
 @SuppressWarnings("restriction")
 class TestCaseReference extends AndroidTestReference {
-
     private final String mClassName;
     private final String mTestName;
     private final String mDeviceName;
 
     /**
-     * Creates a TestCaseReference from a class and method name
-     */
-    TestCaseReference(String className, String testName, String deviceName) {
-        mClassName = className;
-        mTestName = testName;
-        mDeviceName = deviceName == null ? "?" : deviceName;        //$NON-NLS-1$
-    }
-
-    /**
      * Creates a TestCaseReference from a {@link TestIdentifier}
      * @param test
      */
-    TestCaseReference(TestIdentifier test) {
-        this(test.getClassName(), test.getTestName(), test.getDeviceName());
+    TestCaseReference(String deviceName, TestIdentifier test) {
+        mDeviceName = deviceName;
+        mClassName = test.getClassName();
+        mTestName = test.getTestName();
     }
 
     /**
diff --git a/eclipse/plugins/com.android.ide.eclipse.adt/src/com/android/ide/eclipse/adt/internal/launch/junit/runtime/TestCollector.java b/eclipse/plugins/com.android.ide.eclipse.adt/src/com/android/ide/eclipse/adt/internal/launch/junit/runtime/TestCollector.java
index 91954b4..806f289 100644
--- a/eclipse/plugins/com.android.ide.eclipse.adt/src/com/android/ide/eclipse/adt/internal/launch/junit/runtime/TestCollector.java
+++ b/eclipse/plugins/com.android.ide.eclipse.adt/src/com/android/ide/eclipse/adt/internal/launch/junit/runtime/TestCollector.java
@@ -19,26 +19,25 @@ package com.android.ide.eclipse.adt.internal.launch.junit.runtime;
 import com.android.ddmlib.testrunner.ITestRunListener;
 import com.android.ddmlib.testrunner.TestIdentifier;
 
-import org.eclipse.jdt.internal.junit.runner.ITestReference;
-import org.eclipse.jdt.internal.junit.runner.IVisitsTestTrees;
-
-import java.util.HashMap;
 import java.util.Map;
 
 /**
  * Collects info about tests to be executed by listening to the results of an Android test run.
  */
-@SuppressWarnings("restriction")
 class TestCollector implements ITestRunListener {
+    private final String mDeviceName;
+    private final TestSuiteReference mDeviceSuiteRef;
 
     private int mTotalTestCount;
     /** test name to test suite reference map. */
-    private Map<String, TestSuiteReference> mTestTree;
+
     private String mErrorMessage = null;
 
-    TestCollector() {
+    TestCollector(String deviceName) {
+        mDeviceName = deviceName;
+        mDeviceSuiteRef = new TestSuiteReference(deviceName);
+
         mTotalTestCount = 0;
-        mTestTree = new HashMap<String, TestSuiteReference>();
     }
 
     @Override
@@ -91,33 +90,13 @@ class TestCollector implements ITestRunListener {
      */
     @Override
     public synchronized void testStarted(TestIdentifier test) {
-        TestSuiteReference suiteRef;
-        TestSuiteReference deviceSuiteRef;
-
-        String deviceName = test.getDeviceName();
-        if (deviceName != null) {
-            // if the device name is available, nest the test under a per device test suite
-            deviceSuiteRef = mTestTree.get(deviceName);
-            if (deviceSuiteRef == null) {
-                deviceSuiteRef = new TestSuiteReference(deviceName);
-                mTestTree.put(deviceName, deviceSuiteRef);
-            }
-
-            suiteRef = deviceSuiteRef.getTestSuite(test.getClassName());
-            if (suiteRef == null) {
-                suiteRef = new TestSuiteReference(test.getClassName());
-                deviceSuiteRef.addTest(suiteRef);
-            }
-        } else {
-            suiteRef = mTestTree.get(test.getClassName());
-            if (suiteRef == null) {
-                // this test suite has not been seen before, create it
-                suiteRef = new TestSuiteReference(test.getClassName());
-                mTestTree.put(test.getClassName(), suiteRef);
-            }
+        TestSuiteReference suiteRef = mDeviceSuiteRef.getTestSuite(test.getClassName());
+        if (suiteRef == null) {
+            suiteRef = new TestSuiteReference(test.getClassName());
+            mDeviceSuiteRef.addTest(suiteRef);
         }
 
-        suiteRef.addTest(new TestCaseReference(test));
+        suiteRef.addTest(new TestCaseReference(mDeviceName, test));
     }
 
     /**
@@ -128,21 +107,14 @@ class TestCollector implements ITestRunListener {
     }
 
     /**
-     * Sends info about the test tree to be executed (ie the suites and their enclosed tests)
-     *
-     * @param notified the {@link IVisitsTestTrees} to send test data to
-     */
-    public synchronized void sendTrees(IVisitsTestTrees notified) {
-        for (ITestReference ref : mTestTree.values()) {
-            ref.sendTree(notified);
-        }
-    }
-
-    /**
      * Returns the error message that was reported when collecting test info.
      * Returns <code>null</code> if no error occurred.
      */
     public synchronized String getErrorMessage() {
         return mErrorMessage;
     }
+
+    public TestSuiteReference getDeviceSuite() {
+        return mDeviceSuiteRef;
+    }
 }