Merge "Introduce testModuleStart/End callbacks" into oc-dev am: 82e24a3a7c

am: 71fa109b10

Change-Id: Ice8bff46082f8781c4477db28f61f85b2755133f

5 files changed, 71 insertions, 2 deletions

diff --git a/src/com/android/tradefed/invoker/ShardListener.java b/src/com/android/tradefed/invoker/ShardListener.java
index ea0f1fb31..8aa56251c 100644
--- a/src/com/android/tradefed/invoker/ShardListener.java
+++ b/src/com/android/tradefed/invoker/ShardListener.java
@@ -113,7 +113,19 @@ public class ShardListener extends CollectingTestListener {
         super.invocationEnded(elapsedTime);
         synchronized (mMasterListener) {
             logShardContent(getRunResults());
+            IInvocationContext moduleContext = null;
             for (TestRunResult runResult : getRunResults()) {
+                // Stop or start the module
+                if (moduleContext != null
+                        && !getModuleContextForRunResult(runResult).equals(moduleContext)) {
+                    mMasterListener.testModuleEnded();
+                    moduleContext = null;
+                }
+                if (moduleContext == null && getModuleContextForRunResult(runResult) != null) {
+                    moduleContext = getModuleContextForRunResult(runResult);
+                    mMasterListener.testModuleStarted(moduleContext);
+                }
+
                 mMasterListener.testRunStarted(runResult.getName(), runResult.getNumTests());
                 forwardTestResults(runResult.getTestResults());
                 if (runResult.isRunFailure()) {
@@ -121,6 +133,11 @@ public class ShardListener extends CollectingTestListener {
                 }
                 mMasterListener.testRunEnded(runResult.getElapsedTime(), runResult.getRunMetrics());
             }
+            // Close the last module
+            if (moduleContext != null) {
+                mMasterListener.testModuleEnded();
+                moduleContext = null;
+            }
             mMasterListener.invocationEnded(elapsedTime);
         }
     }
diff --git a/src/com/android/tradefed/result/CollectingTestListener.java b/src/com/android/tradefed/result/CollectingTestListener.java
index 80d5987eb..edecf7258 100644
--- a/src/com/android/tradefed/result/CollectingTestListener.java
+++ b/src/com/android/tradefed/result/CollectingTestListener.java
@@ -42,7 +42,10 @@ public class CollectingTestListener implements ITestInvocationListener {
     // Uses a LinkedHashmap to have predictable iteration order
     private Map<String, TestRunResult> mRunResultsMap =
         Collections.synchronizedMap(new LinkedHashMap<String, TestRunResult>());
+    private Map<TestRunResult, IInvocationContext> mModuleContextMap =
+            Collections.synchronizedMap(new LinkedHashMap<TestRunResult, IInvocationContext>());
     private TestRunResult mCurrentResults =  new TestRunResult();
+    private IInvocationContext mCurrentModuleContext = null;
 
     /** represents sums of tests in each TestStatus state for all runs.
      * Indexed by TestStatus.ordinal() */
@@ -116,6 +119,16 @@ public class CollectingTestListener implements ITestInvocationListener {
         mBuildInfo = buildInfo;
     }
 
+    @Override
+    public void testModuleStarted(IInvocationContext moduleContext) {
+        mCurrentModuleContext = moduleContext;
+    }
+
+    @Override
+    public void testModuleEnded() {
+        mCurrentModuleContext = null;
+    }
+
     /**
      * {@inheritDoc}
      */
@@ -130,6 +143,10 @@ public class CollectingTestListener implements ITestInvocationListener {
             mCurrentResults.setAggregateMetrics(mIsAggregateMetrics);
 
             mRunResultsMap.put(name, mCurrentResults);
+            // track the module context associated with the results.
+            if (mCurrentModuleContext != null) {
+                mModuleContextMap.put(mCurrentResults, mCurrentModuleContext);
+            }
         }
         mCurrentResults.testRunStarted(name, numTests);
         mIsCountDirty = true;
@@ -232,6 +249,14 @@ public class CollectingTestListener implements ITestInvocationListener {
         return mRunResultsMap.values();
     }
 
+    /**
+     * Returns the {@link IInvocationContext} of the module associated with the results or null if
+     * it was not associated with any module.
+     */
+    public IInvocationContext getModuleContextForRunResult(TestRunResult res) {
+        return mModuleContextMap.get(res);
+    }
+
     /** Returns True if the result map already has an entry for the run name. */
     public boolean hasResultFor(String runName) {
         return mRunResultsMap.containsKey(runName);
diff --git a/src/com/android/tradefed/result/ITestInvocationListener.java b/src/com/android/tradefed/result/ITestInvocationListener.java
index c8fbd7c2d..5e951d8db 100644
--- a/src/com/android/tradefed/result/ITestInvocationListener.java
+++ b/src/com/android/tradefed/result/ITestInvocationListener.java
@@ -20,6 +20,7 @@ import com.android.ddmlib.testrunner.TestIdentifier;
 import com.android.tradefed.command.ICommandScheduler;
 import com.android.tradefed.invoker.IInvocationContext;
 import com.android.tradefed.log.ITestLogger;
+import com.android.tradefed.testtype.suite.ITestSuite;
 
 import java.util.Map;
 
@@ -95,6 +96,18 @@ public interface ITestInvocationListener extends ITestRunListener, ITestLogger {
     }
 
     /**
+     * Reports the beginning of a module running. This callback is associated with {@link
+     * #testModuleEnded()} and is optional in the sequence. It is only used during a run that uses
+     * modules: {@link ITestSuite} based runners.
+     *
+     * @param moduleContext the {@link IInvocationContext} of the module.
+     */
+    public default void testModuleStarted(IInvocationContext moduleContext) {}
+
+    /** Reports the end of a module run. */
+    public default void testModuleEnded() {}
+
+    /**
      * {@inheritDoc}
      */
     @Override
diff --git a/src/com/android/tradefed/result/ResultForwarder.java b/src/com/android/tradefed/result/ResultForwarder.java
index 494fc7aef..41a3e8809 100644
--- a/src/com/android/tradefed/result/ResultForwarder.java
+++ b/src/com/android/tradefed/result/ResultForwarder.java
@@ -297,4 +297,18 @@ public class ResultForwarder implements ITestInvocationListener {
             }
         }
     }
+
+    @Override
+    public void testModuleStarted(IInvocationContext moduleContext) {
+        for (ITestInvocationListener listener : mListeners) {
+            listener.testModuleStarted(moduleContext);
+        }
+    }
+
+    @Override
+    public void testModuleEnded() {
+        for (ITestInvocationListener listener : mListeners) {
+            listener.testModuleEnded();
+        }
+    }
 }
diff --git a/src/com/android/tradefed/testtype/suite/ITestSuite.java b/src/com/android/tradefed/testtype/suite/ITestSuite.java
index cdef58310..90ed34c7c 100644
--- a/src/com/android/tradefed/testtype/suite/ITestSuite.java
+++ b/src/com/android/tradefed/testtype/suite/ITestSuite.java
@@ -233,7 +233,7 @@ public abstract class ITestSuite
                 }
 
                 try {
-                    mContext.setModuleInvocationContext(module.getModuleInvocationContext());
+                    listener.testModuleStarted(module.getModuleInvocationContext());
                     // Populate the module context with devices and builds
                     for (String deviceName : mContext.getDeviceConfigNames()) {
                         module.getModuleInvocationContext()
@@ -245,7 +245,7 @@ public abstract class ITestSuite
                 } finally {
                     // clear out module invocation context since we are now done with module
                     // execution
-                    mContext.setModuleInvocationContext(null);
+                    listener.testModuleEnded();
                 }
             }
         } catch (DeviceNotAvailableException e) {