diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/CommonAssertions.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/CommonAssertions.kt
index bb03237ca4c60..bd7e8d7683d18 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/CommonAssertions.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/CommonAssertions.kt
@@ -48,11 +48,12 @@ fun WmAssertion.navBarWindowIsAlwaysVisible(
 }
 
 fun WmAssertion.visibleWindowsShownMoreThanOneConsecutiveEntry(
+    ignoreWindows: List<String> = emptyList(),
     bugId: Int = 0,
     enabled: Boolean = bugId == 0
 ) {
-    all("visibleWindowShownMoreThanOneConsecutiveEntry", bugId, enabled) {
-        this.visibleWindowsShownMoreThanOneConsecutiveEntry()
+    all("visibleWindowsShownMoreThanOneConsecutiveEntry", bugId, enabled) {
+        this.visibleWindowsShownMoreThanOneConsecutiveEntry(ignoreWindows)
     }
 }
 
@@ -204,11 +205,12 @@ fun LayersAssertion.statusBarLayerRotatesScales(
 }
 
 fun LayersAssertion.visibleLayersShownMoreThanOneConsecutiveEntry(
+    ignoreLayers: List<String> = emptyList(),
     bugId: Int = 0,
     enabled: Boolean = bugId == 0
 ) {
     all("visibleLayersShownMoreThanOneConsecutiveEntry", bugId, enabled) {
-        this.visibleLayersShownMoreThanOneConsecutiveEntry()
+        this.visibleLayersShownMoreThanOneConsecutiveEntry(ignoreLayers)
     }
 }
 
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeAppHelper.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeAppHelper.kt
index f4de36e99952c..3565ca40474ee 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeAppHelper.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeAppHelper.kt
@@ -22,6 +22,7 @@ import android.support.test.launcherhelper.LauncherStrategyFactory
 import androidx.test.uiautomator.By
 import androidx.test.uiautomator.UiDevice
 import androidx.test.uiautomator.Until
+import com.android.server.wm.traces.parser.windowmanager.WindowManagerStateHelper
 import org.junit.Assert
 
 open class ImeAppHelper(
@@ -38,7 +39,7 @@ open class ImeAppHelper(
         Assert.assertNotNull("Text field not found, this usually happens when the device " +
                 "was left in an unknown state (e.g. in split screen)", editText)
         editText.click()
-        if (!device.waitForIME()) {
+        if (!WindowManagerStateHelper().waitImeWindowShown()) {
             Assert.fail("IME did not appear")
         }
     }
@@ -46,6 +47,8 @@ open class ImeAppHelper(
     open fun closeIME(device: UiDevice) {
         device.pressBack()
         // Using only the AccessibilityInfo it is not possible to identify if the IME is active
-        device.waitForIdle(1000)
+        if (!WindowManagerStateHelper().waitImeWindowGone()) {
+            Assert.fail("IME did not close")
+        }
     }
 }
\ No newline at end of file
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeAutoOpenWindowToAppTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeAutoOpenWindowToAppTest.kt
index a73264d304f81..b341e621d9ed9 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeAutoOpenWindowToAppTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeAutoOpenWindowToAppTest.kt
@@ -27,6 +27,8 @@ import com.android.server.wm.flicker.helpers.ImeAppAutoFocusHelper
 import com.android.server.wm.flicker.helpers.buildTestTag
 import com.android.server.wm.flicker.helpers.setRotation
 import com.android.server.wm.flicker.helpers.wakeUpAndGoToHomeScreen
+import com.android.server.wm.flicker.visibleWindowsShownMoreThanOneConsecutiveEntry
+import com.android.server.wm.flicker.visibleLayersShownMoreThanOneConsecutiveEntry
 import com.android.server.wm.flicker.navBarLayerIsAlwaysVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
 import com.android.server.wm.flicker.navBarWindowIsAlwaysVisible
@@ -66,8 +68,10 @@ class CloseImeAutoOpenWindowToAppTest(
                     withTag { buildTestTag("imeToAppAutoOpen", testApp, configuration) }
                     repeat { configuration.repetitions }
                     setup {
-                        eachRun {
+                        test {
                             device.wakeUpAndGoToHomeScreen()
+                        }
+                        eachRun {
                             this.setRotation(configuration.startRotation)
                             testApp.open()
                             testApp.openIME(device)
@@ -80,14 +84,15 @@ class CloseImeAutoOpenWindowToAppTest(
                         }
                     }
                     transitions {
-                        device.pressBack()
-                        device.waitForIdle()
+                        testApp.closeIME(device)
                     }
                     assertions {
                         windowManagerTrace {
                             navBarWindowIsAlwaysVisible()
                             statusBarWindowIsAlwaysVisible()
-                            imeAppWindowIsAlwaysVisible(testApp, bugId = 141458352)
+                            visibleWindowsShownMoreThanOneConsecutiveEntry(listOf("InputMethod"))
+
+                            imeAppWindowIsAlwaysVisible(testApp)
                         }
 
                         layersTrace {
@@ -96,8 +101,10 @@ class CloseImeAutoOpenWindowToAppTest(
                             noUncoveredRegions(configuration.startRotation)
                             navBarLayerRotatesAndScales(configuration.startRotation)
                             statusBarLayerRotatesScales(configuration.startRotation)
+                            visibleLayersShownMoreThanOneConsecutiveEntry()
+
                             imeLayerBecomesInvisible(bugId = 141458352)
-                            imeAppLayerIsAlwaysVisible(testApp, bugId = 141458352)
+                            imeAppLayerIsAlwaysVisible(testApp)
                         }
                     }
                 }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeAutoOpenWindowToHomeTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeAutoOpenWindowToHomeTest.kt
index 6cc2e2236127b..51a4ca86681b2 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeAutoOpenWindowToHomeTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeAutoOpenWindowToHomeTest.kt
@@ -27,6 +27,8 @@ import com.android.server.wm.flicker.helpers.ImeAppAutoFocusHelper
 import com.android.server.wm.flicker.helpers.buildTestTag
 import com.android.server.wm.flicker.helpers.setRotation
 import com.android.server.wm.flicker.helpers.wakeUpAndGoToHomeScreen
+import com.android.server.wm.flicker.visibleWindowsShownMoreThanOneConsecutiveEntry
+import com.android.server.wm.flicker.visibleLayersShownMoreThanOneConsecutiveEntry
 import com.android.server.wm.flicker.navBarLayerIsAlwaysVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
 import com.android.server.wm.flicker.navBarWindowIsAlwaysVisible
@@ -68,8 +70,10 @@ class CloseImeAutoOpenWindowToHomeTest(
                     }
                     repeat { configuration.repetitions }
                     setup {
-                        eachRun {
+                        test {
                             device.wakeUpAndGoToHomeScreen()
+                        }
+                        eachRun {
                             this.setRotation(configuration.startRotation)
                             testApp.open()
                             testApp.openIME(device)
@@ -89,6 +93,8 @@ class CloseImeAutoOpenWindowToHomeTest(
                         windowManagerTrace {
                             navBarWindowIsAlwaysVisible()
                             statusBarWindowIsAlwaysVisible()
+                            visibleWindowsShownMoreThanOneConsecutiveEntry(listOf(IME_WINDOW_TITLE))
+
                             imeWindowBecomesInvisible(bugId = 141458352)
                             imeAppWindowBecomesInvisible(testApp, bugId = 157449248)
                         }
@@ -100,10 +106,14 @@ class CloseImeAutoOpenWindowToHomeTest(
                                 Surface.ROTATION_0, bugId = 140855415)
                             statusBarLayerRotatesScales(configuration.startRotation,
                                 Surface.ROTATION_0)
-                            navBarLayerIsAlwaysVisible(enabled = false)
-                            statusBarLayerIsAlwaysVisible(enabled = false)
+                            navBarLayerIsAlwaysVisible(
+                                    enabled = Surface.ROTATION_0 == configuration.startRotation)
+                            statusBarLayerIsAlwaysVisible(
+                                    enabled = Surface.ROTATION_0 == configuration.startRotation)
+                            visibleLayersShownMoreThanOneConsecutiveEntry(listOf(IME_WINDOW_TITLE))
+
                             imeLayerBecomesInvisible(bugId = 141458352)
-                            imeAppLayerBecomesInvisible(testApp, bugId = 153739621)
+                            imeAppLayerBecomesInvisible(testApp)
                         }
                     }
                 }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeWindowToAppTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeWindowToAppTest.kt
index 136cf86037a1f..663d45607f934 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeWindowToAppTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeWindowToAppTest.kt
@@ -31,6 +31,8 @@ import com.android.server.wm.flicker.navBarLayerIsAlwaysVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
 import com.android.server.wm.flicker.navBarWindowIsAlwaysVisible
 import com.android.server.wm.flicker.noUncoveredRegions
+import com.android.server.wm.flicker.visibleWindowsShownMoreThanOneConsecutiveEntry
+import com.android.server.wm.flicker.visibleLayersShownMoreThanOneConsecutiveEntry
 import com.android.server.wm.flicker.repetitions
 import com.android.server.wm.flicker.startRotation
 import com.android.server.wm.flicker.statusBarLayerIsAlwaysVisible
@@ -65,8 +67,10 @@ class CloseImeWindowToAppTest(
                     withTestName { buildTestTag("imeToApp", testApp, configuration) }
                     repeat { configuration.repetitions }
                     setup {
-                        eachRun {
+                        test {
                             device.wakeUpAndGoToHomeScreen()
+                        }
+                        eachRun {
                             this.setRotation(configuration.startRotation)
                             testApp.open()
                             testApp.openIME(device)
@@ -79,13 +83,14 @@ class CloseImeWindowToAppTest(
                         }
                     }
                     transitions {
-                        device.pressBack()
-                        device.waitForIdle()
+                        testApp.closeIME(device)
                     }
                     assertions {
                         windowManagerTrace {
                             navBarWindowIsAlwaysVisible()
                             statusBarWindowIsAlwaysVisible()
+                            visibleWindowsShownMoreThanOneConsecutiveEntry(listOf("InputMethod"))
+
                             imeAppWindowIsAlwaysVisible(testApp)
                         }
 
@@ -95,6 +100,8 @@ class CloseImeWindowToAppTest(
                             noUncoveredRegions(configuration.startRotation)
                             navBarLayerRotatesAndScales(configuration.startRotation)
                             statusBarLayerRotatesScales(configuration.startRotation)
+                            visibleLayersShownMoreThanOneConsecutiveEntry()
+
                             imeLayerBecomesInvisible(enabled = false)
                             imeAppLayerIsAlwaysVisible(testApp)
                         }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeWindowToHomeTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeWindowToHomeTest.kt
index 8d9881ec2063c..bfe5264ed3587 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeWindowToHomeTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeWindowToHomeTest.kt
@@ -33,6 +33,8 @@ import com.android.server.wm.flicker.navBarLayerIsAlwaysVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
 import com.android.server.wm.flicker.navBarWindowIsAlwaysVisible
 import com.android.server.wm.flicker.noUncoveredRegions
+import com.android.server.wm.flicker.visibleWindowsShownMoreThanOneConsecutiveEntry
+import com.android.server.wm.flicker.visibleLayersShownMoreThanOneConsecutiveEntry
 import com.android.server.wm.flicker.repetitions
 import com.android.server.wm.flicker.startRotation
 import com.android.server.wm.flicker.statusBarLayerIsAlwaysVisible
@@ -95,6 +97,8 @@ class CloseImeWindowToHomeTest(
                         windowManagerTrace {
                             navBarWindowIsAlwaysVisible()
                             statusBarWindowIsAlwaysVisible()
+                            visibleWindowsShownMoreThanOneConsecutiveEntry(listOf(IME_WINDOW_TITLE))
+
                             imeWindowBecomesInvisible()
                             imeAppWindowBecomesInvisible(testApp)
                         }
@@ -106,10 +110,14 @@ class CloseImeWindowToHomeTest(
                                 Surface.ROTATION_0, bugId = 140855415)
                             statusBarLayerRotatesScales(configuration.startRotation,
                                 Surface.ROTATION_0)
-                            navBarLayerIsAlwaysVisible(enabled = false)
-                            statusBarLayerIsAlwaysVisible(enabled = false)
+                            navBarLayerIsAlwaysVisible(
+                                    enabled = Surface.ROTATION_0 == configuration.startRotation)
+                            statusBarLayerIsAlwaysVisible(
+                                    enabled = Surface.ROTATION_0 == configuration.startRotation)
+                            visibleLayersShownMoreThanOneConsecutiveEntry(listOf(IME_WINDOW_TITLE))
+
                             imeLayerBecomesInvisible(bugId = 153739621)
-                            imeAppLayerBecomesInvisible(testApp, bugId = 153739621)
+                            imeAppLayerBecomesInvisible(testApp)
                         }
                     }
                 }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowTest.kt
index 5767a94025e9f..64a8f304298b1 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowTest.kt
@@ -30,6 +30,8 @@ import com.android.server.wm.flicker.helpers.wakeUpAndGoToHomeScreen
 import com.android.server.wm.flicker.navBarLayerIsAlwaysVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
 import com.android.server.wm.flicker.navBarWindowIsAlwaysVisible
+import com.android.server.wm.flicker.visibleWindowsShownMoreThanOneConsecutiveEntry
+import com.android.server.wm.flicker.visibleLayersShownMoreThanOneConsecutiveEntry
 import com.android.server.wm.flicker.noUncoveredRegions
 import com.android.server.wm.flicker.repetitions
 import com.android.server.wm.flicker.startRotation
@@ -89,6 +91,7 @@ class OpenImeWindowTest(
                         windowManagerTrace {
                             navBarWindowIsAlwaysVisible()
                             statusBarWindowIsAlwaysVisible()
+                            visibleWindowsShownMoreThanOneConsecutiveEntry()
 
                             all("imeWindowBecomesVisible") {
                                 this.skipUntilFirstAssertion()
@@ -104,6 +107,7 @@ class OpenImeWindowTest(
                             noUncoveredRegions(configuration.startRotation)
                             navBarLayerRotatesAndScales(configuration.startRotation)
                             statusBarLayerRotatesScales(configuration.startRotation)
+                            visibleLayersShownMoreThanOneConsecutiveEntry()
 
                             imeLayerBecomesVisible()
                         }