Merge "modify CpuUsageListener to include cpu utilzation." into qt-dev
diff --git a/libraries/collectors-helper/memory/src/com/android/helpers/TotalPssHelper.java b/libraries/collectors-helper/memory/src/com/android/helpers/TotalPssHelper.java
new file mode 100644
index 0000000..4e20c74
--- /dev/null
+++ b/libraries/collectors-helper/memory/src/com/android/helpers/TotalPssHelper.java
@@ -0,0 +1,218 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.helpers;
+
+import static com.android.helpers.MetricUtility.constructKey;
+
+import android.app.ActivityManager;
+import android.app.ActivityManager.RunningAppProcessInfo;
+import android.content.Context;
+import android.os.Debug.MemoryInfo;
+import android.util.Log;
+
+import androidx.test.InstrumentationRegistry;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Helper to collect totalpss memory usage per process tracked by the ActivityManager
+ * memoryinfo.
+ */
+public class TotalPssHelper implements ICollectorHelper<Long> {
+
+    private static final String TAG = TotalPssHelper.class.getSimpleName();
+
+    private static final int DEFAULT_THRESHOLD = 1024;
+    private static final int DEFAULT_MIN_ITERATIONS = 6;
+    private static final int DEFAULT_MAX_ITERATIONS = 20;
+    private static final int DEFAULT_SLEEP_TIME = 1000;
+    private static final String PSS_METRIC_PREFIX = "AM_TOTAL_PSS";
+
+    private String[] mProcessNames;
+    // Minimum number of iterations needed before deciding on the memory usage.
+    private int mMinIterations;
+    // Maximum number of iterations needed waiting for memory usage to be stabilized.
+    private int mMaxIterations;
+    // Sleep time in between the iterations.
+    private int mSleepTime;
+    // Threshold in kb to use whether the data is stabilized.
+    private int mThreshold;
+    // Map to maintain the pss memory size.
+    private Map<String, Long> mPssFinalMap = new HashMap<>();
+
+    public void setUp(String... processNames) {
+        mProcessNames = processNames;
+        // Minimum iterations should be atleast 3 to check for the
+        // stabilization of the memory usage.
+        mMinIterations = DEFAULT_MIN_ITERATIONS;
+        mMaxIterations = DEFAULT_MAX_ITERATIONS;
+        mSleepTime = DEFAULT_SLEEP_TIME;
+        mThreshold = DEFAULT_THRESHOLD;
+    }
+
+    @Override
+    public boolean startCollecting() {
+        return true;
+    }
+
+    @Override
+    public Map<String, Long> getMetrics() {
+        if (mMinIterations < 3) {
+            Log.w(TAG, "Need atleast 3 iterations to check memory usage stabilization.");
+            return mPssFinalMap;
+        }
+        if (mProcessNames != null) {
+            for (String processName : mProcessNames) {
+                if (!processName.isEmpty()) {
+                    measureMemory(processName);
+                }
+            }
+        }
+        return mPssFinalMap;
+    }
+
+    @Override
+    public boolean stopCollecting() {
+        return true;
+    }
+
+    /**
+     * Measure memory info of the given process name tracked by the activity manager
+     * MemoryInfo(i.e getTotalPss).
+     *
+     * @param processName to calculate the memory info.
+     */
+    private void measureMemory(String processName) {
+        Log.i(TAG, "Tracking memory usage of the process - " + processName);
+        List<Long> pssData = new ArrayList<Long>();
+        long pss = 0;
+        int iteration = 0;
+        while (iteration < mMaxIterations) {
+            sleep(mSleepTime);
+            pss = getPss(processName);
+            pssData.add(pss);
+            if (iteration >= mMinIterations && stabilized(pssData)) {
+                Log.i(TAG, "Memory usage stabilized at iteration count = " + iteration);
+                mPssFinalMap.put(constructKey(PSS_METRIC_PREFIX, processName), pss);
+                return;
+            }
+            iteration++;
+        }
+
+        Log.i(TAG, processName + " memory usage did not stabilize."
+                + " Returning the average of the pss data collected.");
+        mPssFinalMap.put(constructKey(PSS_METRIC_PREFIX, processName), average(pssData));
+    }
+
+    /**
+     * Time to sleep in between the iterations.
+     *
+     * @param time in ms to sleep.
+     */
+    private void sleep(int time) {
+        try {
+            Thread.sleep(time);
+        } catch (InterruptedException e) {
+            // ignore
+        }
+    }
+
+    /**
+     * Get the total pss memory of the given process name.
+     *
+     * @param processName of the process to measure the memory.
+     * @return the memory in KB.
+     */
+    private long getPss(String processName) {
+        ActivityManager am = (ActivityManager) InstrumentationRegistry.getInstrumentation()
+                .getContext().getSystemService(Context.ACTIVITY_SERVICE);
+        List<RunningAppProcessInfo> apps = am.getRunningAppProcesses();
+        for (RunningAppProcessInfo proc : apps) {
+            if (!proc.processName.equals(processName)) {
+                continue;
+            }
+            MemoryInfo meminfo = am.getProcessMemoryInfo(new int[] {
+                proc.pid
+            })[0];
+            Log.i(TAG,
+                    String.format("Memory usage of process - %s is %d", processName,
+                            meminfo.getTotalPss()));
+            return meminfo.getTotalPss();
+        }
+        Log.w(TAG, "Not able to find the process id for the process = " + processName);
+        return 0;
+    }
+
+    /**
+     * Checks whether the memory usage is stabilized by calculating the sum of the difference
+     * between the last 3 values and comparing that to the threshold.
+     *
+     * @param pssData list of pssData of the given process name.
+     * @return true if the memory is stabilized.
+     */
+    private boolean stabilized(List<Long> pssData) {
+        long diff1 = Math.abs(pssData.get(pssData.size() - 1) - pssData.get(pssData.size() - 2));
+        long diff2 = Math.abs(pssData.get(pssData.size() - 2) - pssData.get(pssData.size() - 3));
+        Log.i(TAG, "diff1=" + diff1 + " diff2=" + diff2);
+        return (diff1 + diff2) < mThreshold;
+    }
+
+    /**
+     * Returns the average of the pssData collected for the maxIterations.
+     *
+     * @param pssData list of pssData.
+     * @return
+     */
+    private long average(List<Long> pssData) {
+        long sum = 0;
+        for (long sample : pssData) {
+            sum += sample;
+        }
+        return sum / pssData.size();
+    }
+
+    /**
+     * @param minIterations before starting to check for memory is stabilized.
+     */
+    public void setMinIterations(int minIterations) {
+        mMinIterations = minIterations;
+    }
+
+    /**
+     * @param maxIterations to wait for memory to be stabilized.
+     */
+    public void setMaxIterations(int maxIterations) {
+        mMaxIterations = maxIterations;
+    }
+
+    /**
+     * @param sleepTime in between the iterations.
+     */
+    public void setSleepTime(int sleepTime) {
+        mSleepTime = sleepTime;
+    }
+
+    /**
+     * @param threshold for difference in memory usage between two successive iterations in kb
+     */
+    public void setThreshold(int threshold) {
+        mThreshold = threshold;
+    }
+}
diff --git a/libraries/collectors-helper/memory/test/src/com/android/helpers/tests/TotalPssHelperTest.java b/libraries/collectors-helper/memory/test/src/com/android/helpers/tests/TotalPssHelperTest.java
new file mode 100644
index 0000000..b3dd138
--- /dev/null
+++ b/libraries/collectors-helper/memory/test/src/com/android/helpers/tests/TotalPssHelperTest.java
@@ -0,0 +1,105 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.helpers.tests;
+
+import static com.android.helpers.MetricUtility.constructKey;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.helpers.TotalPssHelper;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.Map;
+
+/**
+ * Android Unit tests for {@link TotalPssHelper}.
+ *
+ * To run:
+ * atest CollectorsHelperTest:com.android.helpers.tests.TotalPssHelperTest
+ */
+@RunWith(AndroidJUnit4.class)
+public class TotalPssHelperTest {
+
+    // Process name used for testing
+    private static final String TEST_PROCESS_NAME = "com.android.systemui";
+    // Second process name used for testing
+    private static final String TEST_PROCESS_NAME_2 = "com.google.android.apps.nexuslauncher";
+    // Second process name used for testing
+    private static final String INVALID_PROCESS_NAME = "abc";
+    // Pss prefix in Key.
+    private static final String PSS_METRIC_PREFIX = "AM_TOTAL_PSS";
+
+    private TotalPssHelper mTotalPssHelper;
+
+    @Before
+    public void setUp() {
+         mTotalPssHelper = new TotalPssHelper();
+    }
+
+    /** Test no metrics are sampled if process name is empty. */
+    @Test
+    public void testEmptyProcessName() {
+        mTotalPssHelper.setUp("");
+        Map<String, Long> pssMetrics = mTotalPssHelper.getMetrics();
+        assertTrue(pssMetrics.isEmpty());
+    }
+
+    /** Test no metrics are sampled if process names is null */
+    @Test
+    public void testNullProcessName() {
+        mTotalPssHelper.setUp(null);
+        Map<String, Long> pssMetrics = mTotalPssHelper.getMetrics();
+        assertTrue(pssMetrics.isEmpty());
+    }
+
+    /** Test getting metrics for single process. */
+    @Test
+    public void testGetMetrics_OneProcess() {
+        mTotalPssHelper.setUp(TEST_PROCESS_NAME);
+        Map<String, Long> pssMetrics = mTotalPssHelper.getMetrics();
+        assertFalse(pssMetrics.isEmpty());
+        assertTrue(pssMetrics.containsKey(constructKey(PSS_METRIC_PREFIX, TEST_PROCESS_NAME)));
+        assertTrue(pssMetrics.get(constructKey(PSS_METRIC_PREFIX, TEST_PROCESS_NAME)) > 0);
+    }
+
+    /** Test getting metrics for multiple process. */
+    @Test
+    public void testGetMetrics_MultipleProcesses() {
+        mTotalPssHelper.setUp(TEST_PROCESS_NAME, TEST_PROCESS_NAME_2);
+        Map<String, Long> pssMetrics = mTotalPssHelper.getMetrics();
+        assertFalse(pssMetrics.isEmpty());
+        assertTrue(pssMetrics.containsKey(constructKey(PSS_METRIC_PREFIX, TEST_PROCESS_NAME)));
+        assertTrue(pssMetrics.containsKey(constructKey(PSS_METRIC_PREFIX, TEST_PROCESS_NAME_2)));
+        assertTrue(pssMetrics.get(constructKey(PSS_METRIC_PREFIX, TEST_PROCESS_NAME)) > 0);
+        assertTrue(pssMetrics.get(constructKey(PSS_METRIC_PREFIX, TEST_PROCESS_NAME_2)) > 0);
+    }
+
+    /** Test pss metric is 0 for invalid process name. */
+    @Test
+    public void testGetMetrics_InvalidProcess() {
+        mTotalPssHelper.setUp(INVALID_PROCESS_NAME);
+        Map<String, Long> pssMetrics = mTotalPssHelper.getMetrics();
+        assertTrue(pssMetrics.containsKey(constructKey(PSS_METRIC_PREFIX, INVALID_PROCESS_NAME)));
+        assertTrue(pssMetrics.get(constructKey(PSS_METRIC_PREFIX, INVALID_PROCESS_NAME)) == 0);
+    }
+}
diff --git a/libraries/device-collectors/src/main/java/android/device/collectors/BaseCollectionListener.java b/libraries/device-collectors/src/main/java/android/device/collectors/BaseCollectionListener.java
index fc45c1e..8003aed 100644
--- a/libraries/device-collectors/src/main/java/android/device/collectors/BaseCollectionListener.java
+++ b/libraries/device-collectors/src/main/java/android/device/collectors/BaseCollectionListener.java
@@ -16,12 +16,14 @@
 package android.device.collectors;
 
 import android.os.Bundle;
+import android.util.Log;
 
 import androidx.annotation.VisibleForTesting;
 
 import com.android.helpers.ICollectorHelper;
 
 import org.junit.runner.Description;
+import org.junit.runner.notification.Failure;
 import org.junit.runner.Result;
 
 import java.util.Map;
@@ -42,7 +44,11 @@
     protected ICollectorHelper mHelper;
     // Collect per run if it is set to true otherwise collect per test.
     public static final String COLLECT_PER_RUN = "per_run";
+    // Skip failure metrics collection if this flag is set to true.
+    public static final String SKIP_TEST_FAILURE_METRICS = "skip_test_failure_metrics";
     protected boolean mIsCollectPerRun;
+    protected boolean mSkipTestFailureMetrics;
+    private boolean mIsTestFailed = false;
 
     public BaseCollectionListener() {
         super();
@@ -58,6 +64,8 @@
     public void onTestRunStart(DataRecord runData, Description description) {
         Bundle args = getArgsBundle();
         mIsCollectPerRun = "true".equals(args.getString(COLLECT_PER_RUN));
+        // By default this flag is set to false to collect the metrics on test failure.
+        mSkipTestFailureMetrics = "true".equals(args.getString(SKIP_TEST_FAILURE_METRICS));
 
         // Setup additional args before starting the collection.
         setupAdditionalArgs();
@@ -69,18 +77,32 @@
     }
 
     @Override
-    public void onTestStart(DataRecord testData, Description description) {
+    public final void onTestStart(DataRecord testData, Description description) {
+        mIsTestFailed = false;
         if (!mIsCollectPerRun) {
             mHelper.startCollecting();
         }
     }
 
     @Override
-    public void onTestEnd(DataRecord testData, Description description) {
+    public void onTestFail(DataRecord testData, Description description, Failure failure) {
+        mIsTestFailed = true;
+    }
+
+    @Override
+    public final void onTestEnd(DataRecord testData, Description description) {
         if (!mIsCollectPerRun) {
-            Map<String, T> metrics = mHelper.getMetrics();
-            for (Map.Entry<String, T> entry : metrics.entrySet()) {
-                testData.addStringMetric(entry.getKey(), entry.getValue().toString());
+            // Skip adding the metrics collected during the test failure
+            // if the skip metrics on test failure flag is enabled and the
+            // current test is failed.
+            if (mSkipTestFailureMetrics && mIsTestFailed) {
+                Log.i(getTag(), "Skipping the metric collection.");
+            } else {
+                // Collect the metrics.
+                Map<String, T> metrics = mHelper.getMetrics();
+                for (Map.Entry<String, T> entry : metrics.entrySet()) {
+                    testData.addStringMetric(entry.getKey(), entry.getValue().toString());
+                }
             }
             mHelper.stopCollecting();
         }
@@ -98,15 +120,14 @@
     }
 
     /**
-     * To add listener specific extra args implement this method in the sub class
-     * and add the listener specific args.
+     * To add listener specific extra args implement this method in the sub class and add the
+     * listener specific args.
      */
     public void setupAdditionalArgs() {
-       // NO-OP by default
+        // NO-OP by default
     }
 
     protected void createHelperInstance(ICollectorHelper helper) {
         mHelper = helper;
     }
-
 }
diff --git a/libraries/device-collectors/src/main/java/android/device/collectors/TotalPssMetricListener.java b/libraries/device-collectors/src/main/java/android/device/collectors/TotalPssMetricListener.java
new file mode 100644
index 0000000..bb5a8a9
--- /dev/null
+++ b/libraries/device-collectors/src/main/java/android/device/collectors/TotalPssMetricListener.java
@@ -0,0 +1,93 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.device.collectors;
+
+import android.device.collectors.annotations.OptionClass;
+import android.os.Bundle;
+import android.util.Log;
+
+import androidx.annotation.VisibleForTesting;
+
+import com.android.helpers.TotalPssHelper;
+
+/**
+ * A {@link TotalPssMetricListener} that measures process total pss tracked per
+ * process in activity manaager.
+ *
+ * Options:
+ * -e process-names [processName] : the process from the test case that we want to
+ * measure memory for.
+ */
+@OptionClass(alias = "totalpss-collector")
+public class TotalPssMetricListener extends BaseCollectionListener<Long> {
+
+    private static final String TAG = TotalPssMetricListener.class.getSimpleName();
+    @VisibleForTesting static final String PROCESS_SEPARATOR = ",";
+    @VisibleForTesting static final String PROCESS_NAMES_KEY = "process-names";
+    @VisibleForTesting static final String MIN_ITERATIONS_KEY = "min_iterations";
+    @VisibleForTesting static final String MAX_ITERATIONS_KEY = "max_iterations";
+    @VisibleForTesting static final String SLEEP_TIME_KEY = "sleep_time_ms";
+    @VisibleForTesting static final String THRESHOLD_KEY = "threshold_kb";
+    private TotalPssHelper mTotalPssHelper = new TotalPssHelper();
+
+    public TotalPssMetricListener() {
+        createHelperInstance(mTotalPssHelper);
+    }
+
+    /**
+     * Constructor to simulate receiving the instrumentation arguments. Should not be used except
+     * for testing.
+     */
+    @VisibleForTesting
+    public TotalPssMetricListener(Bundle args, TotalPssHelper helper) {
+        super(args, helper);
+        mTotalPssHelper = helper;
+        createHelperInstance(mTotalPssHelper);
+    }
+
+    /**
+     * Adds the options for total pss collector.
+     */
+    @Override
+    public void setupAdditionalArgs() {
+        Bundle args = getArgsBundle();
+        String procsString = args.getString(PROCESS_NAMES_KEY);
+        if (procsString == null) {
+            Log.e(TAG, "No processes provided to sample");
+            return;
+        }
+
+        String[] procs = procsString.split(PROCESS_SEPARATOR);
+        mTotalPssHelper.setUp(procs);
+
+        if (args.getString(MIN_ITERATIONS_KEY) != null) {
+            mTotalPssHelper.setMinIterations(Integer.parseInt(args.getString(MIN_ITERATIONS_KEY)));
+        }
+
+        if (args.getString(MAX_ITERATIONS_KEY) != null) {
+            mTotalPssHelper.setMaxIterations(Integer.parseInt(args.getString(MAX_ITERATIONS_KEY)));
+        }
+
+        if (args.getString(SLEEP_TIME_KEY) != null) {
+            mTotalPssHelper.setSleepTime(Integer.parseInt(args.getString(SLEEP_TIME_KEY)));
+        }
+
+        if (args.getString(THRESHOLD_KEY) != null) {
+            mTotalPssHelper.setThreshold(Integer.parseInt(args.getString(THRESHOLD_KEY)));
+        }
+    }
+}
diff --git a/libraries/device-collectors/src/test/java/android/device/collectors/BaseCollectionListenerTest.java b/libraries/device-collectors/src/test/java/android/device/collectors/BaseCollectionListenerTest.java
index ff7f8f8..9e23918 100644
--- a/libraries/device-collectors/src/test/java/android/device/collectors/BaseCollectionListenerTest.java
+++ b/libraries/device-collectors/src/test/java/android/device/collectors/BaseCollectionListenerTest.java
@@ -25,6 +25,7 @@
 import org.junit.runner.Description;
 import org.junit.runner.Result;
 import org.junit.runner.RunWith;
+import org.junit.runner.notification.Failure;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -65,8 +66,8 @@
     }
 
     /**
-     * Verify start and stop collection happens only during test run started
-     * and test run ended when per_run option is enabled.
+     * Verify start and stop collection happens only during test run started and test run ended when
+     * per_run option is enabled.
      */
     @Test
     public void testPerRunFlow() throws Exception {
@@ -85,9 +86,8 @@
     }
 
     /**
-     * Verify start and stop collection happens before and after each test
-     * and not during test run started and test run ended when per_run option is
-     * disabled.
+     * Verify start and stop collection happens before and after each test and not during test run
+     * started and test run ended when per_run option is disabled.
      */
     @Test
     public void testPerTestFlow() throws Exception {
@@ -100,18 +100,20 @@
         mListener.onTestStart(mListener.createDataRecord(), FAKE_DESCRIPTION);
         verify(helper, times(1)).startCollecting();
         mListener.onTestEnd(mListener.createDataRecord(), FAKE_DESCRIPTION);
+        verify(helper, times(1)).getMetrics();
         verify(helper, times(1)).stopCollecting();
         mListener.onTestStart(mListener.createDataRecord(), FAKE_DESCRIPTION);
         verify(helper, times(2)).startCollecting();
         mListener.onTestEnd(mListener.createDataRecord(), FAKE_DESCRIPTION);
         verify(helper, times(2)).stopCollecting();
+        verify(helper, times(2)).getMetrics();
         mListener.onTestRunEnd(mListener.createDataRecord(), new Result());
         verify(helper, times(2)).stopCollecting();
     }
 
     /**
-     * Verify start and stop collection happens before and after each test
-     * and not during test run started and test run ended by default.
+     * Verify start and stop collection happens before and after each test and not during test run
+     * started and test run ended by default.
      */
     @Test
     public void testDefaultOptionFlow() throws Exception {
@@ -131,4 +133,109 @@
         mListener.onTestRunEnd(mListener.createDataRecord(), new Result());
         verify(helper, times(2)).stopCollecting();
     }
+
+    /**
+     * Verify metrics is collected when skip on test failure is explictly set
+     * to false.
+     */
+    @Test
+    public void testPerTestFailureFlowNotCollectMetrics() throws Exception {
+        Bundle b = new Bundle();
+        b.putString(BaseCollectionListener.COLLECT_PER_RUN, "false");
+        b.putString(BaseCollectionListener.SKIP_TEST_FAILURE_METRICS, "false");
+        mListener = initListener(b);
+
+        mListener.onTestRunStart(mListener.createDataRecord(), FAKE_DESCRIPTION);
+        verify(helper, times(0)).startCollecting();
+        mListener.onTestStart(mListener.createDataRecord(), FAKE_DESCRIPTION);
+        verify(helper, times(1)).startCollecting();
+        Failure failureDesc = new Failure(Description.createSuiteDescription("run"),
+                new Exception());
+        mListener.testFailure(failureDesc);
+        mListener.onTestEnd(mListener.createDataRecord(), FAKE_DESCRIPTION);
+        verify(helper, times(1)).getMetrics();
+        verify(helper, times(1)).stopCollecting();
+    }
+
+    /**
+     * Verify default behaviour to collect the metrics on test failure.
+     */
+    @Test
+    public void testPerTestFailureFlowDefault() throws Exception {
+        Bundle b = new Bundle();
+        b.putString(BaseCollectionListener.COLLECT_PER_RUN, "false");
+        mListener = initListener(b);
+
+        mListener.onTestRunStart(mListener.createDataRecord(), FAKE_DESCRIPTION);
+        verify(helper, times(0)).startCollecting();
+        mListener.onTestStart(mListener.createDataRecord(), FAKE_DESCRIPTION);
+        verify(helper, times(1)).startCollecting();
+        Failure failureDesc = new Failure(Description.createSuiteDescription("run"),
+                new Exception());
+        mListener.testFailure(failureDesc);
+        mListener.onTestEnd(mListener.createDataRecord(), FAKE_DESCRIPTION);
+        // Metrics should be called by default on test failure by default.
+        verify(helper, times(1)).getMetrics();
+        verify(helper, times(1)).stopCollecting();
+    }
+
+    /**
+     * Verify metrics collection is skipped if the skip on failure metrics
+     * is enabled and if the test is failed.
+     */
+    @Test
+    public void testPerTestFailureSkipMetrics() throws Exception {
+        Bundle b = new Bundle();
+        b.putString(BaseCollectionListener.COLLECT_PER_RUN, "false");
+        b.putString(BaseCollectionListener.SKIP_TEST_FAILURE_METRICS, "true");
+        mListener = initListener(b);
+
+        mListener.onTestRunStart(mListener.createDataRecord(), FAKE_DESCRIPTION);
+        verify(helper, times(0)).startCollecting();
+        mListener.onTestStart(mListener.createDataRecord(), FAKE_DESCRIPTION);
+        verify(helper, times(1)).startCollecting();
+        Failure failureDesc = new Failure(Description.createSuiteDescription("run"),
+                new Exception());
+        mListener.testFailure(failureDesc);
+        mListener.onTestEnd(mListener.createDataRecord(), FAKE_DESCRIPTION);
+        // Metrics should not be collected.
+        verify(helper, times(0)).getMetrics();
+        verify(helper, times(1)).stopCollecting();
+    }
+
+    /**
+     * Verify metrics not collected for test failure in between two test that
+     * succeeded when skip metrics on test failure is enabled.
+     */
+    @Test
+    public void testInterleavingTestFailureMetricsSkip() throws Exception {
+        Bundle b = new Bundle();
+        b.putString(BaseCollectionListener.COLLECT_PER_RUN, "false");
+        b.putString(BaseCollectionListener.SKIP_TEST_FAILURE_METRICS, "true");
+        mListener = initListener(b);
+
+        mListener.onTestRunStart(mListener.createDataRecord(), FAKE_DESCRIPTION);
+        verify(helper, times(0)).startCollecting();
+        mListener.testStarted(FAKE_DESCRIPTION);
+        verify(helper, times(1)).startCollecting();
+        mListener.onTestEnd(mListener.createDataRecord(), FAKE_DESCRIPTION);
+        verify(helper, times(1)).getMetrics();
+        verify(helper, times(1)).stopCollecting();
+
+        mListener.testStarted(FAKE_DESCRIPTION);
+        verify(helper, times(2)).startCollecting();
+        Failure failureDesc = new Failure(Description.createSuiteDescription("run"),
+                new Exception());
+        mListener.testFailure(failureDesc);
+        mListener.onTestEnd(mListener.createDataRecord(), FAKE_DESCRIPTION);
+        // Metric collection should not be done on failure.
+        verify(helper, times(1)).getMetrics();
+        verify(helper, times(2)).stopCollecting();
+
+        mListener.testStarted(FAKE_DESCRIPTION);
+        verify(helper, times(3)).startCollecting();
+        mListener.onTestEnd(mListener.createDataRecord(), FAKE_DESCRIPTION);
+        verify(helper, times(2)).getMetrics();
+        verify(helper, times(3)).stopCollecting();
+    }
 }
diff --git a/libraries/device-collectors/src/test/java/android/device/collectors/TotalPssMetricListenerTest.java b/libraries/device-collectors/src/test/java/android/device/collectors/TotalPssMetricListenerTest.java
new file mode 100644
index 0000000..6653970
--- /dev/null
+++ b/libraries/device-collectors/src/test/java/android/device/collectors/TotalPssMetricListenerTest.java
@@ -0,0 +1,100 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.device.collectors;
+
+import static android.device.collectors.TotalPssMetricListener.PROCESS_NAMES_KEY;
+import static android.device.collectors.TotalPssMetricListener.PROCESS_SEPARATOR;
+import static android.device.collectors.TotalPssMetricListener.MIN_ITERATIONS_KEY;
+import static android.device.collectors.TotalPssMetricListener.MAX_ITERATIONS_KEY;
+import static android.device.collectors.TotalPssMetricListener.SLEEP_TIME_KEY;
+import static android.device.collectors.TotalPssMetricListener.THRESHOLD_KEY;
+
+import static org.mockito.Mockito.verify;
+
+import android.app.Instrumentation;
+import android.os.Bundle;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.helpers.TotalPssHelper;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.Description;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Android Unit tests for {@link TotalPssMetricListener}.
+ *
+ * To run:
+ * atest CollectorDeviceLibTest:android.device.collectors.TotalPssMetricListenerTest
+ */
+@RunWith(AndroidJUnit4.class)
+public class TotalPssMetricListenerTest {
+
+    @Mock
+    private Instrumentation mInstrumentation;
+    @Mock
+    private TotalPssHelper mTotalPssMetricHelper;
+
+    private TotalPssMetricListener mListener;
+    private Description mRunDesc;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mRunDesc = Description.createSuiteDescription("run");
+    }
+
+    private TotalPssMetricListener initListener(Bundle b) {
+        TotalPssMetricListener listener = new TotalPssMetricListener(b, mTotalPssMetricHelper);
+        listener.setInstrumentation(mInstrumentation);
+        return listener;
+    }
+
+    @Test
+    public void testHelperReceivesProcessNames() throws Exception {
+        Bundle b = new Bundle();
+        b.putString(PROCESS_NAMES_KEY, "process1" + PROCESS_SEPARATOR + "process2");
+        mListener = initListener(b);
+
+        mListener.testRunStarted(mRunDesc);
+
+        verify(mTotalPssMetricHelper).setUp("process1", "process2");
+    }
+
+    @Test
+    public void testAdditionalPssOptions() throws Exception {
+        Bundle b = new Bundle();
+        b.putString(PROCESS_NAMES_KEY, "process1");
+        b.putString(MIN_ITERATIONS_KEY, "50");
+        b.putString(MAX_ITERATIONS_KEY, "102");
+        b.putString(SLEEP_TIME_KEY, "2000");
+        b.putString(THRESHOLD_KEY, "2048");
+        mListener = initListener(b);
+
+        mListener.testRunStarted(mRunDesc);
+
+        verify(mTotalPssMetricHelper).setUp("process1");
+        verify(mTotalPssMetricHelper).setMinIterations(50);
+        verify(mTotalPssMetricHelper).setMaxIterations(102);
+        verify(mTotalPssMetricHelper).setSleepTime(2000);
+        verify(mTotalPssMetricHelper).setThreshold(2048);
+    }
+}
diff --git a/libraries/longevity/README.md b/libraries/longevity/README.md
index 0554aea..1c394d6 100644
--- a/libraries/longevity/README.md
+++ b/libraries/longevity/README.md
@@ -68,6 +68,8 @@
 *   `timeout_msec <long>` - a timeout for individual test methods.
 *   `quitter <bool>` - quit the suite if any test errors are encountered.
 *   `profile <string>` - use a profile under assets/ or at your own path.
+*   `rename-iterations <bool>` - rename each iteration by appending the iteration number to the
+    class name.
 
 ## Tests
 
diff --git a/libraries/longevity/platform/samples/src/android/platform/test/longevity/samples/SimpleSuite.java b/libraries/longevity/platform/samples/src/android/platform/test/longevity/samples/SimpleSuite.java
index dc5ae3c..9187c7e 100644
--- a/libraries/longevity/platform/samples/src/android/platform/test/longevity/samples/SimpleSuite.java
+++ b/libraries/longevity/platform/samples/src/android/platform/test/longevity/samples/SimpleSuite.java
@@ -28,11 +28,9 @@
 @RunWith(LongevitySuite.class)
 @SuiteClasses({
     SimpleSuite.PassingTest.class,
-    SimpleSuite.FailingTest.class
+    SimpleSuite.FailingTest.class,
 })
-/**
- * Sample device-side test cases.
- */
+/** Sample device-side test cases. */
 public class SimpleSuite {
     // no local test cases.
 
diff --git a/libraries/longevity/platform/src/android/platform/test/longevity/LongevityClassRunner.java b/libraries/longevity/platform/src/android/platform/test/longevity/LongevityClassRunner.java
index 0097d7b..34382b2 100644
--- a/libraries/longevity/platform/src/android/platform/test/longevity/LongevityClassRunner.java
+++ b/libraries/longevity/platform/src/android/platform/test/longevity/LongevityClassRunner.java
@@ -32,6 +32,7 @@
 import org.junit.BeforeClass;
 import org.junit.internal.runners.statements.RunAfters;
 import org.junit.internal.runners.statements.RunBefores;
+import org.junit.runner.Description;
 import org.junit.runners.BlockJUnit4ClassRunner;
 import org.junit.runners.model.FrameworkMethod;
 import org.junit.runners.model.InitializationError;
@@ -45,11 +46,16 @@
  */
 public class LongevityClassRunner extends BlockJUnit4ClassRunner {
     @VisibleForTesting static final String FILTER_OPTION = "exclude-class";
+    @VisibleForTesting static final String ITERATION_SEP = "@";
+    // A constant to indicate that the iteration number is not set.
+    @VisibleForTesting static final int ITERATION_NOT_SET = -1;
 
     private String[] mExcludedClasses;
 
     private boolean mTestFailed = true;
     private boolean mTestAttempted = false;
+    // Iteration number.
+    private int mIteration = ITERATION_NOT_SET;
 
     public LongevityClassRunner(Class<?> klass) throws InitializationError {
         this(klass, InstrumentationRegistry.getArguments());
@@ -64,6 +70,19 @@
                         : new String[] {};
     }
 
+    /** Set the iteration of the test that this runner is running. */
+    public void setIteration(int iteration) {
+        mIteration = iteration;
+    }
+
+    /**
+     * Utilized by tests to check that the iteration is set, independent of the description logic.
+     */
+    @VisibleForTesting
+    int getIteration() {
+        return mIteration;
+    }
+
     /**
      * Override the parent {@code withBeforeClasses} method to be a no-op.
      *
@@ -231,4 +250,25 @@
         }
         return errors;
     }
+
+    /**
+     * Rename the child class name to add iterations if the renaming iteration option is enabled.
+     *
+     * <p>Renaming the class here is chosen over renaming the method name because
+     *
+     * <ul>
+     *   <li>Conceptually, the runner is running a class multiple times, as opposed to a method.
+     *   <li>When instrumenting a suite in command line, by default the instrumentation command
+     *       outputs the class name only. Renaming the class helps with interpretation in this case.
+     */
+    @Override
+    protected Description describeChild(FrameworkMethod method) {
+        Description original = super.describeChild(method);
+        if (mIteration == ITERATION_NOT_SET) {
+            return original;
+        }
+        return Description.createTestDescription(
+                String.join(ITERATION_SEP, original.getClassName(), String.valueOf(mIteration)),
+                original.getMethodName());
+    }
 }
diff --git a/libraries/longevity/platform/src/android/platform/test/longevity/LongevitySuite.java b/libraries/longevity/platform/src/android/platform/test/longevity/LongevitySuite.java
index 76ada9a..850a830 100644
--- a/libraries/longevity/platform/src/android/platform/test/longevity/LongevitySuite.java
+++ b/libraries/longevity/platform/src/android/platform/test/longevity/LongevitySuite.java
@@ -35,6 +35,7 @@
 import java.util.List;
 import java.util.Map;
 
+import org.junit.runner.Description;
 import org.junit.runner.Runner;
 import org.junit.runner.notification.RunNotifier;
 import org.junit.runners.BlockJUnit4ClassRunner;
@@ -49,12 +50,17 @@
 public class LongevitySuite extends android.host.test.longevity.LongevitySuite {
     private static final String LOG_TAG = LongevitySuite.class.getSimpleName();
 
+    public static final String RENAME_ITERATION_OPTION = "rename-iterations";
+    private boolean mRenameIterations;
+
     private Instrumentation mInstrumentation;
     private Context mContext;
 
     // Cached {@link TimeoutTerminator} instance.
     private TimeoutTerminator mTimeoutTerminator;
 
+    private Map<Description, Integer> mIterations = new HashMap<>();
+
     /**
      * Takes a {@link Bundle} and maps all String K/V pairs into a {@link Map<String, String>}.
      *
@@ -92,6 +98,9 @@
         super(klass, runners, toMap(args));
         mInstrumentation = InstrumentationRegistry.getInstrumentation();
         mContext = InstrumentationRegistry.getContext();
+
+        // Parse out additional options.
+        mRenameIterations = Boolean.valueOf(args.getString(RENAME_ITERATION_OPTION));
     }
 
     /**
@@ -155,7 +164,15 @@
 
     @Override
     protected void runChild(Runner runner, final RunNotifier notifier) {
-        super.runChild(getSuiteRunner(runner), notifier);
+        // Update iterations.
+        mIterations.computeIfPresent(runner.getDescription(), (k, v) -> v + 1);
+        mIterations.computeIfAbsent(runner.getDescription(), k -> 1);
+
+        LongevityClassRunner suiteRunner = getSuiteRunner(runner);
+        if (mRenameIterations) {
+            suiteRunner.setIteration(mIterations.get(runner.getDescription()));
+        }
+        super.runChild(suiteRunner, notifier);
     }
 
     /**
@@ -193,7 +210,7 @@
      * Returns a {@link Runner} specific for the suite, if any. Can be overriden by subclasses to
      * supply different runner implementations.
      */
-    protected Runner getSuiteRunner(Runner runner) {
+    protected LongevityClassRunner getSuiteRunner(Runner runner) {
         try {
             // Cast is safe as we verified the runner is BlockJUnit4Runner at initialization.
             return new LongevityClassRunner(
diff --git a/libraries/longevity/platform/src/android/platform/test/longevity/ProfileSuite.java b/libraries/longevity/platform/src/android/platform/test/longevity/ProfileSuite.java
index 76746bb..3e6bee1 100644
--- a/libraries/longevity/platform/src/android/platform/test/longevity/ProfileSuite.java
+++ b/libraries/longevity/platform/src/android/platform/test/longevity/ProfileSuite.java
@@ -136,7 +136,7 @@
      * features expand.
      */
     @Override
-    protected Runner getSuiteRunner(Runner runner) {
+    protected LongevityClassRunner getSuiteRunner(Runner runner) {
         if (mProfile.getConfiguration() == null) {
             return super.getSuiteRunner(runner);
         }
diff --git a/libraries/longevity/platform/tests/AndroidTest.xml b/libraries/longevity/platform/tests/AndroidTest.xml
index 68acaf4..530392c 100644
--- a/libraries/longevity/platform/tests/AndroidTest.xml
+++ b/libraries/longevity/platform/tests/AndroidTest.xml
@@ -24,6 +24,7 @@
     <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
         <option name="package" value="android.platform.test.longevity.tests" />
         <option name="exclude-filter" value="android.platform.test.longevity.samples.testing.SampleProfileSuite" />
+        <option name="exclude-filter" value="android.platform.test.scenario" />
         <option name="runtime-hint" value="1m30s" />
     </test>
 </configuration>
diff --git a/libraries/longevity/platform/tests/src/android/platform/test/longevity/LongevityClassRunnerTest.java b/libraries/longevity/platform/tests/src/android/platform/test/longevity/LongevityClassRunnerTest.java
index 9db4e26..70669b4 100644
--- a/libraries/longevity/platform/tests/src/android/platform/test/longevity/LongevityClassRunnerTest.java
+++ b/libraries/longevity/platform/tests/src/android/platform/test/longevity/LongevityClassRunnerTest.java
@@ -21,6 +21,7 @@
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
@@ -37,6 +38,7 @@
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
+import org.junit.runner.Description;
 import org.junit.runner.RunWith;
 import org.junit.runner.notification.Failure;
 import org.junit.runner.notification.RunListener;
@@ -46,7 +48,6 @@
 import org.junit.runners.model.Statement;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
-import org.mockito.Mockito;
 import org.mockito.exceptions.base.MockitoAssertionError;
 
 
@@ -359,8 +360,8 @@
     /** Test that excluded classes are not executed. */
     @Test
     public void testIgnore_excludedClasses() throws Throwable {
-        RunNotifier notifier = Mockito.spy(new RunNotifier());
-        RunListener listener = Mockito.mock(RunListener.class);
+        RunNotifier notifier = spy(new RunNotifier());
+        RunListener listener = mock(RunListener.class);
         notifier.addListener(listener);
         Bundle ignores = new Bundle();
         ignores.putString(LongevityClassRunner.FILTER_OPTION, FailingTest.class.getCanonicalName());
@@ -369,6 +370,38 @@
         verify(listener, times(1)).testIgnored(any());
     }
 
+    /** Test that the runner does not report iteration when iteration is not set. */
+    @Test
+    public void testReportIteration_noIterationSet() throws Throwable {
+        ArgumentCaptor<Description> captor = ArgumentCaptor.forClass(Description.class);
+        RunNotifier notifier = mock(RunNotifier.class);
+        mRunner = spy(new LongevityClassRunner(NoOpTest.class));
+        mRunner.run(notifier);
+        verify(notifier).fireTestStarted(captor.capture());
+        Assert.assertFalse(
+                "Description class name should not contain the iteration number.",
+                captor.getValue()
+                        .getClassName()
+                        .matches(
+                                String.join(LongevityClassRunner.ITERATION_SEP, "^.*", "[0-9]+$")));
+    }
+
+    /** Test that the runner reports iteration when set. */
+    @Test
+    public void testReportIteration_withIteration() throws Throwable {
+        ArgumentCaptor<Description> captor = ArgumentCaptor.forClass(Description.class);
+        RunNotifier notifier = mock(RunNotifier.class);
+        mRunner = spy(new LongevityClassRunner(NoOpTest.class));
+        mRunner.setIteration(7);
+        mRunner.run(notifier);
+        verify(notifier).fireTestStarted(captor.capture());
+        Assert.assertTrue(
+                "Description class name should contain the iteration number.",
+                captor.getValue()
+                        .getClassName()
+                        .matches(String.join(LongevityClassRunner.ITERATION_SEP, "^.*", "7$")));
+    }
+
     private List<FrameworkMethod> getMethodNameMatcher(String methodName) {
         return argThat(
                 l ->
diff --git a/libraries/longevity/platform/tests/src/android/platform/test/longevity/LongevitySuiteTest.java b/libraries/longevity/platform/tests/src/android/platform/test/longevity/LongevitySuiteTest.java
index 65d7041..d50d536 100644
--- a/libraries/longevity/platform/tests/src/android/platform/test/longevity/LongevitySuiteTest.java
+++ b/libraries/longevity/platform/tests/src/android/platform/test/longevity/LongevitySuiteTest.java
@@ -16,6 +16,7 @@
 package android.platform.test.longevity;
 
 import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -26,9 +27,11 @@
 import android.os.BatteryManager;
 import android.os.Bundle;
 
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.internal.builders.AllDefaultPossibilitiesBuilder;
+import org.junit.runner.Runner;
 import org.junit.runner.RunWith;
 import org.junit.runner.notification.RunNotifier;
 import org.junit.runners.JUnit4;
@@ -36,6 +39,9 @@
 import org.junit.runners.Suite.SuiteClasses;
 import org.mockito.Mockito;
 
+import java.util.ArrayList;
+import java.util.List;
+
 /**
  * Unit tests for the {@link LongevitySuite} runner.
  */
@@ -50,19 +56,22 @@
     private LongevitySuite mSuite;
 
     @Before
-    public void setUpSuite() throws InitializationError {
+    public void setUp() {
         // Android context mocking.
         when(mContext.registerReceiver(any(), any())).thenReturn(mBatteryIntent);
-        // Create the core suite to test.
-        mSuite = new LongevitySuite(TestSuite.class, new AllDefaultPossibilitiesBuilder(true),
-                mInstrumentation, mContext, new Bundle());
     }
 
-    /**
-     * Tests that devices with batteries terminate on low battery.
-     */
+    /** Tests that devices with batteries terminate on low battery. */
     @Test
-    public void testDeviceWithBattery_registersReceiver() {
+    public void testDeviceWithBattery_registersReceiver() throws InitializationError {
+        // Create the core suite to test.
+        mSuite =
+                new LongevitySuite(
+                        TestSuite.class,
+                        new AllDefaultPossibilitiesBuilder(true),
+                        mInstrumentation,
+                        mContext,
+                        new Bundle());
         mBatteryIntent.putExtra(BatteryManager.EXTRA_PRESENT, true);
         mBatteryIntent.putExtra(BatteryManager.EXTRA_LEVEL, 1);
         mBatteryIntent.putExtra(BatteryManager.EXTRA_SCALE, 100);
@@ -70,11 +79,17 @@
         verify(mRunNotifier).pleaseStop();
     }
 
-    /**
-     * Tests that devices without batteries do not terminate on low battery.
-     */
+    /** Tests that devices without batteries do not terminate on low battery. */
     @Test
-    public void testDeviceWithoutBattery_doesNotRegisterReceiver() {
+    public void testDeviceWithoutBattery_doesNotRegisterReceiver() throws InitializationError {
+        // Create the core suite to test.
+        mSuite =
+                new LongevitySuite(
+                        TestSuite.class,
+                        new AllDefaultPossibilitiesBuilder(true),
+                        mInstrumentation,
+                        mContext,
+                        new Bundle());
         mBatteryIntent.putExtra(BatteryManager.EXTRA_PRESENT, false);
         mBatteryIntent.putExtra(BatteryManager.EXTRA_LEVEL, 1);
         mBatteryIntent.putExtra(BatteryManager.EXTRA_SCALE, 100);
@@ -82,14 +97,82 @@
         verify(mRunNotifier, never()).pleaseStop();
     }
 
+    /** Test that the runner does not report iterations when the option is not set. */
+    @Test
+    public void testReportingIteration_notSet() throws InitializationError {
+        // Create and spy the core suite to test. The option is not set as the args bundle is empty.
+        mSuite =
+                Mockito.spy(
+                        new LongevitySuite(
+                                IterationSuite.class,
+                                new AllDefaultPossibilitiesBuilder(true),
+                                mInstrumentation,
+                                mContext,
+                                new Bundle()));
+        // Store the runners that the tests are executing. Since these are object references,
+        // subsequent modifications to the runners (setting the iteration) will still be observable
+        // here.
+        List<LongevityClassRunner> runners = new ArrayList<>();
+        doAnswer(
+                        invocation -> {
+                            LongevityClassRunner runner =
+                                    (LongevityClassRunner) invocation.callRealMethod();
+                            runners.add(runner);
+                            return runner;
+                        })
+                .when(mSuite)
+                .getSuiteRunner(any(Runner.class));
+        mSuite.run(mRunNotifier);
+        Assert.assertEquals(runners.size(), 3);
+        // No runner should have a iteration number set.
+        Assert.assertEquals(runners.get(0).getIteration(), LongevityClassRunner.ITERATION_NOT_SET);
+        Assert.assertEquals(runners.get(1).getIteration(), LongevityClassRunner.ITERATION_NOT_SET);
+        Assert.assertEquals(runners.get(2).getIteration(), LongevityClassRunner.ITERATION_NOT_SET);
+    }
 
+    /** Test that the runner reports iterations when the option is set. */
+    @Test
+    public void testReportingIteration_set() throws InitializationError {
+        Bundle args = new Bundle();
+        args.putString(LongevitySuite.RENAME_ITERATION_OPTION, String.valueOf(true));
+        // Create and spy the core suite to test.
+        mSuite =
+                Mockito.spy(
+                        new LongevitySuite(
+                                IterationSuite.class,
+                                new AllDefaultPossibilitiesBuilder(true),
+                                mInstrumentation,
+                                mContext,
+                                args));
+        // Store the runners that the tests are executing. Since these are object references,
+        // subsequent modifications to the runners (setting the iteration) will still be observable
+        // here.
+        List<LongevityClassRunner> runners = new ArrayList<>();
+        doAnswer(
+                        invocation -> {
+                            LongevityClassRunner runner =
+                                    (LongevityClassRunner) invocation.callRealMethod();
+                            runners.add(runner);
+                            return runner;
+                        })
+                .when(mSuite)
+                .getSuiteRunner(any(Runner.class));
+        mSuite.run(mRunNotifier);
+        Assert.assertEquals(runners.size(), 3);
+        // Check the runners and their corresponding iterations.
+        Assert.assertTrue(runners.get(0).getDescription().getDisplayName().contains("TestOne"));
+        Assert.assertEquals(runners.get(0).getIteration(), 1);
+        Assert.assertTrue(runners.get(1).getDescription().getDisplayName().contains("TestTwo"));
+        Assert.assertEquals(runners.get(1).getIteration(), 1);
+        Assert.assertTrue(runners.get(2).getDescription().getDisplayName().contains("TestOne"));
+        Assert.assertEquals(runners.get(2).getIteration(), 2);
+    }
+
+    /** Sample device-side test cases. */
     @RunWith(LongevitySuite.class)
     @SuiteClasses({
         TestSuite.BasicTest.class,
     })
-    /**
-     * Sample device-side test cases.
-     */
     public static class TestSuite {
         // no local test cases.
 
@@ -98,4 +181,25 @@
             public void testNothing() { }
         }
     }
+
+    /** Sample test class with multiple iterations of the same test. */
+    @RunWith(LongevitySuite.class)
+    @SuiteClasses({
+        IterationSuite.TestOne.class,
+        IterationSuite.TestTwo.class,
+        IterationSuite.TestOne.class,
+    })
+    public static class IterationSuite {
+        // no local test cases.
+
+        public static class TestOne {
+            @Test
+            public void testNothing() {}
+        }
+
+        public static class TestTwo {
+            @Test
+            public void testNothing() {}
+        }
+    }
 }