diff --git a/apct-tests/perftests/windowmanager/src/android/wm/RecentsAnimationPerfTest.java b/apct-tests/perftests/windowmanager/src/android/wm/RecentsAnimationPerfTest.java
index 1667c165..6122ef2 100644
--- a/apct-tests/perftests/windowmanager/src/android/wm/RecentsAnimationPerfTest.java
+++ b/apct-tests/perftests/windowmanager/src/android/wm/RecentsAnimationPerfTest.java
@@ -23,6 +23,7 @@
 import static org.hamcrest.core.AnyOf.anyOf;
 import static org.hamcrest.core.Is.is;
 
+import android.app.ActivityManager;
 import android.app.ActivityManager.TaskSnapshot;
 import android.app.ActivityTaskManager;
 import android.app.IActivityTaskManager;
@@ -121,6 +122,12 @@
     @AfterClass
     public static void tearDownClass() {
         sSetUpClassException = null;
+        try {
+            // Recents activity may stop app switches. Restore the state to avoid affecting
+            // the next test.
+            ActivityManager.resumeAppSwitches();
+        } catch (RemoteException ignored) {
+        }
         sUiAutomation.dropShellPermissionIdentity();
     }
 
diff --git a/apct-tests/perftests/windowmanager/src/android/wm/RelayoutPerfTest.java b/apct-tests/perftests/windowmanager/src/android/wm/RelayoutPerfTest.java
index 8139a2e..f04e555 100644
--- a/apct-tests/perftests/windowmanager/src/android/wm/RelayoutPerfTest.java
+++ b/apct-tests/perftests/windowmanager/src/android/wm/RelayoutPerfTest.java
@@ -88,10 +88,7 @@
     public void testRelayout() throws Throwable {
         final Activity activity = mActivityRule.getActivity();
         final ContentView contentView = new ContentView(activity);
-        mActivityRule.runOnUiThread(() -> {
-            activity.getWindow().addFlags(WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON);
-            activity.setContentView(contentView);
-        });
+        mActivityRule.runOnUiThread(() -> activity.setContentView(contentView));
         getInstrumentation().waitForIdleSync();
 
         final RelayoutRunner relayoutRunner = new RelayoutRunner(activity, contentView.getWindow(),
diff --git a/apct-tests/perftests/windowmanager/src/android/wm/WindowManagerPerfTestBase.java b/apct-tests/perftests/windowmanager/src/android/wm/WindowManagerPerfTestBase.java
index 9e17e94..655d2f7 100644
--- a/apct-tests/perftests/windowmanager/src/android/wm/WindowManagerPerfTestBase.java
+++ b/apct-tests/perftests/windowmanager/src/android/wm/WindowManagerPerfTestBase.java
@@ -19,11 +19,13 @@
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
 import android.app.Activity;
+import android.app.KeyguardManager;
 import android.app.UiAutomation;
 import android.content.Context;
 import android.content.Intent;
 import android.os.BatteryManager;
 import android.os.ParcelFileDescriptor;
+import android.os.PowerManager;
 import android.perftests.utils.PerfTestActivity;
 import android.provider.Settings;
 
@@ -61,24 +63,32 @@
     @BeforeClass
     public static void setUpOnce() {
         final Context context = getInstrumentation().getContext();
-        sOriginalStayOnWhilePluggedIn = Settings.Global.getInt(context.getContentResolver(),
+        final int stayOnWhilePluggedIn = Settings.Global.getInt(context.getContentResolver(),
                 Settings.Global.STAY_ON_WHILE_PLUGGED_IN, 0);
-        // Keep the device awake during testing.
-        setStayOnWhilePluggedIn(BatteryManager.BATTERY_PLUGGED_USB);
+        sOriginalStayOnWhilePluggedIn = -1;
+        if (stayOnWhilePluggedIn != BatteryManager.BATTERY_PLUGGED_ANY) {
+            sOriginalStayOnWhilePluggedIn = stayOnWhilePluggedIn;
+            // Keep the device awake during testing.
+            setStayOnWhilePluggedIn(BatteryManager.BATTERY_PLUGGED_ANY);
+        }
 
         if (!BASE_OUT_PATH.exists()) {
             executeShellCommand("mkdir -p " + BASE_OUT_PATH);
         }
-        // In order to be closer to the real use case.
-        executeShellCommand("input keyevent KEYCODE_WAKEUP");
-        executeShellCommand("wm dismiss-keyguard");
+        if (!context.getSystemService(PowerManager.class).isInteractive()
+                || context.getSystemService(KeyguardManager.class).isKeyguardLocked()) {
+            executeShellCommand("input keyevent KEYCODE_WAKEUP");
+            executeShellCommand("wm dismiss-keyguard");
+        }
         context.startActivity(new Intent(Intent.ACTION_MAIN)
                 .addCategory(Intent.CATEGORY_HOME).setFlags(Intent.FLAG_ACTIVITY_NEW_TASK));
     }
 
     @AfterClass
     public static void tearDownOnce() {
-        setStayOnWhilePluggedIn(sOriginalStayOnWhilePluggedIn);
+        if (sOriginalStayOnWhilePluggedIn != -1) {
+            setStayOnWhilePluggedIn(sOriginalStayOnWhilePluggedIn);
+        }
     }
 
     private static void setStayOnWhilePluggedIn(int value) {