diff --git a/tests/InputMethodStressTest/Android.bp b/tests/InputMethodStressTest/Android.bp
index 0ad3876..27640a5 100644
--- a/tests/InputMethodStressTest/Android.bp
+++ b/tests/InputMethodStressTest/Android.bp
@@ -32,5 +32,8 @@
         "general-tests",
         "vts",
     ],
-    sdk_version: "31",
+    data: [
+        ":SimpleTestIme",
+    ],
+    sdk_version: "current",
 }
diff --git a/tests/InputMethodStressTest/AndroidManifest.xml b/tests/InputMethodStressTest/AndroidManifest.xml
index 2d183bc..62eee02 100644
--- a/tests/InputMethodStressTest/AndroidManifest.xml
+++ b/tests/InputMethodStressTest/AndroidManifest.xml
@@ -17,7 +17,7 @@
 
 <manifest xmlns:android="http://schemas.android.com/apk/res/android"
           package="com.android.inputmethod.stresstest">
-
+    <uses-permission android:name="android.permission.POST_NOTIFICATIONS"/>
     <application>
         <activity android:name=".ImeStressTestUtil$TestActivity"
                   android:configChanges="orientation|screenSize"/>
diff --git a/tests/InputMethodStressTest/AndroidTest.xml b/tests/InputMethodStressTest/AndroidTest.xml
index 9ac4135..bedf099 100644
--- a/tests/InputMethodStressTest/AndroidTest.xml
+++ b/tests/InputMethodStressTest/AndroidTest.xml
@@ -25,6 +25,7 @@
 
     <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
         <option name="cleanup-apks" value="true" />
+        <option name="test-file-name" value="SimpleTestIme.apk" />
         <option name="test-file-name" value="InputMethodStressTest.apk" />
     </target_preparer>
 
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java
index 9c70e6e..3d257b2 100644
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java
@@ -61,14 +61,10 @@
 @RunWith(Parameterized.class)
 public final class AutoShowTest {
 
-    @Rule(order = 0) public DisableLockScreenRule mDisableLockScreenRule =
-            new DisableLockScreenRule();
-    @Rule(order = 1) public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
-    @Rule(order = 2) public ScreenOrientationRule mScreenOrientationRule =
-            new ScreenOrientationRule(true /* isPortrait */);
-    @Rule(order = 3) public PressHomeBeforeTestRule mPressHomeBeforeTestRule =
-            new PressHomeBeforeTestRule();
-    @Rule(order = 4) public ScreenCaptureRule mScreenCaptureRule =
+    @Rule(order = 0) public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
+    @Rule(order = 1) public ImeStressTestRule mImeStressTestRule =
+        new ImeStressTestRule(true /* useSimpleTestIme */);
+    @Rule(order = 2) public ScreenCaptureRule mScreenCaptureRule =
             new ScreenCaptureRule("/sdcard/InputMethodStressTest");
     @Parameterized.Parameters(
             name = "windowFocusFlags={0}, softInputVisibility={1}, softInputAdjustment={2}")
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/DisableLockScreenRule.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/DisableLockScreenRule.java
deleted file mode 100644
index d95decf..0000000
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/DisableLockScreenRule.java
+++ /dev/null
@@ -1,53 +0,0 @@
-/*
- * 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 9d4aefb..7632ab0 100644
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeOpenCloseStressTest.java
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeOpenCloseStressTest.java
@@ -68,14 +68,10 @@
     private static final String TAG = "ImeOpenCloseStressTest";
     private static final int NUM_TEST_ITERATIONS = 10;
 
-    @Rule(order = 0) public DisableLockScreenRule mDisableLockScreenRule =
-            new DisableLockScreenRule();
-    @Rule(order = 1) public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
-    @Rule(order = 2) public ScreenOrientationRule mScreenOrientationRule =
-            new ScreenOrientationRule(true /* isPortrait */);
-    @Rule(order = 3) public PressHomeBeforeTestRule mPressHomeBeforeTestRule =
-            new PressHomeBeforeTestRule();
-    @Rule(order = 4) public ScreenCaptureRule mScreenCaptureRule =
+    @Rule(order = 0) public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
+    @Rule(order = 1) public ImeStressTestRule mImeStressTestRule =
+            new ImeStressTestRule(true /* useSimpleTestIme */);
+    @Rule(order = 2) public ScreenCaptureRule mScreenCaptureRule =
             new ScreenCaptureRule("/sdcard/InputMethodStressTest");
 
     private final Instrumentation mInstrumentation;
@@ -499,8 +495,6 @@
 
     @Test
     public void testRotateScreenWithKeyboardOn() throws Exception {
-        // TODO(b/256739702): Keyboard disappears after rotating screen to landscape mode if
-        // android:configChanges="orientation|screenSize" is not set
         Intent intent =
                 createIntent(
                         mWindowFocusFlags,
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeStressTestRule.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeStressTestRule.java
new file mode 100644
index 0000000..12104b2
--- /dev/null
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeStressTestRule.java
@@ -0,0 +1,153 @@
+/*
+ * Copyright (C) 2023 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.app.Instrumentation;
+import android.os.RemoteException;
+import android.support.test.uiautomator.UiDevice;
+
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.junit.rules.TestWatcher;
+import org.junit.runner.Description;
+
+import java.io.IOException;
+
+/**
+ * Do setup and cleanup for Ime stress tests, including disabling lock and auto-rotate screen,
+ * pressing home and enabling a simple test Ime during the tests.
+ */
+public class ImeStressTestRule 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 static final String SET_PORTRAIT_MODE_COMMAND = "settings put system user_rotation 0";
+    private static final String SET_LANDSCAPE_MODE_COMMAND = "settings put system user_rotation 1";
+    private static final String SIMPLE_IME_ID =
+            "com.android.apps.inputmethod.simpleime/.SimpleInputMethodService";
+    private static final String ENABLE_IME_COMMAND = "ime enable " + SIMPLE_IME_ID;
+    private static final String SET_IME_COMMAND = "ime set " + SIMPLE_IME_ID;
+    private static final String DISABLE_IME_COMMAND = "ime disable " + SIMPLE_IME_ID;
+    private static final String RESET_IME_COMMAND = "ime reset";
+
+    @NonNull private final Instrumentation mInstrumentation;
+    @NonNull private final UiDevice mUiDevice;
+    // Whether the screen orientation is set to portrait.
+    private boolean mIsPortrait;
+    // Whether to use a simple test Ime or system default Ime for test.
+    private final boolean mUseSimpleTestIme;
+
+    public ImeStressTestRule(boolean useSimpleTestIme) {
+        mInstrumentation = InstrumentationRegistry.getInstrumentation();
+        mUiDevice = UiDevice.getInstance(mInstrumentation);
+        // Default is portrait mode
+        mIsPortrait = true;
+        mUseSimpleTestIme = useSimpleTestIme;
+    }
+
+    public void setIsPortrait(boolean isPortrait) {
+        mIsPortrait = isPortrait;
+    }
+
+    @Override
+    protected void starting(Description description) {
+        disableLockScreen();
+        setOrientation();
+        mUiDevice.pressHome();
+        if (mUseSimpleTestIme) {
+            enableSimpleIme();
+        } else {
+            resetImeToDefault();
+        }
+
+        mInstrumentation.waitForIdleSync();
+    }
+
+    @Override
+    protected void finished(Description description) {
+        if (mUseSimpleTestIme) {
+            disableSimpleIme();
+        }
+        unfreezeRotation();
+        restoreLockScreen();
+    }
+
+    private void disableLockScreen() {
+        try {
+            executeShellCommand(LOCK_SCREEN_OFF_COMMAND);
+        } catch (IOException e) {
+            throw new RuntimeException("Could not disable lock screen.", e);
+        }
+    }
+
+    private void restoreLockScreen() {
+        try {
+            executeShellCommand(LOCK_SCREEN_ON_COMMAND);
+        } catch (IOException e) {
+            throw new RuntimeException("Could not enable lock screen.", e);
+        }
+    }
+
+    private void setOrientation() {
+        try {
+            mUiDevice.freezeRotation();
+            executeShellCommand(
+                    mIsPortrait ? SET_PORTRAIT_MODE_COMMAND : SET_LANDSCAPE_MODE_COMMAND);
+        } catch (IOException e) {
+            throw new RuntimeException("Could not set screen orientation.", e);
+        } catch (RemoteException e) {
+            throw new RuntimeException("Could not freeze rotation.", e);
+        }
+    }
+
+    private void unfreezeRotation() {
+        try {
+            mUiDevice.unfreezeRotation();
+        } catch (RemoteException e) {
+            throw new RuntimeException("Could not unfreeze screen rotation.", e);
+        }
+    }
+
+    private void enableSimpleIme() {
+        try {
+            executeShellCommand(ENABLE_IME_COMMAND);
+            executeShellCommand(SET_IME_COMMAND);
+        } catch (IOException e) {
+            throw new RuntimeException("Could not enable SimpleTestIme.", e);
+        }
+    }
+
+    private void disableSimpleIme() {
+        try {
+            executeShellCommand(DISABLE_IME_COMMAND);
+        } catch (IOException e) {
+            throw new RuntimeException("Could not disable SimpleTestIme.", e);
+        }
+    }
+
+    private void resetImeToDefault() {
+        try {
+            executeShellCommand(RESET_IME_COMMAND);
+        } catch (IOException e) {
+            throw new RuntimeException("Could not reset Ime to default.", e);
+        }
+    }
+
+    private @NonNull String executeShellCommand(@NonNull String cmd) throws IOException {
+        return mUiDevice.executeShellCommand(cmd);
+    }
+}
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java
index d2708ad..f4a04a1 100644
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java
@@ -77,11 +77,10 @@
     private static final BySelector REPLY_SEND_BUTTON_SELECTOR =
             By.res("com.android.systemui", "remote_input_send").enabled(true);
 
-    @Rule
-    public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
-
-    @Rule
-    public ScreenCaptureRule mScreenCaptureRule =
+    @Rule(order = 0) public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
+    @Rule(order = 1) public ImeStressTestRule mImeStressTestRule =
+            new ImeStressTestRule(true /* useSimpleTestIme */);
+    @Rule(order = 2) public ScreenCaptureRule mScreenCaptureRule =
             new ScreenCaptureRule("/sdcard/InputMethodStressTest");
 
     private Context mContext;
@@ -141,7 +140,8 @@
 
         // Post inline reply notification.
         PendingIntent pendingIntent = PendingIntent.getBroadcast(
-                mContext, REPLY_REQUEST_CODE, new Intent().setAction(ACTION_REPLY),
+                mContext, REPLY_REQUEST_CODE,
+                new Intent().setAction(ACTION_REPLY).setClass(mContext, NotificationTest.class),
                 PendingIntent.FLAG_UPDATE_CURRENT | PendingIntent.FLAG_MUTABLE);
         RemoteInput remoteInput = new RemoteInput.Builder(REPLY_INPUT_KEY)
                 .setLabel(REPLY_INPUT_LABEL)
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/PressHomeBeforeTestRule.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/PressHomeBeforeTestRule.java
deleted file mode 100644
index 6586f63..0000000
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/PressHomeBeforeTestRule.java
+++ /dev/null
@@ -1,34 +0,0 @@
-/*
- * Copyright (C) 2023 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;
-
-/** This rule will press home before a test case. */
-public class PressHomeBeforeTestRule extends TestWatcher {
-    private final UiDevice mUiDevice =
-            UiDevice.getInstance(InstrumentationRegistry.getInstrumentation());
-
-    @Override
-    protected void starting(Description description) {
-        mUiDevice.pressHome();
-    }
-}
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ScreenOrientationRule.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ScreenOrientationRule.java
deleted file mode 100644
index bc3b1ef..0000000
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ScreenOrientationRule.java
+++ /dev/null
@@ -1,66 +0,0 @@
-/*
- * 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);
-        }
-    }
-}