diff --git a/tests/InputMethodStressTest/Android.bp b/tests/InputMethodStressTest/Android.bp
index 131611d..0ad3876 100644
--- a/tests/InputMethodStressTest/Android.bp
+++ b/tests/InputMethodStressTest/Android.bp
@@ -25,6 +25,7 @@
         "androidx.test.uiautomator_uiautomator",
         "compatibility-device-util-axt",
         "platform-test-annotations",
+        "platform-test-rules",
         "truth-prebuilt",
     ],
     test_suites: [
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java
index 33cad78..f0f78740 100644
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/AutoShowTest.java
@@ -29,6 +29,7 @@
 import android.content.Intent;
 import android.os.Bundle;
 import android.platform.test.annotations.RootPermissionTest;
+import android.platform.test.rule.UnlockScreenRule;
 import android.widget.EditText;
 import android.widget.LinearLayout;
 
@@ -36,6 +37,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.platform.app.InstrumentationRegistry;
 
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -43,6 +45,9 @@
 @RunWith(AndroidJUnit4.class)
 public final class AutoShowTest {
 
+    @Rule
+    public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
+
     @Test
     public void autoShow() {
         Instrumentation instrumentation = InstrumentationRegistry.getInstrumentation();
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeOpenCloseStressTest.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeOpenCloseStressTest.java
index b3c63c8..4be07fb 100644
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeOpenCloseStressTest.java
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeOpenCloseStressTest.java
@@ -29,6 +29,7 @@
 import android.os.Bundle;
 import android.os.SystemClock;
 import android.platform.test.annotations.RootPermissionTest;
+import android.platform.test.rule.UnlockScreenRule;
 import android.view.WindowInsets;
 import android.view.WindowInsetsAnimation;
 import android.view.inputmethod.InputMethodManager;
@@ -39,6 +40,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.platform.app.InstrumentationRegistry;
 
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -50,6 +52,9 @@
 
     private static final int NUM_TEST_ITERATIONS = 10;
 
+    @Rule
+    public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
+
     @Test
     public void test() {
         Instrumentation instrumentation = InstrumentationRegistry.getInstrumentation();
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java
index e0215f55..356c470 100644
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java
@@ -30,6 +30,7 @@
 import android.content.Intent;
 import android.graphics.drawable.Icon;
 import android.platform.test.annotations.RootPermissionTest;
+import android.platform.test.rule.UnlockScreenRule;
 import android.provider.Settings;
 import android.view.KeyEvent;
 
@@ -42,6 +43,7 @@
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -70,6 +72,9 @@
     private static final BySelector REPLY_SEND_BUTTON_SELECTOR =
             By.res("com.android.systemui", "remote_input_send");
 
+    @Rule
+    public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
+
     private Context mContext;
     private NotificationManager mNotificationManager;
     private UiDevice mUiDevice;