Snap for 5706892 from c8de01185ab893df04d9d8a1d4dc5e057aa77838 to qt-release

Change-Id: Ia614204ba2508e946d42cd0e33fda0e304eaa217

6 files changed, 660 insertions, 16 deletions

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 000000000..4e20c74ce
--- /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 000000000..b3dd138a7
--- /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 fc45c1e9a..8003aeddb 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 @@ public class BaseCollectionListener<T> extends BaseMetricListener {
     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 class BaseCollectionListener<T> extends BaseMetricListener {
     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 @@ public class BaseCollectionListener<T> extends BaseMetricListener {
     }
 
     @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 @@ public class BaseCollectionListener<T> extends BaseMetricListener {
     }
 
     /**
-     * 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 000000000..bb5a8a9e0
--- /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 ff7f8f8fb..9e239181e 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.Test;
 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 @@ public class BaseCollectionListenerTest {
     }
 
     /**
-     * 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 @@ public class BaseCollectionListenerTest {
     }
 
     /**
-     * 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 @@ public class BaseCollectionListenerTest {
         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 @@ public class BaseCollectionListenerTest {
         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 000000000..665397089
--- /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);
+    }
+}