Optionally turn screen back on after the ScreenOff CUJ.

Also updated the CUJ to use platform-test-options.

Bug: 144721773, 118122395
Test: Manual
Change-Id: Id0bff1ed3bf93bf91a82147d423de4dfebac4614
(cherry picked from commit 706448be9cb8d00abb389936dafe8ba80ac36fd2)

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();
+        }
     }
 }