diff --git a/tests/autofillservice/src/android/autofillservice/cts/dialog/LoginActivityTest.java b/tests/autofillservice/src/android/autofillservice/cts/dialog/LoginActivityTest.java
index 990a97b..0f6ce62 100644
--- a/tests/autofillservice/src/android/autofillservice/cts/dialog/LoginActivityTest.java
+++ b/tests/autofillservice/src/android/autofillservice/cts/dialog/LoginActivityTest.java
@@ -26,6 +26,7 @@
 import static android.autofillservice.cts.testcore.Helper.assertFillEventForDatasetSelected;
 import static android.autofillservice.cts.testcore.Helper.assertFillEventForDatasetShown;
 import static android.autofillservice.cts.testcore.Helper.assertHasFlags;
+import static android.autofillservice.cts.testcore.Helper.assertMockImeStatus;
 import static android.autofillservice.cts.testcore.Helper.assertNoDeprecatedClientState;
 import static android.autofillservice.cts.testcore.Helper.assertNoFlags;
 import static android.autofillservice.cts.testcore.Helper.enableFillDialogFeature;
@@ -113,7 +114,7 @@
         mUiBot.waitForIdleSync();
 
         // Verify IME is shown
-        assertThat(isImeShowing(activity.getRootWindowInsets())).isTrue();
+        assertMockImeStatus(activity.getRootWindowInsets(), true);
     }
 
     @Test
diff --git a/tests/autofillservice/src/android/autofillservice/cts/testcore/Helper.java b/tests/autofillservice/src/android/autofillservice/cts/testcore/Helper.java
index 687c219..9502925 100644
--- a/tests/autofillservice/src/android/autofillservice/cts/testcore/Helper.java
+++ b/tests/autofillservice/src/android/autofillservice/cts/testcore/Helper.java
@@ -1693,6 +1693,20 @@
         return false;
     }
 
+    /**
+     * Asserts whether mock IME is showing
+     */
+    public static void assertMockImeStatus(WindowInsets rootWindowInsets,
+            boolean expectedImeShow) throws Exception {
+        Timeouts.MOCK_IME_TIMEOUT.run("assertMockImeStatus(" + expectedImeShow + ")",
+                () -> {
+                    final boolean actual = isImeShowing(rootWindowInsets);
+                    Log.v(TAG, "assertMockImeStatus(): expected=" + expectedImeShow + ", actual="
+                            + actual);
+                    return actual == expectedImeShow ? "expected" : null;
+                });
+    }
+
     private Helper() {
         throw new UnsupportedOperationException("contain static methods only");
     }
diff --git a/tests/autofillservice/src/android/autofillservice/cts/testcore/Timeouts.java b/tests/autofillservice/src/android/autofillservice/cts/testcore/Timeouts.java
index 7d4e140..ac259df 100644
--- a/tests/autofillservice/src/android/autofillservice/cts/testcore/Timeouts.java
+++ b/tests/autofillservice/src/android/autofillservice/cts/testcore/Timeouts.java
@@ -131,6 +131,12 @@
             "UI_SCREEN_ORIENTATION_TIMEOUT", ONE_TIMEOUT_TO_RULE_THEN_ALL_MS, 2F,
             ONE_TIMEOUT_TO_RULE_THEN_ALL_MS);
 
+    /**
+     * Timeout for changing the mock ime status.
+     */
+    public static final Timeout MOCK_IME_TIMEOUT = new Timeout(
+            "MOCK_IME_TIMEOUT", MOCK_IME_TIMEOUT_MS, 2F, MOCK_IME_TIMEOUT_MS);
+
     private Timeouts() {
         throw new UnsupportedOperationException("contain static methods only");
     }