diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java
index 8e4ecf1b2a264..9b0f952133aa0 100644
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java
@@ -61,10 +61,11 @@ import java.util.List;
 public final class AutoShowTest {
 
     @Rule public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
-
-    @Rule
-    public ScreenCaptureRule mScreenCaptureRule =
+    @Rule public ScreenCaptureRule mScreenCaptureRule =
             new ScreenCaptureRule("/sdcard/InputMethodStressTest");
+    @Rule public DisableLockScreenRule mDisableLockScreenRule = new DisableLockScreenRule();
+    @Rule public ScreenOrientationRule mScreenOrientationRule =
+            new ScreenOrientationRule(true /* isPortrait */);
 
     // TODO(b/240359838): add test case {@code Configuration.SCREENLAYOUT_SIZE_LARGE}.
     @Parameterized.Parameters(
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/DisableLockScreenRule.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/DisableLockScreenRule.java
new file mode 100644
index 0000000000000..d95decff2d865
--- /dev/null
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/DisableLockScreenRule.java
@@ -0,0 +1,53 @@
+/*
+ * Copyright (C) 2022 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.inputmethod.stresstest;
+
+import android.support.test.uiautomator.UiDevice;
+
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import org.junit.rules.TestWatcher;
+import org.junit.runner.Description;
+
+import java.io.IOException;
+
+/** Disable lock screen during the test. */
+public class DisableLockScreenRule extends TestWatcher {
+    private static final String LOCK_SCREEN_OFF_COMMAND = "locksettings set-disabled true";
+    private static final String LOCK_SCREEN_ON_COMMAND = "locksettings set-disabled false";
+
+    private final UiDevice mUiDevice =
+            UiDevice.getInstance(InstrumentationRegistry.getInstrumentation());
+
+    @Override
+    protected void starting(Description description) {
+        try {
+            mUiDevice.executeShellCommand(LOCK_SCREEN_OFF_COMMAND);
+        } catch (IOException e) {
+            throw new RuntimeException("Could not disable lock screen.", e);
+        }
+    }
+
+    @Override
+    protected void finished(Description description) {
+        try {
+            mUiDevice.executeShellCommand(LOCK_SCREEN_ON_COMMAND);
+        } catch (IOException e) {
+            throw new RuntimeException("Could not enable lock screen.", e);
+        }
+    }
+}
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeOpenCloseStressTest.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeOpenCloseStressTest.java
index 82acfb646d191..a6131ce45c8bc 100644
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeOpenCloseStressTest.java
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeOpenCloseStressTest.java
@@ -67,10 +67,11 @@ public final class ImeOpenCloseStressTest {
     private static final int NUM_TEST_ITERATIONS = 10;
 
     @Rule public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
-
-    @Rule
-    public ScreenCaptureRule mScreenCaptureRule =
+    @Rule public ScreenCaptureRule mScreenCaptureRule =
             new ScreenCaptureRule("/sdcard/InputMethodStressTest");
+    @Rule public DisableLockScreenRule mDisableLockScreenRule = new DisableLockScreenRule();
+    @Rule public ScreenOrientationRule mScreenOrientationRule =
+            new ScreenOrientationRule(true /* isPortrait */);
 
     private final Instrumentation mInstrumentation;
     private final int mSoftInputFlags;
@@ -485,7 +486,6 @@ public final class ImeOpenCloseStressTest {
 
         UiDevice uiDevice = UiDevice.getInstance(mInstrumentation);
 
-        uiDevice.freezeRotation();
         uiDevice.setOrientationRight();
         uiDevice.waitForIdle();
         Thread.sleep(1000);
@@ -502,7 +502,6 @@ public final class ImeOpenCloseStressTest {
 
         uiDevice.setOrientationNatural();
         uiDevice.waitForIdle();
-        uiDevice.unfreezeRotation();
     }
 
     private static void verifyShowBehavior(TestActivity activity) {
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ScreenOrientationRule.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ScreenOrientationRule.java
new file mode 100644
index 0000000000000..bc3b1efcf6b5f
--- /dev/null
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ScreenOrientationRule.java
@@ -0,0 +1,66 @@
+/*
+ * Copyright (C) 2022 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.inputmethod.stresstest;
+
+import android.os.RemoteException;
+import android.support.test.uiautomator.UiDevice;
+
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import org.junit.rules.TestWatcher;
+import org.junit.runner.Description;
+
+import java.io.IOException;
+
+/**
+ * Disable auto-rotate during the test and set the screen orientation to portrait or landscape
+ * before the test starts.
+ */
+public class ScreenOrientationRule extends TestWatcher {
+    private static final String SET_PORTRAIT_MODE_CMD = "settings put system user_rotation 0";
+    private static final String SET_LANDSCAPE_MODE_CMD = "settings put system user_rotation 1";
+
+    private final boolean mIsPortrait;
+    private final UiDevice mUiDevice =
+            UiDevice.getInstance(InstrumentationRegistry.getInstrumentation());
+
+    ScreenOrientationRule(boolean isPortrait) {
+        mIsPortrait = isPortrait;
+    }
+
+    @Override
+    protected void starting(Description description) {
+        try {
+            mUiDevice.freezeRotation();
+            mUiDevice.executeShellCommand(mIsPortrait ? SET_PORTRAIT_MODE_CMD :
+                    SET_LANDSCAPE_MODE_CMD);
+        } catch (IOException e) {
+            throw new RuntimeException("Could not set screen orientation.", e);
+        } catch (RemoteException e) {
+            throw new RuntimeException("Could not freeze rotation.", e);
+        }
+    }
+
+    @Override
+    protected void finished(Description description) {
+        try {
+            mUiDevice.unfreezeRotation();
+        } catch (RemoteException e) {
+            throw new RuntimeException("Could not unfreeze screen rotation.", e);
+        }
+    }
+}