diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/SwitchImeWindowsFromGestureNavTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/SwitchImeWindowsFromGestureNavTest.kt
index 56ec80cd5caaa..a7a9fe289a4aa 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/SwitchImeWindowsFromGestureNavTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/SwitchImeWindowsFromGestureNavTest.kt
@@ -107,9 +107,13 @@ class SwitchImeWindowsFromGestureNavTest(private val testSpec: FlickerTestParame
     fun imeAppWindowVisibility() {
         testSpec.assertWm {
             isAppWindowVisible(imeTestApp.component)
+                .then()
+                .isAppSnapshotStartingWindowVisibleFor(testApp.component, isOptional = true)
                 .then()
                 .isAppWindowVisible(testApp.component)
                 .then()
+                .isAppSnapshotStartingWindowVisibleFor(imeTestApp.component, isOptional = true)
+                .then()
                 .isAppWindowVisible(imeTestApp.component)
         }
     }