Snap for 6123600 from 2212ad8f30dd5d9c44d17a14ca483c52cc02fbee to qt-qpr2-releaseandroid-10.0.0_r36android-10.0.0_r35android-10.0.0_r34android-10.0.0_r33android-10.0.0_r32android-10.0.0_r31android-10.0.0_r30android10-qpr2-s4-releaseandroid10-qpr2-s3-releaseandroid10-qpr2-s2-releaseandroid10-qpr2-s1-releaseandroid10-qpr2-release

Change-Id: I999c1271707185411b043bfdd74cac87bcbc6ccc

1 files changed, 21 insertions, 16 deletions

diff --git a/tests/health/scenarios/src/android/platform/test/scenario/system/ScreenOff.java b/tests/health/scenarios/src/android/platform/test/scenario/system/ScreenOff.java
index d3dc78022..1f5e23475 100644
--- a/tests/health/scenarios/src/android/platform/test/scenario/system/ScreenOff.java
+++ b/tests/health/scenarios/src/android/platform/test/scenario/system/ScreenOff.java
@@ -18,44 +18,49 @@ package android.platform.test.scenario.system;
 
 import android.os.RemoteException;
 import android.os.SystemClock;
+import android.platform.test.option.BooleanOption;
+import android.platform.test.option.LongOption;
 import android.platform.test.scenario.annotation.Scenario;
 import android.support.test.uiautomator.UiDevice;
 import androidx.test.InstrumentationRegistry;
 
+import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
-/**
- * Shuts the screen off and waits for a specified amount of time.
- */
+/** Shuts the screen off and waits for a specified amount of time. */
 @Scenario
 @RunWith(JUnit4.class)
 public class ScreenOff {
-    private static final String DURATION_OPTION = "screenOffDurationMs";
-    private static final String DURATION_DEFAULT = "1000";
+    @Rule public LongOption mDurationMs = new LongOption("screenOffDurationMs").setDefault(1000L);
+
+    @Rule
+    public BooleanOption mTurnScreenBackOn =
+            new BooleanOption("screenOffTurnScreenBackOnAfterTest").setDefault(false);
 
-    private long mDurationMs = 0L;
     private UiDevice mDevice;
 
     @Before
     public void setUp() {
         mDevice = UiDevice.getInstance(InstrumentationRegistry.getInstrumentation());
-        String durationMsString =
-                InstrumentationRegistry.getArguments().getString(DURATION_OPTION, DURATION_DEFAULT);
-        try {
-            mDurationMs = Long.parseLong(durationMsString);
-        } catch (NumberFormatException e) {
-            throw new IllegalArgumentException(
-                    String.format(
-                            "Failed to parse option %s: %s", DURATION_OPTION, durationMsString));
-        }
     }
 
     @Test
     public void testScreenOff() throws RemoteException {
         mDevice.sleep();
-        SystemClock.sleep(mDurationMs);
+        SystemClock.sleep(mDurationMs.get());
+    }
+
+    @After
+    public void tearDown() throws RemoteException {
+        if (mTurnScreenBackOn.get()) {
+            mDevice.wakeUp();
+            mDevice.waitForIdle();
+            mDevice.pressMenu();
+            mDevice.waitForIdle();
+        }
     }
 }