Create callbacks for setting up and cleaning up metric collectors.

Subclasses should implement these callbacks to ensure that components
that use metric collectors programatically can have them set up and
cleaned up correctly. These classes should not call these callbacks
directly, however, but through setUp and cleanUp().

Bug: 206859096
Test: atest baseMetricListenerInstrumentedTest (updated)

Original change: https://android-review.googlesource.com/c/platform/platform_testing/+/1918957

Change-Id: I6d74fe767c73199bfc216f15519e6949a05af3fb
(cherry picked from commit 5c8241e95c8ad9e84afd55d9edba9d7d1763dde0)
Merged-In: I6d74fe767c73199bfc216f15519e6949a05af3fb
Merged-In: Ica467d9fb143c8f25e495d4c1d4d4be8632a72ae
Merged-In: Idcb8af0a76e03ea14fe929fd9eff29a23134c942
diff --git a/libraries/device-collectors/src/main/java/android/device/collectors/BaseMetricListener.java b/libraries/device-collectors/src/main/java/android/device/collectors/BaseMetricListener.java
index 48a61e0..fcd5963 100644
--- a/libraries/device-collectors/src/main/java/android/device/collectors/BaseMetricListener.java
+++ b/libraries/device-collectors/src/main/java/android/device/collectors/BaseMetricListener.java
@@ -141,6 +141,7 @@
                 Log.e(getTag(), "Exception during onTestRunEnd.", e);
             }
         }
+        cleanUp();
         super.testRunFinished(result);
     }
 
@@ -214,6 +215,17 @@
     public final void setUp() {
         parseArguments();
         setupAdditionalArgs();
+        onSetUp();
+    }
+
+    /**
+     * Clean up the metric collector.
+     *
+     * <p>If another class is invoking the metric collector's callbacks directly, it should call
+     * this method to make sure that the metric collector is cleaned up properly after collection.
+     */
+    public final void cleanUp() {
+        onCleanUp();
     }
 
     /**
@@ -224,6 +236,17 @@
         return new DataRecord();
     }
 
+    // ---------- Interfaces that can be implemented to set up and clean up metric collection.
+
+    /** Called if custom set-up is needed for this metric collector. */
+    protected void onSetUp() {
+        // Does nothing by default.
+    }
+
+    protected void onCleanUp() {
+        // Does nothing by default.
+    }
+
     // ---------- Interfaces that can be implemented to take action on each test state.
 
     /**
diff --git a/libraries/device-collectors/src/test/java/android/device/collectors/BaseMetricListenerInstrumentedTest.java b/libraries/device-collectors/src/test/java/android/device/collectors/BaseMetricListenerInstrumentedTest.java
index eba8501..5e73d15 100644
--- a/libraries/device-collectors/src/test/java/android/device/collectors/BaseMetricListenerInstrumentedTest.java
+++ b/libraries/device-collectors/src/test/java/android/device/collectors/BaseMetricListenerInstrumentedTest.java
@@ -351,6 +351,101 @@
         assertEquals(0, resultBundle.size());
     }
 
+    @MetricOption(group = "testGroup")
+    @Test
+    public void testSetUpAndCleanUpWithTestCycle() throws Exception {
+        // We use this bundle to mimic device state. We modify it in setUp() and
+        // check that this is maintained throughout the test. We remove the
+        // modification during cleanUp() and check if the bundle has indeed
+        // become empty again.
+        Bundle data = new Bundle();
+        final String arg = "arg";
+        final String value = "value";
+
+        BaseMetricListener listener =
+                new BaseMetricListener() {
+                    public static final String SETUP_ARG = "arg";
+                    private static final String SETUP_VALUE = "value";
+
+                    @Override
+                    protected void onSetUp() {
+                        data.putString(arg, value);
+                    }
+
+                    @Override
+                    protected void onCleanUp() {
+                        data.remove(arg);
+                    }
+
+                    @Override
+                    public void onTestRunStart(DataRecord runData, Description description) {
+                        assertEquals(value, data.getString(arg));
+                    }
+
+                    @Override
+                    public void onTestRunEnd(DataRecord runData, Result result) {
+                        assertEquals(value, data.getString(arg));
+                    }
+
+                    @Override
+                    public void onTestStart(DataRecord testData, Description description) {
+                        assertEquals(value, data.getString(arg));
+                    }
+
+                    @Override
+                    public void onTestEnd(DataRecord testData, Description description) {
+                        assertEquals(value, data.getString(arg));
+                    }
+                };
+
+        Description runDescription = Description.createSuiteDescription("run");
+        Description test1Description = Description.createTestDescription("class", "method1");
+        Description test2Description = Description.createTestDescription("class", "method2");
+        // Simulate a test cycle.
+        listener.testRunStarted(runDescription);
+        listener.testStarted(test1Description);
+        listener.testFinished(test1Description);
+        listener.testStarted(test2Description);
+        listener.testFinished(test2Description);
+        listener.testRunFinished(new Result());
+        listener.instrumentationRunFinished(System.out, new Bundle(), new Result());
+
+        assertFalse(data.containsKey(arg));
+    }
+
+    @MetricOption(group = "testGroup")
+    @Test
+    public void testSetUpAndCleanUpWithCallbacks() throws Exception {
+        // We use this bundle to mimic device state. We modify it in setUp() and
+        // check that this is maintained throughout the test. We remove the
+        // modification during cleanUp() and check if the bundle has indeed
+        // become empty again.
+        Bundle data = new Bundle();
+        final String arg = "arg";
+        final String value = "value";
+
+        BaseMetricListener listener =
+                new BaseMetricListener() {
+                    public static final String SETUP_ARG = "arg";
+                    private static final String SETUP_VALUE = "value";
+
+                    @Override
+                    protected void onSetUp() {
+                        data.putString(arg, value);
+                    }
+
+                    @Override
+                    protected void onCleanUp() {
+                        data.remove(arg);
+                    }
+                };
+
+        listener.setUp();
+        assertEquals(value, data.getString(arg));
+        listener.cleanUp();
+        assertFalse(data.containsKey(arg));
+    }
+
     /**
      * Test annotation that allows to instantiate {@link MetricOption} for testing purpose.
      */