New rule to cool down DUT with display-off.

In order to speed up DUT to cool down, added a new rule to turn off the DUT display while waiting for DUT to cool down.

Test: atest PlatformRuleTests:android.platform.test.rule.CoolDownRuleTest &&
      atest CollectorsHelperTest:com.android.helpers.ThermalHelperTest#testGetTemperature,testGetTemperatureFailed,testParseTemperatureMock
DUT: aosp_cf_x86_64_phone, aosp_oriole-userdebug

Change-Id: I197e82fb4b336bfcc727a6fff981b65ca991747a
Bug: 187311027

6 files changed, 819 insertions, 10 deletions

diff --git a/libraries/collectors-helper/statsd/src/com/android/helpers/ThermalHelper.java b/libraries/collectors-helper/statsd/src/com/android/helpers/ThermalHelper.java
index 874fb2e0f..b8a27100f 100644
--- a/libraries/collectors-helper/statsd/src/com/android/helpers/ThermalHelper.java
+++ b/libraries/collectors-helper/statsd/src/com/android/helpers/ThermalHelper.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019 The Android Open Source Project
+ * Copyright (C) 2021 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.
@@ -47,6 +47,12 @@ import java.util.regex.Pattern;
 public class ThermalHelper implements ICollectorHelper<StringBuilder> {
     private static final String LOG_TAG = ThermalHelper.class.getSimpleName();
 
+    @VisibleForTesting static final String DUMP_THERMALSERVICE_CMD = "dumpsys thermalservice";
+    private static final String METRIC_KEY_TEMPERATURE = "temperature";
+    private static final String METRIC_KEY_VALUE = "value";
+    private static final String METRIC_KEY_TYPE = "type";
+    private static final String METRIC_KEY_STATUS = "status";
+
     private static final int UNDEFINED_SEVERITY = -1;
     private static final Pattern SEVERITY_DUMPSYS_PATTERN =
             Pattern.compile("Thermal Status: (\\d+)");
@@ -63,7 +69,7 @@ public class ThermalHelper implements ICollectorHelper<StringBuilder> {
         // Add an initial value because this only detects changes.
         mInitialSeverity = UNDEFINED_SEVERITY;
         try {
-            String[] output = getDevice().executeShellCommand("dumpsys thermalservice").split("\n");
+            String[] output = getDevice().executeShellCommand(DUMP_THERMALSERVICE_CMD).split("\n");
             for (String line : output) {
                 Matcher severityMatcher = SEVERITY_DUMPSYS_PATTERN.matcher(line);
                 if (severityMatcher.matches()) {
@@ -115,24 +121,36 @@ public class ThermalHelper implements ICollectorHelper<StringBuilder> {
             }
         }
 
+        updateTemperatureMetrics(results);
+
+        return results;
+    }
+
+    /** Collect temperature metrics into result map. */
+    private void updateTemperatureMetrics(Map<String, StringBuilder> results) {
+
         try {
-            String[] output = getDevice().executeShellCommand("dumpsys thermalservice").split("\n");
+            String output = getDevice().executeShellCommand(DUMP_THERMALSERVICE_CMD);
+            String[] lines = output.split("\n");
             boolean inCurrentTempSection = false;
-            for (String line : output) {
+            for (String line : lines) {
                 Matcher temperatureMatcher = TEMPERATURE_DUMPSYS_PATTERN.matcher(line);
                 if (inCurrentTempSection && temperatureMatcher.matches()) {
                     Log.v(LOG_TAG, "Matched " + line);
                     String name = temperatureMatcher.group(3);
                     MetricUtility.addMetric(
-                            MetricUtility.constructKey("temperature", name, "value"),
+                            MetricUtility.constructKey(
+                                    METRIC_KEY_TEMPERATURE, name, METRIC_KEY_VALUE),
                             Double.parseDouble(temperatureMatcher.group(1)), // value group
                             results);
                     MetricUtility.addMetric(
-                            MetricUtility.constructKey("temperature", name, "type"),
+                            MetricUtility.constructKey(
+                                    METRIC_KEY_TEMPERATURE, name, METRIC_KEY_TYPE),
                             Integer.parseInt(temperatureMatcher.group(2)), // type group
                             results);
                     MetricUtility.addMetric(
-                            MetricUtility.constructKey("temperature", name, "status"),
+                            MetricUtility.constructKey(
+                                    METRIC_KEY_TEMPERATURE, name, METRIC_KEY_STATUS),
                             Integer.parseInt(temperatureMatcher.group(4)), // status group
                             results);
                 }
@@ -151,8 +169,27 @@ public class ThermalHelper implements ICollectorHelper<StringBuilder> {
         } catch (IOException ioe) {
             Log.e(LOG_TAG, String.format("Failed to query thermalservice. Error: %s", ioe));
         }
+    }
 
-        return results;
+    /**
+     * Get latest temperature value for needed name. Return temperature value is in unit of degree
+     * Celsius
+     */
+    public double getTemperature(String name) {
+        Map<String, StringBuilder> results = new HashMap<>();
+        updateTemperatureMetrics(results);
+        String temperatureKey =
+                MetricUtility.constructKey(METRIC_KEY_TEMPERATURE, name, METRIC_KEY_VALUE);
+        List<Double> values = MetricUtility.getMetricDoubles(temperatureKey, results);
+        if (values.size() > 0) {
+            double value =
+                    values.get(values.size() - 1).doubleValue(); // last value is the latest value.
+            Log.v(LOG_TAG, String.format("Got temperature of %s: %,.6f", name, value));
+            return value;
+        } else {
+            throw new IllegalArgumentException(
+                    String.format("Failed to get temperature of %s", name));
+        }
     }
 
     /** Remove the statsd config used to track thermal events. */
diff --git a/libraries/collectors-helper/statsd/test/src/com/android/helpers/ThermalHelperTest.java b/libraries/collectors-helper/statsd/test/src/com/android/helpers/ThermalHelperTest.java
index e10051995..f187e028a 100644
--- a/libraries/collectors-helper/statsd/test/src/com/android/helpers/ThermalHelperTest.java
+++ b/libraries/collectors-helper/statsd/test/src/com/android/helpers/ThermalHelperTest.java
@@ -18,6 +18,7 @@ package com.android.helpers;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertThrows;
 import static org.mockito.Mockito.when;
 
 import android.os.nano.OsProtoEnums;
@@ -50,6 +51,31 @@ public class ThermalHelperTest {
     private static final String THROTTLING_KEY =
             MetricUtility.constructKey("thermal", "throttling", "severity");
     private static final String FAKE_SERVICE_DUMP = "F\nA\nK\nE\nThermal Status: 2\nO\nK";
+    private static final String SERVICE_DUMP_TEMPLATE =
+            "IsStatusOverride: false\n"
+                    + "ThermalEventListeners:\n"
+                    + "        callbacks: 1\n"
+                    + "        killed: false\n"
+                    + "        broadcasts count: -1\n"
+                    + "ThermalStatusListeners:\n"
+                    + "        callbacks: 1\n"
+                    + "        killed: false\n"
+                    + "        broadcasts count: -1\n"
+                    + "Thermal Status: 0\n"
+                    + "Cached temperatures:\n"
+                    + "        Temperature{mValue=45.67, mType=3, mName=cached temperature sensor,"
+                    + " mStatus=0}\n"
+                    + "HAL Ready: true\n"
+                    + "HAL connection:\n"
+                    + "        ThermalHAL 2.0 connected: yes\n"
+                    + "Current temperatures from HAL:\n"
+                    + "        Temperature{mValue=%s, mType=3, mName=%s, mStatus=0}\n"
+                    + "        Temperature{mValue=45.6, mType=3, mName=test temperature sensor2,"
+                    + " mStatus=0}\n"
+                    + "        Temperature{mValue=56.789, mType=3, mName=test temperature sensor3,"
+                    + " mStatus=0}\n"
+                    + "Current cooling devices from HAL:\n"
+                    + "        CoolingDevice{mValue=100, mType=0, mName=test cooling device}";
 
     private ThermalHelper mThermalHelper;
     private StatsdHelper mStatsdHelper;
@@ -149,7 +175,8 @@ public class ThermalHelperTest {
     /** Test that the temperature section is parsed correctly. */
     @Test
     public void testParseTemperature() throws Exception {
-        // Use real data for this test. It should work everywhere.
+        // Use real data for this test. It should work everywhere, except for
+        // aosp_cf_x86_64_phone-userdebug.
         mThermalHelper = new ThermalHelper();
         mThermalHelper.setStatsdHelper(mStatsdHelper);
         assertTrue(mThermalHelper.startCollecting());
@@ -178,6 +205,97 @@ public class ThermalHelperTest {
         assertTrue(mThermalHelper.stopCollecting());
     }
 
+    /** Test that the mock temperature section is parsed correctly. */
+    @Test
+    public void testParseTemperatureMock() throws Exception {
+        // Use mock data for this test.
+        final String correctName = "test temperature sensor";
+        final double correctValue = 32.1;
+        final String correctOutput =
+                String.format(SERVICE_DUMP_TEMPLATE, String.valueOf(correctValue), correctName);
+        when(mDevice.executeShellCommand(ThermalHelper.DUMP_THERMALSERVICE_CMD))
+                .thenReturn(correctOutput);
+        Map<String, StringBuilder> metrics = mThermalHelper.getMetrics();
+        // Validate at least 2 temperature keys exist with all 3 metrics.
+        int statusMetricsFound = 0;
+        int valueMetricsFound = 0;
+        int typeMetricsFound = 0;
+        for (String key : metrics.keySet()) {
+            if (!key.startsWith("temperature")) {
+                continue;
+            }
+
+            if (key.endsWith("status")) {
+                statusMetricsFound++;
+            } else if (key.endsWith("value")) {
+                valueMetricsFound++;
+            } else if (key.endsWith("type")) {
+                typeMetricsFound++;
+            }
+        }
+
+        assertTrue(
+                "Didn't find at least 2 status, value, and type temperature metrics.",
+                statusMetricsFound >= 2 && valueMetricsFound >= 2 && typeMetricsFound >= 2);
+    }
+
+    /** Test getting temperature value from DUT correctly. */
+    @Test
+    public void testGetTemperature() throws Exception {
+        final double THRESHOLD = 0.0001;
+        final String correctName = "test temperature sensor";
+        final double correctValue = 32.1;
+        final String correctOutput =
+                String.format(SERVICE_DUMP_TEMPLATE, String.valueOf(correctValue), correctName);
+
+        when(mDevice.executeShellCommand(ThermalHelper.DUMP_THERMALSERVICE_CMD))
+                .thenReturn(correctOutput);
+        assertTrue(Math.abs(mThermalHelper.getTemperature(correctName) - correctValue) < THRESHOLD);
+    }
+
+    /** Test failing to get temperature value from DUT. */
+    @Test
+    public void testGetTemperatureFailed() throws Exception {
+        final String correctName = "test temperature sensor";
+        final double correctValue = 32.1;
+        final String correctOutput =
+                String.format(SERVICE_DUMP_TEMPLATE, String.valueOf(correctValue), correctName);
+        final String expectedMessage = "Failed to get temperature of ";
+
+        final String badName = "bad temperature sensor";
+        when(mDevice.executeShellCommand(ThermalHelper.DUMP_THERMALSERVICE_CMD))
+                .thenReturn(correctOutput);
+        Exception exception1 =
+                assertThrows(
+                        IllegalArgumentException.class,
+                        () -> {
+                            mThermalHelper.getTemperature(badName);
+                        });
+        assertTrue(exception1.getMessage().contains(expectedMessage));
+
+        final String badOutput = String.format(SERVICE_DUMP_TEMPLATE, "bad", correctName);
+        when(mDevice.executeShellCommand(ThermalHelper.DUMP_THERMALSERVICE_CMD))
+                .thenReturn(badOutput);
+        Exception exception2 =
+                assertThrows(
+                        IllegalArgumentException.class,
+                        () -> {
+                            mThermalHelper.getTemperature(correctName);
+                        });
+        assertTrue(exception2.getMessage().contains(expectedMessage));
+
+        final String badOutputEmpty = String.format(SERVICE_DUMP_TEMPLATE, "", correctName);
+        when(mDevice.executeShellCommand(ThermalHelper.DUMP_THERMALSERVICE_CMD))
+                .thenReturn(badOutputEmpty);
+        Exception exception3 =
+                assertThrows(
+                        IllegalArgumentException.class,
+                        () -> {
+                            mThermalHelper.getTemperature(correctName);
+                        });
+        assertTrue(exception3.getMessage().contains(expectedMessage));
+    }
+
     /**
      * Returns a list of {@link com.android.os.nano.StatsLog.EventMetricData} that statsd returns.
      */
diff --git a/libraries/collectors-helper/utilities/src/com/android/helpers/MetricUtility.java b/libraries/collectors-helper/utilities/src/com/android/helpers/MetricUtility.java
index 5e7e8e5a0..def11392f 100644
--- a/libraries/collectors-helper/utilities/src/com/android/helpers/MetricUtility.java
+++ b/libraries/collectors-helper/utilities/src/com/android/helpers/MetricUtility.java
@@ -5,9 +5,12 @@ import android.os.ParcelFileDescriptor;
 import android.util.Log;
 
 import java.io.ByteArrayOutputStream;
-import java.io.IOException;
 import java.io.InputStream;
+import java.io.IOException;
 import java.text.DecimalFormat;
+import java.text.ParseException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -73,6 +76,43 @@ public class MetricUtility {
     }
 
     /**
+     * Get metric values from result map.
+     *
+     * @param metricKey Unique key to track the metric.
+     * @param resultMap Map of all the metrics.
+     * @return Double List of metric values for metric key
+     */
+    public static List<Double> getMetricDoubles(
+            String metricKey, Map<String, StringBuilder> resultMap) {
+        List<Double> result = new ArrayList<Double>();
+        if (!resultMap.containsKey(metricKey)) {
+            Log.e(TAG, String.format("No such metric key %s", metricKey));
+            return result;
+        } else {
+            String value = resultMap.get(metricKey).toString();
+            if (value.length() == 0) {
+                Log.e(TAG, String.format("Missed value of metric key %s", metricKey));
+                return result;
+            } else {
+                String[] values = value.split(METRIC_SEPARATOR);
+                for (int i = 0; i < values.length; i++) {
+                    try {
+                        result.add(DOUBLE_FORMAT.parse(values[i]).doubleValue());
+                    } catch (ParseException e) {
+                        Log.e(
+                                TAG,
+                                String.format(
+                                        "Error parsing value of metric key %s: #%d of value %s",
+                                        metricKey, i, value));
+                        return new ArrayList<Double>();
+                    }
+                }
+            }
+        }
+        return result;
+    }
+
+    /**
      * Turn executeShellCommand into a blocking operation.
      *
      * @param command shell command to be executed.
diff --git a/libraries/health/rules/Android.bp b/libraries/health/rules/Android.bp
index c72e8a8e1..7236494bd 100644
--- a/libraries/health/rules/Android.bp
+++ b/libraries/health/rules/Android.bp
@@ -28,6 +28,7 @@ java_library {
         "package-helper",
         "launcher-aosp-tapl",
         "flickerlib",
+        "statsd-helper",
     ],
     srcs: ["src/**/*.java"],
 }
@@ -42,6 +43,7 @@ java_library {
         "guava",
         "memory-helper",
         "package-helper",
+        "statsd-helper",
         "launcher-aosp-tapl",
     ],
     srcs: ["src/**/*.java"],
diff --git a/libraries/health/rules/src/android/platform/test/rule/CoolDownRule.java b/libraries/health/rules/src/android/platform/test/rule/CoolDownRule.java
new file mode 100644
index 000000000..fe52efb17
--- /dev/null
+++ b/libraries/health/rules/src/android/platform/test/rule/CoolDownRule.java
@@ -0,0 +1,224 @@
+/*
+ * Copyright (C) 2021 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.platform.test.rule;
+
+import android.os.RemoteException;
+import android.os.SystemClock;
+import android.util.Log;
+import androidx.annotation.VisibleForTesting;
+
+import com.android.helpers.ThermalHelper;
+
+import java.text.CharacterIterator;
+import java.text.StringCharacterIterator;
+
+import org.junit.runner.Description;
+
+/**
+ * This rule will cool DUT phone before a test case.
+ *
+ * <p>The rule takes a temperature name from command-line argument
+ * "cooldown-device-temperature-name" to check temperature. The spaces in temperature name must be
+ * replaced with '#' and '#' should be replaced with '##'. The rule unescapes argument value string
+ * of temperature name accordingly.
+ */
+public class CoolDownRule extends TestWatcher {
+
+    private static final String LOG_TAG = CoolDownRule.class.getSimpleName();
+
+    // Interval in seconds, to poll for device temperature; defaults to 30s
+    @VisibleForTesting static final String POLL_INTERVAL_OPTION = "cooldown-poll-interval";
+    private long mPollIntervalSecs = 30;
+
+    // Max wait time in seconds, for device cool down to target temperature; defaults to 20 minutes
+    @VisibleForTesting static final String MAX_WAIT_OPTION = "cooldown-max-wait";
+    private long mMaxWaitSecs = 60 * 20;
+
+    // If test should be aborted if device is still above expected temperature; defaults to false
+    @VisibleForTesting static final String ABORT_ON_TIMEOUT_OPTION = "cooldown-abort-on-timeout";
+    private boolean mAbortOnTimeout = false;
+
+    // Additional time to wait in seconds, after temperature has reached to target; defaults to 30s
+    @VisibleForTesting static final String POST_IDLE_WAIT_OPTION = "cooldown-post-idle-wait";
+    private long mPostIdleWaitSecs = 30;
+
+    @VisibleForTesting
+    static final String DEVICE_TEMPERATURE_NAME_OPTION = "cooldown-device-temperature-name";
+
+    private String mDeviceTemperatureName = "";
+
+    // Target Temperature that device should have; defaults to 35 degree Celsius
+    @VisibleForTesting
+    static final String TARGET_TEMPERATURE_OPTION = "cooldown-target-temperature";
+
+    private int mTargetTemperature = 35;
+
+    private ThermalHelper mThermalHelper;
+
+    @Override
+    protected void starting(Description description) {
+        mDeviceTemperatureName =
+                CoolDownRule.unescapeOptionStr(
+                        getArguments().getString(DEVICE_TEMPERATURE_NAME_OPTION, ""));
+        if (mDeviceTemperatureName.isEmpty()) {
+            throw new IllegalArgumentException("Missed device temperature name.");
+        }
+        mPollIntervalSecs = Long.valueOf(getArguments().getString(POLL_INTERVAL_OPTION, "30"));
+        mMaxWaitSecs = Long.valueOf(getArguments().getString(MAX_WAIT_OPTION, "1200"));
+        mAbortOnTimeout =
+                Boolean.valueOf(getArguments().getString(ABORT_ON_TIMEOUT_OPTION, "false"));
+        mPostIdleWaitSecs = Long.valueOf(getArguments().getString(POST_IDLE_WAIT_OPTION, "30"));
+        mTargetTemperature =
+                Integer.valueOf(getArguments().getString(TARGET_TEMPERATURE_OPTION, "35"));
+        if (mTargetTemperature > (100) || mTargetTemperature <= 0) {
+            throw new IllegalArgumentException(
+                    String.format(
+                            "Invalid target target temperature: %d degree Celsius",
+                            mTargetTemperature));
+        }
+        mThermalHelper = initThermalHelper();
+
+        try {
+            // Turn off the screen if necessary.
+            final boolean screenOnOriginal = getUiDevice().isScreenOn();
+            if (screenOnOriginal) {
+                getUiDevice().sleep();
+            }
+
+            waitTemperature();
+
+            // Turn on the screen if necessary.
+            if (screenOnOriginal && !getUiDevice().isScreenOn()) {
+                getUiDevice().wakeUp();
+            }
+        } catch (RemoteException e) {
+            throw new RuntimeException("Could not cool down device in time.", e);
+        }
+    }
+
+    @VisibleForTesting
+    ThermalHelper initThermalHelper() {
+        return new ThermalHelper();
+    }
+
+    private void waitTemperature() {
+        long start = System.currentTimeMillis();
+        long maxWaitMs = mMaxWaitSecs * 1000;
+        long intervalMs = mPollIntervalSecs * 1000;
+        int deviceTemperature = getDeviceTemperature(mDeviceTemperatureName);
+
+        while ((deviceTemperature > mTargetTemperature) && (elapsedMs(start) < maxWaitMs)) {
+            Log.i(
+                    LOG_TAG,
+                    String.format(
+                            "Temperature is still high actual %d/expected %d",
+                            deviceTemperature, mTargetTemperature));
+            sleepMs(intervalMs);
+            deviceTemperature = getDeviceTemperature(mDeviceTemperatureName);
+        }
+
+        if (deviceTemperature <= mTargetTemperature) {
+            Log.i(
+                    LOG_TAG,
+                    String.format(
+                            "Total time elapsed to get to %dc : %ds",
+                            mTargetTemperature, (System.currentTimeMillis() - start) / 1000));
+        } else {
+            Log.w(
+                    LOG_TAG,
+                    String.format(
+                            "Temperature is still high, actual %d/expected %d; waiting after %ds",
+                            deviceTemperature,
+                            mTargetTemperature,
+                            (System.currentTimeMillis() - start) / 1000));
+            if (mAbortOnTimeout) {
+                throw new IllegalStateException(
+                        String.format(
+                                "Temperature is still high after wait timeout; actual %d/expected"
+                                        + " %d",
+                                deviceTemperature, mTargetTemperature));
+            }
+        }
+
+        // Extra idle time after reaching the target to stabilize the system
+        sleepMs(mPostIdleWaitSecs * 1000);
+        Log.i(
+                LOG_TAG,
+                String.format(
+                        "Done waiting, total time elapsed: %ds",
+                        (System.currentTimeMillis() - start) / 1000));
+    }
+
+    @VisibleForTesting
+    void sleepMs(long milliSeconds) {
+        SystemClock.sleep(milliSeconds);
+    }
+
+    @VisibleForTesting
+    long elapsedMs(long start) {
+        return System.currentTimeMillis() - start;
+    }
+
+    /**
+     * @param name : temperature name in need.
+     * @return Device temperature in unit of millidegree Celsius
+     */
+    @VisibleForTesting
+    int getDeviceTemperature(String name) {
+        return (int) mThermalHelper.getTemperature(name);
+    }
+
+    /**
+     * @param input : Option value string to be unescaped. Option value string should be escaped by
+     *     replacing ' ' (space) with '#' and '#' with '##'. Not support to unescaped consecutive
+     *     spaces like ' ' and space before '#' like ' #'.
+     * @return Unescaped string.
+     */
+    @VisibleForTesting
+    static String unescapeOptionStr(String input) {
+        final StringBuilder result = new StringBuilder();
+        final StringCharacterIterator iterator = new StringCharacterIterator(input);
+        final char escape_char = '#';
+        final char space_char = ' ';
+        char character = iterator.current();
+        boolean escapeFlag = false;
+
+        while (character != CharacterIterator.DONE) {
+            if (character == space_char) {
+                throw new IllegalArgumentException(
+                        "Non-escaped option value string (please replace ' ' to '#'; '#' to '##'): "
+                                + input);
+            } else if (escapeFlag) {
+                if (character == escape_char) {
+                    result.append(escape_char);
+                } else {
+                    result.append(space_char);
+                    result.append(character);
+                }
+                escapeFlag = false;
+            } else if (character == escape_char) {
+                escapeFlag = true;
+            } else {
+                result.append(character);
+            }
+            character = iterator.next();
+        }
+        if (escapeFlag) {
+            result.append(space_char);
+        }
+        return result.toString();
+    }
+}
diff --git a/libraries/health/rules/tests/src/android/platform/test/rule/CoolDownRuleTest.java b/libraries/health/rules/tests/src/android/platform/test/rule/CoolDownRuleTest.java
new file mode 100644
index 000000000..3bb3104b3
--- /dev/null
+++ b/libraries/health/rules/tests/src/android/platform/test/rule/CoolDownRuleTest.java
@@ -0,0 +1,388 @@
+/*
+ * Copyright (C) 2021 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.platform.test.rule;
+
+import android.os.Bundle;
+import android.os.RemoteException;
+import androidx.test.uiautomator.UiDevice;
+
+import com.android.helpers.ThermalHelper;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.junit.runner.Description;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.junit.runners.model.Statement;
+import org.junit.Test;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.fail;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/** Unit test the logic for {@link CoolDownRule} */
+@RunWith(JUnit4.class)
+public class CoolDownRuleTest {
+    private static final int TARGET_TEMPERATURE = 35;
+    private static final int POLL_INTERVAL = 987;
+    private static final int IDLE_WAIT = 123;
+    private static final String TEMPERATURE_NAME = "test temperature sensor";
+    private static final String TEMPERATURE_NAME_ESCAPED = "test#temperature#sensor";
+    private static final int OVERHEAT_COUNT = 3;
+
+    private static final String OPS_SCREEN_ON = "screen On";
+    private static final String OPS_SCREEN_OFF = "screen Off";
+    private static final String OPS_TEST = "test";
+    private static final String OPS_SLEEP_INTERVAL = "sleep 987000 milli seconds";
+    private static final String OPS_SLEEP_IDLE = "sleep 123000 milli seconds";
+
+    private final ThermalHelper mThermalHelper = mock(ThermalHelper.class);
+
+    /** Tests that this rule will complete cool down as expected steps. */
+    @Test
+    public void testCoolDownNormal() throws Throwable {
+        boolean screenOn = true;
+        boolean abortOnTimeout = false;
+        int maxWait = (OVERHEAT_COUNT * POLL_INTERVAL) + 5;
+        TestableRule rule = getDefaultRule(screenOn, maxWait, abortOnTimeout);
+
+        doAnswer(new TemperatureAnswer(TARGET_TEMPERATURE, OVERHEAT_COUNT))
+                .when(mThermalHelper)
+                .getTemperature(TEMPERATURE_NAME);
+        rule.apply(rule.getTestStatement(), Description.createTestDescription("clzz", "mthd"))
+                .evaluate();
+        assertThat(rule.getOperations())
+                .containsExactly(
+                        OPS_SCREEN_OFF,
+                        "get device temperature degree: 46",
+                        OPS_SLEEP_INTERVAL,
+                        "get device temperature degree: 41",
+                        OPS_SLEEP_INTERVAL,
+                        "get device temperature degree: 37",
+                        OPS_SLEEP_INTERVAL,
+                        "get device temperature degree: 33",
+                        OPS_SLEEP_IDLE,
+                        OPS_SCREEN_ON,
+                        OPS_TEST)
+                .inOrder();
+    }
+
+    /** Tests that this rule will fail to cool down due to timeout as expected steps. */
+    @Test
+    public void testCoolDownTimeout() throws Throwable {
+        boolean screenOn = true;
+        boolean abortOnTimeout = false;
+        int maxWait = (OVERHEAT_COUNT * POLL_INTERVAL) - 5;
+        TestableRule rule = getDefaultRule(screenOn, maxWait, abortOnTimeout);
+
+        double cooldownOffset = -1.0; // heat up instead of cool down
+        doAnswer(new TemperatureAnswer(TARGET_TEMPERATURE, OVERHEAT_COUNT, cooldownOffset))
+                .when(mThermalHelper)
+                .getTemperature(TEMPERATURE_NAME);
+        rule.apply(rule.getTestStatement(), Description.createTestDescription("clzz", "mthd"))
+                .evaluate();
+        assertThat(rule.getOperations())
+                .containsExactly(
+                        OPS_SCREEN_OFF,
+                        "get device temperature degree: 37",
+                        OPS_SLEEP_INTERVAL,
+                        "get device temperature degree: 38",
+                        OPS_SLEEP_INTERVAL,
+                        "get device temperature degree: 39",
+                        OPS_SLEEP_INTERVAL,
+                        "get device temperature degree: 40",
+                        OPS_SLEEP_IDLE,
+                        OPS_SCREEN_ON,
+                        OPS_TEST)
+                .inOrder();
+    }
+
+    /**
+     * Tests that this rule will fail to cool down due to timeout and throw exception to abort the
+     * test as expected steps.
+     */
+    @Test
+    public void testCoolDownTimeoutAbort() throws Throwable {
+        boolean screenOn = true;
+        boolean abortOnTimeout = true;
+        int maxWait = (OVERHEAT_COUNT * POLL_INTERVAL) - 5;
+        TestableRule rule = getDefaultRule(screenOn, maxWait, abortOnTimeout);
+
+        double cooldownOffset = -1.0; // heat up instead of cool down
+        doAnswer(new TemperatureAnswer(TARGET_TEMPERATURE, OVERHEAT_COUNT, cooldownOffset))
+                .when(mThermalHelper)
+                .getTemperature(TEMPERATURE_NAME);
+        try {
+            rule.apply(rule.getTestStatement(), Description.createTestDescription("clzz", "mthd"))
+                    .evaluate();
+            fail("An exception should have been thrown.");
+        } catch (IllegalStateException e) {
+            assertThat(rule.getOperations())
+                    .containsExactly(
+                            OPS_SCREEN_OFF,
+                            "get device temperature degree: 37",
+                            OPS_SLEEP_INTERVAL,
+                            "get device temperature degree: 38",
+                            OPS_SLEEP_INTERVAL,
+                            "get device temperature degree: 39",
+                            OPS_SLEEP_INTERVAL,
+                            "get device temperature degree: 40")
+                    .inOrder();
+        }
+    }
+
+    /**
+     * Tests that this rule will complete cool down without turning off screen as expected steps.
+     */
+    @Test
+    public void testCoolDownScreenOff() throws Throwable {
+        boolean screenOn = false;
+        boolean abortOnTimeout = false;
+        int maxWait = (OVERHEAT_COUNT * POLL_INTERVAL) + 5;
+        TestableRule rule = getDefaultRule(screenOn, maxWait, abortOnTimeout);
+
+        doAnswer(new TemperatureAnswer(TARGET_TEMPERATURE, OVERHEAT_COUNT))
+                .when(mThermalHelper)
+                .getTemperature(TEMPERATURE_NAME);
+        rule.apply(rule.getTestStatement(), Description.createTestDescription("clzz", "mthd"))
+                .evaluate();
+        assertThat(rule.getOperations())
+                .containsExactly(
+                        "get device temperature degree: 46",
+                        OPS_SLEEP_INTERVAL,
+                        "get device temperature degree: 41",
+                        OPS_SLEEP_INTERVAL,
+                        "get device temperature degree: 37",
+                        OPS_SLEEP_INTERVAL,
+                        "get device temperature degree: 33",
+                        OPS_SLEEP_IDLE,
+                        OPS_TEST)
+                .inOrder();
+    }
+    /** Tests to confirm option value strings will be escaped correctly. */
+    @Test
+    public void testEscapedOptionStrs() throws Throwable {
+        Map<String, String> escape_test_cases = new HashMap<String, String>();
+        escape_test_cases.put("", "");
+        escape_test_cases.put("#", " ");
+        escape_test_cases.put("##", "#");
+        escape_test_cases.put("prefix#", "prefix ");
+        escape_test_cases.put("prefix##", "prefix#");
+        escape_test_cases.put("#postfix", " postfix");
+        escape_test_cases.put("##postfix", "#postfix");
+        escape_test_cases.put("prefix#postfix", "prefix postfix");
+        escape_test_cases.put("prefix##postfix", "prefix#postfix");
+        escape_test_cases.put("###", "# ");
+        escape_test_cases.put("####", "##");
+        escape_test_cases.put("prefix###", "prefix# ");
+        escape_test_cases.put("prefix####", "prefix##");
+        escape_test_cases.put("###postfix", "# postfix");
+        escape_test_cases.put("####postfix", "##postfix");
+        escape_test_cases.put("prefix###postfix", "prefix# postfix");
+        escape_test_cases.put("prefix####postfix", "prefix##postfix");
+        escape_test_cases.put("#####", "## ");
+        escape_test_cases.put("######", "###");
+        for (String input : escape_test_cases.keySet()) {
+            String expected_output = escape_test_cases.get(input);
+            String actual_output = CoolDownRule.unescapeOptionStr(input);
+            assertThat(expected_output.equals(actual_output)).isTrue();
+        }
+    }
+    /** Tests to detect unescaped option value strings. */
+    @Test
+    public void testNonEscapeOptionStrs() throws Throwable {
+        String inputs[] =
+                new String[] {
+                    " ",
+                    "prefix postfix",
+                    " #",
+                    "# ",
+                    "prefix #",
+                    "# postfix",
+                    "prefix #postfix",
+                    "prefix# postfix",
+                    " ##",
+                    "## ",
+                    "# #",
+                    " prefix##",
+                    "##postfix ",
+                    "prefix ##postfix",
+                    "prefix# #postfix",
+                    "prefix## postfix",
+                    " ###",
+                    "# ##",
+                    "## #",
+                    "### ",
+                };
+
+        for (String input : inputs) {
+            assertThrows(
+                    IllegalArgumentException.class, () -> CoolDownRule.unescapeOptionStr(input));
+        }
+    }
+
+    private TestableRule getDefaultRule(boolean screenOn, int maxWait, boolean abortOnTimeout) {
+        TestableRule rule = new TestableRule(screenOn, mThermalHelper);
+        rule.addArg(CoolDownRule.POLL_INTERVAL_OPTION, String.valueOf(POLL_INTERVAL));
+        rule.addArg(CoolDownRule.MAX_WAIT_OPTION, String.valueOf(maxWait));
+        rule.addArg(CoolDownRule.ABORT_ON_TIMEOUT_OPTION, String.valueOf(abortOnTimeout));
+        rule.addArg(CoolDownRule.POST_IDLE_WAIT_OPTION, String.valueOf(IDLE_WAIT));
+        rule.addArg(CoolDownRule.DEVICE_TEMPERATURE_NAME_OPTION, TEMPERATURE_NAME_ESCAPED);
+        rule.addArg(CoolDownRule.TARGET_TEMPERATURE_OPTION, String.valueOf(TARGET_TEMPERATURE));
+        return rule;
+    }
+
+    private static class TestableRule extends CoolDownRule {
+        private final UiDevice mUiDevice;
+
+        private List<String> mOperations = new ArrayList<>();
+        private Bundle mBundle = new Bundle();
+        private boolean mIsScreenOn = true;
+        private long mTotalSleepMs = 0;
+        private final ThermalHelper mThermalHelper;
+
+        private TestableRule() {
+            this(true, null);
+        }
+
+        public TestableRule(boolean screenOn, ThermalHelper thermalHelper) {
+            mUiDevice = mock(UiDevice.class);
+            mIsScreenOn = screenOn;
+            mThermalHelper = thermalHelper;
+        }
+
+        @Override
+        protected UiDevice getUiDevice() {
+            try {
+                when(mUiDevice.isScreenOn()).thenReturn(mIsScreenOn);
+                doAnswer(
+                                invocation -> {
+                                    return setScreen(false);
+                                })
+                        .when(mUiDevice)
+                        .sleep();
+                doAnswer(
+                                invocation -> {
+                                    return setScreen(true);
+                                })
+                        .when(mUiDevice)
+                        .wakeUp();
+            } catch (RemoteException e) {
+                throw new RuntimeException("Could not unlock device.", e);
+            }
+            return mUiDevice;
+        }
+
+        @Override
+        ThermalHelper initThermalHelper() {
+            return mThermalHelper;
+        }
+
+        @Override
+        protected Bundle getArguments() {
+            return mBundle;
+        }
+
+        @Override
+        int getDeviceTemperature(String name) {
+            int value = super.getDeviceTemperature(name);
+            mOperations.add(String.format("get device temperature degree: %d", value));
+            return value;
+        }
+
+        @Override
+        protected void sleepMs(long milliSeconds) {
+            mOperations.add(String.format("sleep %d milli seconds", milliSeconds));
+            mTotalSleepMs += milliSeconds;
+        }
+
+        @Override
+        protected long elapsedMs(long start) {
+            long ms = super.elapsedMs(start);
+            return (mTotalSleepMs + ms);
+        }
+
+        public List<String> getOperations() {
+            return mOperations;
+        }
+
+        public void addArg(String key, String value) {
+            mBundle.putString(key, value);
+        }
+
+        public Object setScreen(boolean screenOn) {
+            mIsScreenOn = screenOn;
+            mOperations.add(mIsScreenOn ? OPS_SCREEN_ON : OPS_SCREEN_OFF);
+            return null;
+        }
+
+        public Statement getTestStatement() {
+            return new Statement() {
+                @Override
+                public void evaluate() throws Throwable {
+                    mOperations.add(OPS_TEST);
+                }
+            };
+        }
+    }
+
+    private static class TemperatureAnswer implements Answer {
+        private static final double INIT_OFFSET = 2.0;
+        private static final double DEFAULT_COOLDOWN_OFFSET = 4.4;
+
+        private final double targetTemperature;
+        private final double cooldownOffset;
+        private final int overHeatCount;
+
+        private double temperature = 0.0;
+
+        private TemperatureAnswer() {
+            this(36.0, 3, DEFAULT_COOLDOWN_OFFSET);
+        }
+
+        public TemperatureAnswer(double targetTemperatureIn, int overHeatCount) {
+            this(targetTemperatureIn, overHeatCount, DEFAULT_COOLDOWN_OFFSET);
+        }
+
+        public TemperatureAnswer(
+                double targetTemperatureIn, int overHeatCountIn, double cooldownOffsetIn) {
+            targetTemperature = targetTemperatureIn;
+            overHeatCount = overHeatCountIn;
+            cooldownOffset = cooldownOffsetIn;
+
+            if (cooldownOffset > 0) { // cool down by turn
+                temperature = targetTemperature - INIT_OFFSET;
+                temperature += cooldownOffset * overHeatCount;
+            } else { // from warm to warmer
+                temperature = targetTemperature + INIT_OFFSET;
+            }
+        }
+
+        @Override
+        public Double answer(InvocationOnMock invocation) {
+            double result = temperature;
+            temperature -= cooldownOffset;
+            return Double.valueOf(result);
+        }
+    }
+}