Add onTestEnd with test description

To add log files with names that references the test name, it's
necessary to add a TestDirection argument to the onTestEnd function. By
default, the original definition of onTestEnd will be used. If the
version of onTestEnd with a TestDescription is overridden, the other
overloaded definition will be ignored.

Bug: 134790621
Test: tradefed.sh run commandAndExit host \
--class com.android.tradefed.device.metric.BaseDeviceMetricCollectorTest \
--method testOnTestEndWithTestDescription -l verbose
Change-Id: Ia748d5c073fcd00b2f29821c2d722d5b4c56ee9a
Signed-off-by: Maggie White <maggiewhite@google.com>
Merged-In: Ia748d5c073fcd00b2f29821c2d722d5b4c56ee9a

3 files changed, 69 insertions, 1 deletions

diff --git a/src/com/android/tradefed/device/metric/BaseDeviceMetricCollector.java b/src/com/android/tradefed/device/metric/BaseDeviceMetricCollector.java
index 1e4d16ca0..10885cf2f 100644
--- a/src/com/android/tradefed/device/metric/BaseDeviceMetricCollector.java
+++ b/src/com/android/tradefed/device/metric/BaseDeviceMetricCollector.java
@@ -130,6 +130,15 @@ public class BaseDeviceMetricCollector implements IMetricCollector {
         // Does nothing
     }
 
+    @Override
+    public void onTestEnd(
+            DeviceMetricData testData,
+            final Map<String, Metric> currentTestCaseMetrics,
+            TestDescription test) {
+        // Call the default implementation of onTestEnd if not overridden
+        onTestEnd(testData, currentTestCaseMetrics);
+    }
+
     /** =================================== */
     /** Invocation Listeners for forwarding */
     @Override
@@ -233,7 +242,7 @@ public class BaseDeviceMetricCollector implements IMetricCollector {
             TestDescription test, long endTime, HashMap<String, Metric> testMetrics) {
         if (!mSkipTestCase) {
             try {
-                onTestEnd(mTestData, testMetrics);
+                onTestEnd(mTestData, testMetrics, test);
                 mTestData.addToMetrics(testMetrics);
             } catch (Throwable t) {
                 // Prevent exception from messing up the status reporting.
diff --git a/src/com/android/tradefed/device/metric/IMetricCollector.java b/src/com/android/tradefed/device/metric/IMetricCollector.java
index 387e53678..1c8b84fad 100644
--- a/src/com/android/tradefed/device/metric/IMetricCollector.java
+++ b/src/com/android/tradefed/device/metric/IMetricCollector.java
@@ -113,4 +113,18 @@ public interface IMetricCollector extends ITestInvocationListener, IDisableable
      */
     public void onTestEnd(
             DeviceMetricData testData, final Map<String, Metric> currentTestCaseMetrics);
+
+    /**
+     * Callback when a test case is ended. This should be the time for clean up.
+     *
+     * @param testData the {@link DeviceMetricData} holding the data for the test case. Will be the
+     *     same object as during {@link #onTestStart(DeviceMetricData)}.
+     * @param currentTestCaseMetrics the current map of metrics passed to {@link
+     *     #testEnded(TestDescription, Map)}.
+     * @param test the {@link TestDescription} of the test case in progress.
+     */
+    public void onTestEnd(
+            DeviceMetricData testData,
+            final Map<String, Metric> currentTestCaseMetrics,
+            TestDescription test);
 }
diff --git a/tests/src/com/android/tradefed/device/metric/BaseDeviceMetricCollectorTest.java b/tests/src/com/android/tradefed/device/metric/BaseDeviceMetricCollectorTest.java
index dd6bde56a..70bc6b7f2 100644
--- a/tests/src/com/android/tradefed/device/metric/BaseDeviceMetricCollectorTest.java
+++ b/tests/src/com/android/tradefed/device/metric/BaseDeviceMetricCollectorTest.java
@@ -485,4 +485,49 @@ public class BaseDeviceMetricCollectorTest {
         assertTrue(allValues.get(2).containsKey("onteststart"));
         assertTrue(allValues.get(2).containsKey("ontestend"));
     }
+
+    /**
+     * Test that onTestEnd with TestDescription formal supercedes the method signature without a
+     * TestDescription.
+     */
+    @Test
+    public void testOnTestEndWithTestDescription() throws Exception {
+        mBase =
+                new TwoMetricsBaseCollector() {
+                    @Override
+                    public void onTestEnd(
+                            DeviceMetricData testData,
+                            final Map<String, Metric> currentTestCaseMetrics,
+                            TestDescription test) {
+                        testData.addMetric(
+                                test.getTestName(),
+                                Metric.newBuilder()
+                                        .setMeasurements(
+                                                Measurements.newBuilder()
+                                                        .setSingleString("value1")));
+                    }
+                };
+        mBase.init(mContext, mMockListener);
+        mBase.invocationStarted(mContext);
+        mBase.testRunStarted("testRun", 1);
+        TestDescription test = new TestDescription("class", "method");
+        mBase.testStarted(test);
+        mBase.testEnded(test, new HashMap<String, Metric>());
+        mBase.testRunEnded(0L, new HashMap<String, Metric>());
+        mBase.invocationEnded(0L);
+
+        Mockito.verify(mMockListener, times(1)).invocationStarted(Mockito.any());
+        Mockito.verify(mMockListener, times(1)).testRunStarted("testRun", 1);
+        Mockito.verify(mMockListener, times(1)).testStarted(Mockito.eq(test), Mockito.anyLong());
+        Mockito.verify(mMockListener, times(1))
+                .testEnded(Mockito.eq(test), Mockito.anyLong(), mCapturedMetrics.capture());
+
+        Mockito.verify(mMockListener, times(1))
+                .testRunEnded(Mockito.anyLong(), (HashMap<String, Metric>) Mockito.any());
+
+        List<HashMap<String, Metric>> allValues = mCapturedMetrics.getAllValues();
+        assertTrue(allValues.get(0).containsKey("onteststart"));
+        assertTrue(allValues.get(0).containsKey("method"));
+        assertTrue(!allValues.get(0).containsKey("ontestend"));
+    }
 }