diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/bubble/BaseBubbleScreen.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/bubble/BaseBubbleScreen.kt
index 278ba9b0f4db3..5b073038059c6 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/bubble/BaseBubbleScreen.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/bubble/BaseBubbleScreen.kt
@@ -48,7 +48,7 @@ abstract class BaseBubbleScreen(protected val testSpec: FlickerTestParameter) {
             ServiceManager.getService(Context.NOTIFICATION_SERVICE))
 
     protected val uid = context.packageManager.getApplicationInfo(
-            testApp.component.packageName, 0).uid
+            testApp.`package`, 0).uid
 
     protected abstract val transition: FlickerBuilder.() -> Unit
 
@@ -59,7 +59,7 @@ abstract class BaseBubbleScreen(protected val testSpec: FlickerTestParameter) {
         return {
             setup {
                 test {
-                    notifyManager.setBubblesAllowed(testApp.component.packageName,
+                    notifyManager.setBubblesAllowed(testApp.`package`,
                             uid, NotificationManager.BUBBLE_PREFERENCE_ALL)
                     testApp.launchViaIntent(wmHelper)
                     waitAndGetAddBubbleBtn()
@@ -69,7 +69,7 @@ abstract class BaseBubbleScreen(protected val testSpec: FlickerTestParameter) {
 
             teardown {
                 test {
-                    notifyManager.setBubblesAllowed(testApp.component.packageName,
+                    notifyManager.setBubblesAllowed(testApp.`package`,
                         uid, NotificationManager.BUBBLE_PREFERENCE_NONE)
                     testApp.exit()
                 }
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/bubble/LaunchBubbleFromLockScreen.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/bubble/LaunchBubbleFromLockScreen.kt
index 27d65648dbb7c..e2b19ea22b3be 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/bubble/LaunchBubbleFromLockScreen.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/bubble/LaunchBubbleFromLockScreen.kt
@@ -16,10 +16,10 @@
 
 package com.android.wm.shell.flicker.bubble
 
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.view.WindowInsets
 import android.view.WindowManager
-import android.platform.test.annotations.FlakyTest
 import androidx.test.filters.RequiresDevice
 import androidx.test.uiautomator.By
 import androidx.test.uiautomator.Until
@@ -29,8 +29,8 @@ import com.android.server.wm.flicker.annotation.Group4
 import com.android.server.wm.flicker.dsl.FlickerBuilder
 import com.android.server.wm.flicker.helpers.isShellTransitionsEnabled
 import org.junit.Assume
-import org.junit.runner.RunWith
 import org.junit.Test
+import org.junit.runner.RunWith
 import org.junit.runners.Parameterized
 
 /**
@@ -54,20 +54,21 @@ class LaunchBubbleFromLockScreen(testSpec: FlickerTestParameter) : BaseBubbleScr
                     val addBubbleBtn = waitAndGetAddBubbleBtn()
                     addBubbleBtn?.click() ?: error("Bubble widget not found")
                     device.sleep()
-                    wmHelper.waitFor("noAppWindowsOnTop") {
-                        it.wmState.topVisibleAppWindow.isEmpty()
-                    }
+                    wmHelper.StateSyncBuilder()
+                        .withoutTopVisibleAppWindows()
+                        .waitForAndVerify()
                     device.wakeUp()
                 }
             }
             transitions {
                 // Swipe & wait for the notification shade to expand so all can be seen
                 val wm = context.getSystemService(WindowManager::class.java)
-                val metricInsets = wm.getCurrentWindowMetrics().windowInsets
+                    ?: error("Unable to obtain WM service")
+                val metricInsets = wm.currentWindowMetrics.windowInsets
                 val insets = metricInsets.getInsetsIgnoringVisibility(
                         WindowInsets.Type.statusBars()
                         or WindowInsets.Type.displayCutout())
-                device.swipe(100, insets.top + 100, 100, device.getDisplayHeight() / 2, 4)
+                device.swipe(100, insets.top + 100, 100, device.displayHeight / 2, 4)
                 device.waitForIdle(2000)
                 instrumentation.uiAutomation.syncInputTransactions()
 
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/AppPairsHelper.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/AppPairsHelper.kt
index 41cd31aabf05c..f4305ed198244 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/AppPairsHelper.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/AppPairsHelper.kt
@@ -17,44 +17,10 @@
 package com.android.wm.shell.flicker.helpers
 
 import android.app.Instrumentation
-import com.android.server.wm.flicker.Flicker
-import com.android.server.wm.flicker.helpers.WindowUtils
 import com.android.server.wm.traces.common.FlickerComponentName
-import com.android.server.wm.traces.common.region.Region
 
 class AppPairsHelper(
     instrumentation: Instrumentation,
     activityLabel: String,
     component: FlickerComponentName
-) : BaseAppHelper(instrumentation, activityLabel, component) {
-    fun getPrimaryBounds(dividerBounds: Region): Region {
-        val primaryAppBounds = Region.from(0, 0, dividerBounds.bounds.right,
-                dividerBounds.bounds.bottom + WindowUtils.dockedStackDividerInset)
-        return primaryAppBounds
-    }
-
-    fun getSecondaryBounds(dividerBounds: Region): Region {
-        val displayBounds = WindowUtils.displayBounds
-        val secondaryAppBounds = Region.from(0,
-                dividerBounds.bounds.bottom - WindowUtils.dockedStackDividerInset,
-                displayBounds.right, displayBounds.bottom - WindowUtils.navigationBarFrameHeight)
-        return secondaryAppBounds
-    }
-
-    companion object {
-        const val TEST_REPETITIONS = 1
-        const val TIMEOUT_MS = 3_000L
-
-        fun Flicker.waitAppsShown(app1: SplitScreenHelper?, app2: SplitScreenHelper?) {
-            wmHelper.waitFor("primaryAndSecondaryAppsVisible") { dump ->
-                val primaryAppVisible = app1?.let {
-                    dump.wmState.isWindowSurfaceShown(app1.defaultWindowName)
-                } ?: false
-                val secondaryAppVisible = app2?.let {
-                    dump.wmState.isWindowSurfaceShown(app2.defaultWindowName)
-                } ?: false
-                primaryAppVisible && secondaryAppVisible
-            }
-        }
-    }
-}
+) : BaseAppHelper(instrumentation, activityLabel, component)
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/BaseAppHelper.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/BaseAppHelper.kt
index 3dd9e05729477..912ba67285f74 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/BaseAppHelper.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/BaseAppHelper.kt
@@ -40,7 +40,7 @@ abstract class BaseAppHelper(
     component,
     LauncherStrategyFactory.getInstance(instrumentation).launcherStrategy
 ) {
-    private val appSelector = By.pkg(component.packageName).depth(0)
+    private val appSelector = By.pkg(`package`).depth(0)
 
     protected val isTelevision: Boolean
         get() = context.packageManager.run {
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/ImeAppHelper.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/ImeAppHelper.kt
index cc5b9f9eb26d5..2e690de666f4c 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/ImeAppHelper.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/ImeAppHelper.kt
@@ -18,7 +18,6 @@ package com.android.wm.shell.flicker.helpers
 
 import android.app.Instrumentation
 import androidx.test.uiautomator.By
-import androidx.test.uiautomator.UiDevice
 import androidx.test.uiautomator.Until
 import com.android.server.wm.flicker.helpers.FIND_TIMEOUT
 import com.android.server.wm.traces.parser.toFlickerComponent
@@ -35,8 +34,7 @@ open class ImeAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
      *
      * @param wmHelper Helper used to wait for WindowManager states
      */
-    @JvmOverloads
-    open fun openIME(wmHelper: WindowManagerStateHelper? = null) {
+    open fun openIME(wmHelper: WindowManagerStateHelper) {
         if (!isTelevision) {
             val editText = uiDevice.wait(
                 Until.findObject(By.res(getPackage(), "plain_text_input")),
@@ -47,7 +45,9 @@ open class ImeAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
                     "was left in an unknown state (e.g. in split screen)"
             }
             editText.click()
-            waitAndAssertIMEShown(uiDevice, wmHelper)
+            wmHelper.StateSyncBuilder()
+                .withImeShown()
+                .waitForAndVerify()
         } else {
             // If we do the same thing as above - editText.click() - on TV, that's going to force TV
             // into the touch mode. We really don't want that.
@@ -55,36 +55,22 @@ open class ImeAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
         }
     }
 
-    protected fun waitAndAssertIMEShown(
-        device: UiDevice,
-        wmHelper: WindowManagerStateHelper? = null
-    ) {
-        if (wmHelper == null) {
-            device.waitForIdle()
-        } else {
-            wmHelper.waitImeShown()
-        }
-    }
-
     /**
      * Opens the IME and wait for it to be gone
      *
      * @param wmHelper Helper used to wait for WindowManager states
      */
-    @JvmOverloads
-    open fun closeIME(wmHelper: WindowManagerStateHelper? = null) {
+    open fun closeIME(wmHelper: WindowManagerStateHelper) {
         if (!isTelevision) {
             uiDevice.pressBack()
             // Using only the AccessibilityInfo it is not possible to identify if the IME is active
-            if (wmHelper == null) {
-                uiDevice.waitForIdle()
-            } else {
-                wmHelper.waitImeGone()
-            }
+            wmHelper.StateSyncBuilder()
+                .withImeGone()
+                .waitForAndVerify()
         } else {
             // While pressing the back button should close the IME on TV as well, it may also lead
             // to the app closing. So let's instead just ask the app to close the IME.
             launchViaIntent(action = Components.ImeActivity.ACTION_CLOSE_IME)
         }
     }
-}
\ No newline at end of file
+}
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/PipAppHelper.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/PipAppHelper.kt
index f73d191b19171..216445fe93563 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/PipAppHelper.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/PipAppHelper.kt
@@ -19,13 +19,13 @@ package com.android.wm.shell.flicker.helpers
 import android.app.Instrumentation
 import android.media.session.MediaController
 import android.media.session.MediaSessionManager
-import android.os.SystemClock
 import androidx.test.uiautomator.By
 import androidx.test.uiautomator.BySelector
 import androidx.test.uiautomator.Until
 import com.android.server.wm.flicker.helpers.FIND_TIMEOUT
 import com.android.server.wm.flicker.helpers.SYSTEMUI_PACKAGE
 import com.android.server.wm.traces.common.Rect
+import com.android.server.wm.traces.common.WindowManagerConditionsFactory
 import com.android.server.wm.traces.parser.toFlickerComponent
 import com.android.server.wm.traces.parser.windowmanager.WindowManagerStateHelper
 import com.android.wm.shell.flicker.pip.tv.closeTvPipWindow
@@ -43,11 +43,11 @@ class PipAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
 
     private val mediaController: MediaController?
         get() = mediaSessionManager.getActiveSessions(null).firstOrNull {
-            it.packageName == component.packageName
+            it.packageName == `package`
         }
 
     fun clickObject(resId: String) {
-        val selector = By.res(component.packageName, resId)
+        val selector = By.res(`package`, resId)
         val obj = uiDevice.findObject(selector) ?: error("Could not find `$resId` object")
 
         if (!isTelevision) {
@@ -71,8 +71,12 @@ class PipAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
     ) {
         launchViaIntentAndWaitShown(
             wmHelper, expectedWindowName, action, stringExtras,
-            waitConditions = arrayOf(WindowManagerStateHelper.pipShownCondition)
+            waitConditions = arrayOf(WindowManagerConditionsFactory.hasPipWindow())
         )
+
+        wmHelper.StateSyncBuilder()
+            .withPipShown()
+            .waitForAndVerify()
     }
 
     /**
@@ -95,12 +99,13 @@ class PipAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
         return false
     }
 
-    @JvmOverloads
-    fun clickEnterPipButton(wmHelper: WindowManagerStateHelper? = null) {
+    fun clickEnterPipButton(wmHelper: WindowManagerStateHelper) {
         clickObject(ENTER_PIP_BUTTON_ID)
 
         // Wait on WMHelper or simply wait for 3 seconds
-        wmHelper?.waitPipShown() ?: SystemClock.sleep(3_000)
+        wmHelper.StateSyncBuilder()
+            .withPipShown()
+            .waitForAndVerify()
         // when entering pip, the dismiss button is visible at the start. to ensure the pip
         // animation is complete, wait until the pip dismiss button is no longer visible.
         // b/176822698: dismiss-only state will be removed in the future
@@ -116,7 +121,7 @@ class PipAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
     }
 
     fun checkWithCustomActionsCheckbox() = uiDevice
-        .findObject(By.res(component.packageName, WITH_CUSTOM_ACTIONS_BUTTON_ID))
+        .findObject(By.res(`package`, WITH_CUSTOM_ACTIONS_BUTTON_ID))
         ?.takeIf { it.isCheckable }
         ?.apply { if (!isChecked) clickObject(WITH_CUSTOM_ACTIONS_BUTTON_ID) }
         ?: error("'With custom actions' checkbox not found")
@@ -166,8 +171,10 @@ class PipAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
         }
 
         // Wait for animation to complete.
-        wmHelper.waitPipGone()
-        wmHelper.waitForHomeActivityVisible()
+        wmHelper.StateSyncBuilder()
+            .withPipGone()
+            .withHomeActivityVisible()
+            .waitForAndVerify()
     }
 
     /**
@@ -183,8 +190,10 @@ class PipAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
             ?: error("PIP window expand button not found")
         val expandButtonBounds = expandPipObject.visibleBounds
         uiDevice.click(expandButtonBounds.centerX(), expandButtonBounds.centerY())
-        wmHelper.waitPipGone()
-        wmHelper.waitForAppTransitionIdle()
+        wmHelper.StateSyncBuilder()
+            .withPipGone()
+            .withFullScreenApp(component)
+            .waitForAndVerify()
     }
 
     /**
@@ -194,7 +203,9 @@ class PipAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
         val windowRect = getWindowRect(wmHelper)
         uiDevice.click(windowRect.centerX(), windowRect.centerY())
         uiDevice.click(windowRect.centerX(), windowRect.centerY())
-        wmHelper.waitForAppTransitionIdle()
+        wmHelper.StateSyncBuilder()
+            .withAppTransitionIdle()
+            .waitForAndVerify()
     }
 
     companion object {
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/EnterPipOnUserLeaveHintTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/EnterPipOnUserLeaveHintTest.kt
index e7d641e9c66e6..cb84fc441f752 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/EnterPipOnUserLeaveHintTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/EnterPipOnUserLeaveHintTest.kt
@@ -17,7 +17,6 @@
 package com.android.wm.shell.flicker.pip
 
 import androidx.test.filters.RequiresDevice
-import com.android.launcher3.tapl.LauncherInstrumentation
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.annotation.Group3
@@ -51,7 +50,6 @@ import org.junit.runners.Parameterized
 @FixMethodOrder(MethodSorters.NAME_ASCENDING)
 @Group3
 class EnterPipOnUserLeaveHintTest(testSpec: FlickerTestParameter) : EnterPipTest(testSpec) {
-    protected val taplInstrumentation = LauncherInstrumentation()
     /**
      * Defines the transition used to run the test
      */
@@ -70,7 +68,7 @@ class EnterPipOnUserLeaveHintTest(testSpec: FlickerTestParameter) : EnterPipTest
                 }
             }
             transitions {
-                taplInstrumentation.goHome()
+                tapl.goHome()
             }
         }
 
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/EnterPipTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/EnterPipTest.kt
index eb318fb256c99..c18c8ce345fc3 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/EnterPipTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/EnterPipTest.kt
@@ -16,16 +16,16 @@
 
 package com.android.wm.shell.flicker.pip
 
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.view.Surface
-import android.platform.test.annotations.FlakyTest
 import androidx.test.filters.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
-import com.android.server.wm.flicker.LAUNCHER_COMPONENT
 import com.android.server.wm.flicker.annotation.Group3
 import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.traces.common.FlickerComponentName.Companion.LAUNCHER
 import org.junit.FixMethodOrder
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -157,15 +157,15 @@ open class EnterPipTest(testSpec: FlickerTestParameter) : PipTransition(testSpec
     }
 
     /**
-     * Checks [LAUNCHER_COMPONENT] layer remains visible throughout the animation
+     * Checks [LAUNCHER] layer remains visible throughout the animation
      */
     @Presubmit
     @Test
     fun launcherLayerBecomesVisible() {
         testSpec.assertLayers {
-            isInvisible(LAUNCHER_COMPONENT)
+            isInvisible(LAUNCHER)
                 .then()
-                .isVisible(LAUNCHER_COMPONENT)
+                .isVisible(LAUNCHER)
         }
     }
 
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/EnterPipToOtherOrientationTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/EnterPipToOtherOrientationTest.kt
index 75481f2199c26..0333577b0d2ac 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/EnterPipToOtherOrientationTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/EnterPipToOtherOrientationTest.kt
@@ -16,9 +16,9 @@
 
 package com.android.wm.shell.flicker.pip
 
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.view.Surface
-import android.platform.test.annotations.FlakyTest
 import androidx.test.filters.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
@@ -33,8 +33,8 @@ import com.android.server.wm.traces.common.FlickerComponentName
 import com.android.wm.shell.flicker.helpers.FixedAppHelper
 import com.android.wm.shell.flicker.pip.PipTransition.BroadcastActionTrigger.Companion.ORIENTATION_LANDSCAPE
 import com.android.wm.shell.flicker.pip.PipTransition.BroadcastActionTrigger.Companion.ORIENTATION_PORTRAIT
-import com.android.wm.shell.flicker.testapp.Components.PipActivity.ACTION_ENTER_PIP
 import com.android.wm.shell.flicker.testapp.Components.FixedActivity.EXTRA_FIXED_ORIENTATION
+import com.android.wm.shell.flicker.testapp.Components.PipActivity.ACTION_ENTER_PIP
 import org.junit.FixMethodOrder
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -98,10 +98,12 @@ class EnterPipToOtherOrientationTest(
                 // Enter PiP, and assert that the PiP is within bounds now that the device is back
                 // in portrait
                 broadcastActionTrigger.doAction(ACTION_ENTER_PIP)
-                wmHelper.waitPipShown()
-                wmHelper.waitForAppTransitionIdle()
                 // during rotation the status bar becomes invisible and reappears at the end
-                wmHelper.waitForNavBarStatusBarVisible()
+                wmHelper.StateSyncBuilder()
+                    .withPipShown()
+                    .withAppTransitionIdle()
+                    .withNavBarStatusBarVisible()
+                    .waitForAndVerify()
             }
         }
 
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipTransition.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipTransition.kt
index 0b4bc761838dc..47215b9fb8831 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipTransition.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipTransition.kt
@@ -19,10 +19,10 @@ package com.android.wm.shell.flicker.pip
 import android.platform.test.annotations.Presubmit
 import android.view.Surface
 import com.android.server.wm.flicker.FlickerTestParameter
-import com.android.server.wm.flicker.LAUNCHER_COMPONENT
 import com.android.server.wm.flicker.dsl.FlickerBuilder
 import com.android.server.wm.flicker.helpers.isShellTransitionsEnabled
 import com.android.server.wm.flicker.helpers.setRotation
+import com.android.server.wm.traces.common.FlickerComponentName.Companion.LAUNCHER
 import org.junit.Test
 
 /**
@@ -78,7 +78,7 @@ abstract class ExitPipTransition(testSpec: FlickerTestParameter) : PipTransition
     }
 
     /**
-     * Checks that [pipApp] and [LAUNCHER_COMPONENT] layers are visible at the start
+     * Checks that [pipApp] and [LAUNCHER] layers are visible at the start
      * of the transition. Then [pipApp] layer becomes invisible, and remains invisible
      * until the end of the transition
      */
@@ -87,10 +87,10 @@ abstract class ExitPipTransition(testSpec: FlickerTestParameter) : PipTransition
     open fun pipLayerBecomesInvisible() {
         testSpec.assertLayers {
             this.isVisible(pipApp.component)
-                .isVisible(LAUNCHER_COMPONENT)
+                .isVisible(LAUNCHER)
                 .then()
                 .isInvisible(pipApp.component)
-                .isVisible(LAUNCHER_COMPONENT)
+                .isVisible(LAUNCHER)
         }
     }
 }
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipViaExpandButtonClickTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipViaExpandButtonClickTest.kt
index 46424a7684082..32022ad5f018e 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipViaExpandButtonClickTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipViaExpandButtonClickTest.kt
@@ -16,8 +16,8 @@
 
 package com.android.wm.shell.flicker.pip
 
-import android.view.Surface
 import android.platform.test.annotations.FlakyTest
+import android.view.Surface
 import androidx.test.filters.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
@@ -74,7 +74,9 @@ class ExitPipViaExpandButtonClickTest(
                 // This will bring PipApp to fullscreen
                 pipApp.expandPipWindowToApp(wmHelper)
                 // Wait until the other app is no longer visible
-                wmHelper.waitForWindowSurfaceDisappeared(testApp.component)
+                wmHelper.StateSyncBuilder()
+                    .withWindowSurfaceDisappeared(testApp.component)
+                    .waitForAndVerify()
             }
         }
 
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipViaIntentTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipViaIntentTest.kt
index 18711d8fd3c66..3fbea488d091c 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipViaIntentTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipViaIntentTest.kt
@@ -16,9 +16,9 @@
 
 package com.android.wm.shell.flicker.pip
 
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.view.Surface
-import android.platform.test.annotations.FlakyTest
 import androidx.test.filters.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
@@ -73,7 +73,9 @@ class ExitPipViaIntentTest(testSpec: FlickerTestParameter) : ExitPipToAppTransit
                 // This will bring PipApp to fullscreen
                 pipApp.exitPipToFullScreenViaIntent(wmHelper)
                 // Wait until the other app is no longer visible
-                wmHelper.waitForWindowSurfaceDisappeared(testApp.component)
+                wmHelper.StateSyncBuilder()
+                    .withWindowSurfaceDisappeared(testApp.component)
+                    .waitForAndVerify()
             }
         }
 
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipWithSwipeDownTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipWithSwipeDownTest.kt
index 9aae7f3242696..210b1968b2da4 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipWithSwipeDownTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipWithSwipeDownTest.kt
@@ -16,9 +16,9 @@
 
 package com.android.wm.shell.flicker.pip
 
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.view.Surface
-import android.platform.test.annotations.FlakyTest
 import androidx.test.filters.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
@@ -64,9 +64,12 @@ class ExitPipWithSwipeDownTest(testSpec: FlickerTestParameter) : ExitPipTransiti
                 val pipCenterY = pipRegion.centerY()
                 val displayCenterX = device.displayWidth / 2
                 device.swipe(pipCenterX, pipCenterY, displayCenterX, device.displayHeight, 10)
-                wmHelper.waitPipGone()
-                wmHelper.waitForWindowSurfaceDisappeared(pipApp.component)
-                wmHelper.waitForAppTransitionIdle()
+                // Wait until the other app is no longer visible
+                wmHelper.StateSyncBuilder()
+                    .withPipGone()
+                    .withWindowSurfaceDisappeared(pipApp.component)
+                    .withAppTransitionIdle()
+                    .waitForAndVerify()
             }
         }
 
@@ -108,4 +111,4 @@ class ExitPipWithSwipeDownTest(testSpec: FlickerTestParameter) : ExitPipTransiti
                             repetitions = 3)
         }
     }
-}
\ No newline at end of file
+}
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExpandPipOnDoubleClickTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExpandPipOnDoubleClickTest.kt
index ffbb89ee8f6b4..dc1fdde18d1b6 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExpandPipOnDoubleClickTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExpandPipOnDoubleClickTest.kt
@@ -23,9 +23,9 @@ import androidx.test.filters.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
-import com.android.server.wm.flicker.LAUNCHER_COMPONENT
 import com.android.server.wm.flicker.annotation.Group3
 import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.traces.common.FlickerComponentName.Companion.LAUNCHER
 import org.junit.FixMethodOrder
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -147,13 +147,13 @@ class ExpandPipOnDoubleClickTest(testSpec: FlickerTestParameter) : PipTransition
     }
 
     /**
-     * Checks [pipApp] layer remains visible throughout the animation
+     * Checks [LAUNCHER] layer remains visible throughout the animation
      */
     @Presubmit
     @Test
     fun launcherIsAlwaysVisible() {
         testSpec.assertLayers {
-            isVisible(LAUNCHER_COMPONENT)
+            isVisible(LAUNCHER)
         }
     }
 
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/PipTransition.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/PipTransition.kt
index 8d542c8ec9e61..fe5dd8b83cfb9 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/PipTransition.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/PipTransition.kt
@@ -21,6 +21,7 @@ import android.content.Intent
 import android.platform.test.annotations.Presubmit
 import android.view.Surface
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.tapl.LauncherInstrumentation
 import com.android.server.wm.flicker.FlickerBuilderProvider
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.dsl.FlickerBuilder
@@ -41,6 +42,7 @@ import org.junit.Test
 
 abstract class PipTransition(protected val testSpec: FlickerTestParameter) {
     protected val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
+    protected val tapl = LauncherInstrumentation()
     protected val pipApp = PipAppHelper(instrumentation)
     protected val displayBounds = WindowUtils.getDisplayBounds(testSpec.startRotation)
     protected val broadcastActionTrigger = BroadcastActionTrigger(instrumentation)
@@ -117,13 +119,11 @@ abstract class PipTransition(protected val testSpec: FlickerTestParameter) {
                 test {
                     if (!eachRun) {
                         pipApp.launchViaIntentAndWaitForPip(wmHelper, stringExtras = stringExtras)
-                        wmHelper.waitPipShown()
                     }
                 }
                 eachRun {
                     if (eachRun) {
                         pipApp.launchViaIntentAndWaitForPip(wmHelper, stringExtras = stringExtras)
-                        wmHelper.waitPipShown()
                     }
                 }
             }
@@ -171,4 +171,4 @@ abstract class PipTransition(protected val testSpec: FlickerTestParameter) {
     @Presubmit
     @Test
     open fun entireScreenCovered() = testSpec.entireScreenCovered()
-}
\ No newline at end of file
+}
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/SetRequestedOrientationWhilePinnedTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/SetRequestedOrientationWhilePinnedTest.kt
index e44cf3866a094..83d3fe2609280 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/SetRequestedOrientationWhilePinnedTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/SetRequestedOrientationWhilePinnedTest.kt
@@ -16,18 +16,18 @@
 
 package com.android.wm.shell.flicker.pip
 
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.view.Surface
-import android.platform.test.annotations.FlakyTest
 import androidx.test.filters.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
 import com.android.server.wm.flicker.annotation.Group4
 import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.flicker.helpers.WindowUtils
 import com.android.server.wm.flicker.helpers.setRotation
 import com.android.server.wm.flicker.helpers.wakeUpAndGoToHomeScreen
-import com.android.server.wm.flicker.helpers.WindowUtils
 import com.android.server.wm.flicker.rules.RemoveAllTasksButHomeRule.Companion.removeAllTasksButHome
 import com.android.wm.shell.flicker.pip.PipTransition.BroadcastActionTrigger.Companion.ORIENTATION_LANDSCAPE
 import com.android.wm.shell.flicker.testapp.Components
@@ -66,11 +66,12 @@ open class SetRequestedOrientationWhilePinnedTest(
                         EXTRA_FIXED_ORIENTATION to ORIENTATION_LANDSCAPE.toString()))
                     // Enter PiP.
                     broadcastActionTrigger.doAction(Components.PipActivity.ACTION_ENTER_PIP)
-                    wmHelper.waitPipShown()
-                    wmHelper.waitForRotation(Surface.ROTATION_0)
-                    wmHelper.waitForAppTransitionIdle()
                     // System bar may fade out during fixed rotation.
-                    wmHelper.waitForNavBarStatusBarVisible()
+                    wmHelper.StateSyncBuilder()
+                        .withPipShown()
+                        .withRotation(Surface.ROTATION_0)
+                        .withNavBarStatusBarVisible()
+                        .waitForAndVerify()
                 }
             }
             teardown {
@@ -85,11 +86,13 @@ open class SetRequestedOrientationWhilePinnedTest(
             transitions {
                 // Launch the activity back into fullscreen and ensure that it is now in landscape
                 pipApp.launchViaIntent(wmHelper)
-                wmHelper.waitForFullScreenApp(pipApp.component)
-                wmHelper.waitForRotation(Surface.ROTATION_90)
-                wmHelper.waitForAppTransitionIdle()
                 // System bar may fade out during fixed rotation.
-                wmHelper.waitForNavBarStatusBarVisible()
+                wmHelper.StateSyncBuilder()
+                    .withFullScreenApp(pipApp.component)
+                    .withRotation(Surface.ROTATION_90)
+                    .withAppTransitionIdle()
+                    .withNavBarStatusBarVisible()
+                    .waitForAndVerify()
             }
         }
 
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipBasicTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipBasicTest.kt
index 49094e609fbc4..31fb16ffbd3e0 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipBasicTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipBasicTest.kt
@@ -43,7 +43,7 @@ class TvPipBasicTest(
 
         // Set up ratio and enter Pip
         testApp.clickObject(radioButtonId)
-        testApp.clickEnterPipButton()
+        testApp.clickEnterPipButton(wmHelper)
 
         val actualRatio: Float = testApp.ui?.visibleBounds?.ratio
                 ?: fail("Application UI not found")
@@ -84,4 +84,4 @@ class TvPipBasicTest(
             )
         }
     }
-}
\ No newline at end of file
+}
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipMenuTests.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipMenuTests.kt
index 061218a015e44..4be19d61278be 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipMenuTests.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipMenuTests.kt
@@ -244,7 +244,7 @@ class TvPipMenuTests : TvPipTestBase() {
     }
 
     private fun enterPip_openMenu_assertShown(): UiObject2 {
-        testApp.clickEnterPipButton()
+        testApp.clickEnterPipButton(wmHelper)
         // Pressing the Window key should bring up Pip menu
         uiDevice.pressWindowKey()
         return uiDevice.waitForTvPipMenu() ?: fail("Pip menu should have been shown")
@@ -256,4 +256,4 @@ class TvPipMenuTests : TvPipTestBase() {
         uiDevice.findTvPipMenuFullscreenButton()
                 ?: fail("\"Full screen\" button should be shown in Pip menu")
     }
-}
\ No newline at end of file
+}
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipNotificationTests.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipNotificationTests.kt
index bcf38d340867f..134e97bd46e7c 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipNotificationTests.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipNotificationTests.kt
@@ -56,7 +56,7 @@ class TvPipNotificationTests : TvPipTestBase() {
     @Test
     fun pipNotification_postedAndDismissed() {
         testApp.launchViaIntent()
-        testApp.clickEnterPipButton()
+        testApp.clickEnterPipButton(wmHelper)
 
         assertNotNull("Pip notification should have been posted",
                 waitForNotificationToAppear { it.isPipNotificationWithTitle(testApp.appName) })
@@ -70,7 +70,7 @@ class TvPipNotificationTests : TvPipTestBase() {
     @Test
     fun pipNotification_closeIntent() {
         testApp.launchViaIntent()
-        testApp.clickEnterPipButton()
+        testApp.clickEnterPipButton(wmHelper)
 
         val notification: StatusBarNotification = waitForNotificationToAppear {
             it.isPipNotificationWithTitle(testApp.appName)
@@ -87,8 +87,8 @@ class TvPipNotificationTests : TvPipTestBase() {
 
     @Test
     fun pipNotification_menuIntent() {
-        testApp.launchViaIntent()
-        testApp.clickEnterPipButton()
+        testApp.launchViaIntent(wmHelper)
+        testApp.clickEnterPipButton(wmHelper)
 
         val notification: StatusBarNotification = waitForNotificationToAppear {
             it.isPipNotificationWithTitle(testApp.appName)
@@ -106,10 +106,10 @@ class TvPipNotificationTests : TvPipTestBase() {
 
     @Test
     fun pipNotification_mediaSessionTitle_isDisplayed() {
-        testApp.launchViaIntent()
+        testApp.launchViaIntent(wmHelper)
         // Start media session and to PiP
         testApp.clickStartMediaSessionButton()
-        testApp.clickEnterPipButton()
+        testApp.clickEnterPipButton(wmHelper)
 
         // Wait for the correct notification to show up...
         waitForNotificationToAppear {
@@ -170,4 +170,4 @@ private val StatusBarNotification.deleteIntent: PendingIntent?
     get() = tvExtensions?.getParcelable("delete_intent")
 
 private fun StatusBarNotification.isPipNotificationWithTitle(expectedTitle: String): Boolean =
-    tag == "TvPip" && title == expectedTitle
\ No newline at end of file
+    tag == "TvPip" && title == expectedTitle
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipTestBase.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipTestBase.kt
index 9c3b0fa183b68..a97994e7d6c36 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipTestBase.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/tv/TvPipTestBase.kt
@@ -23,6 +23,7 @@ import android.os.SystemClock
 import android.view.Surface.ROTATION_0
 import android.view.Surface.rotationToString
 import com.android.server.wm.flicker.helpers.wakeUpAndGoToHomeScreen
+import com.android.server.wm.traces.parser.windowmanager.WindowManagerStateHelper
 import com.android.wm.shell.flicker.SYSTEM_UI_PACKAGE_NAME
 import com.android.wm.shell.flicker.pip.PipTestBase
 import org.junit.After
@@ -33,6 +34,7 @@ import org.junit.Before
 abstract class TvPipTestBase : PipTestBase(rotationToString(ROTATION_0), ROTATION_0) {
 
     private val systemUiProcessObserver = SystemUiProcessObserver()
+    protected val wmHelper = WindowManagerStateHelper()
 
     @Before
     final override fun televisionSetUp() {
@@ -88,4 +90,4 @@ abstract class TvPipTestBase : PipTestBase(rotationToString(ROTATION_0), ROTATIO
     companion object {
         private const val AFTER_TEXT_PROCESS_CHECK_DELAY = 1_000L // 1 sec
     }
-}
\ No newline at end of file
+}
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/splitscreen/EnterSplitScreenByDragFromAllApps.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/splitscreen/EnterSplitScreenByDragFromAllApps.kt
index 702710caded77..cfba3387825c5 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/splitscreen/EnterSplitScreenByDragFromAllApps.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/splitscreen/EnterSplitScreenByDragFromAllApps.kt
@@ -56,7 +56,7 @@ class EnterSplitScreenByDragFromAllApps(
 ) : SplitScreenBase(testSpec) {
 
     @Before
-    open fun before() {
+    fun before() {
         Assume.assumeTrue(taplInstrumentation.isTablet)
     }
 
@@ -73,8 +73,7 @@ class EnterSplitScreenByDragFromAllApps(
                 taplInstrumentation.launchedAppState.taskbar
                     .openAllApps()
                     .getAppIcon(secondaryApp.appName)
-                    .dragToSplitscreen(secondaryApp.component.packageName,
-                        primaryApp.component.packageName)
+                    .dragToSplitscreen(secondaryApp.`package`, primaryApp.`package`)
             }
         }
 
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 315c40ffa9ba5..896c7309cf4cb 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/CommonAssertions.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/CommonAssertions.kt
@@ -111,8 +111,8 @@ fun FlickerTestParameter.statusBarLayerIsVisible() {
  */
 fun FlickerTestParameter.navBarLayerPositionStart() {
     assertLayersStart {
-        val display = this.entry.displays.minByOrNull { it.id }
-            ?: throw RuntimeException("There is no display!")
+        val display = this.entry.displays.firstOrNull { !it.isVirtual }
+                ?: error("There is no display!")
         this.visibleRegion(FlickerComponentName.NAV_BAR)
             .coversExactly(WindowUtils.getNavigationBarPosition(display, isGesturalNavigation))
     }
@@ -180,7 +180,8 @@ fun FlickerTestParameter.statusBarLayerRotatesScales() {
  * the visibleRegion of the given app component exactly
  */
 fun FlickerTestParameter.snapshotStartingWindowLayerCoversExactlyOnApp(
-        component: FlickerComponentName) {
+    component: FlickerComponentName
+) {
     assertLayers {
         invoke("snapshotStartingWindowLayerCoversExactlyOnApp") {
             val snapshotLayers = it.subjects.filter { subject ->
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/close/CloseAppBackButtonTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/close/CloseAppBackButtonTest.kt
index 6c5f6bace49db..d0f1dfddaed02 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/close/CloseAppBackButtonTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/close/CloseAppBackButtonTest.kt
@@ -72,13 +72,15 @@ class CloseAppBackButtonTest(testSpec: FlickerTestParameter) : CloseAppTransitio
         get() = {
             super.transition(this)
             transitions {
-                device.pressBack()
-                wmHelper.waitForHomeActivityVisible()
+                tapl.pressBack()
+                wmHelper.StateSyncBuilder()
+                    .withHomeActivityVisible()
+                    .waitForAndVerify()
             }
         }
 
     /** {@inheritDoc} */
-    @FlakyTest
+    @FlakyTest(bugId = 206753786)
     @Test
     override fun navBarLayerRotatesAndScales() = super.navBarLayerRotatesAndScales()
 
@@ -118,4 +120,4 @@ class CloseAppBackButtonTest(testSpec: FlickerTestParameter) : CloseAppTransitio
                 .getConfigNonRotationTests(repetitions = 3)
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/close/CloseAppHomeButtonTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/close/CloseAppHomeButtonTest.kt
index 7451a523f2609..b178e8c7ecedb 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/close/CloseAppHomeButtonTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/close/CloseAppHomeButtonTest.kt
@@ -16,7 +16,6 @@
 
 package com.android.server.wm.flicker.close
 
-import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.FlakyTest
 import androidx.test.filters.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
@@ -69,37 +68,29 @@ class CloseAppHomeButtonTest(testSpec: FlickerTestParameter) : CloseAppTransitio
         get() = {
             super.transition(this)
             transitions {
+                // Can't use TAPL at the moment because of rotation test issues
+                // When pressing home, TAPL expects the orientation to remain constant
+                // However, when closing a landscape app back to a portrait-only launcher
+                // this causes an error in verifyActiveContainer();
                 device.pressHome()
-                wmHelper.waitForHomeActivityVisible()
+                wmHelper.StateSyncBuilder()
+                    .withHomeActivityVisible()
+                    .waitForAndVerify()
             }
         }
 
     /** {@inheritDoc} */
-    @FlakyTest
+    @FlakyTest(bugId = 206753786)
     @Test
     override fun navBarLayerRotatesAndScales() = super.navBarLayerRotatesAndScales()
 
     /** {@inheritDoc} */
-    @FlakyTest(bugId = 227430489)
+    @FlakyTest(bugId = 206753786)
     @Test
     override fun statusBarLayerRotatesScales() {
         super.statusBarLayerRotatesScales()
     }
 
-    /** {@inheritDoc} */
-    @Presubmit
-    @Test
-    override fun launcherLayerReplacesApp() {
-        super.launcherLayerReplacesApp()
-    }
-
-    /** {@inheritDoc} */
-    @Presubmit
-    @Test
-    override fun entireScreenCovered() {
-        super.entireScreenCovered()
-    }
-
     /** {@inheritDoc} */
     @FlakyTest(bugId = 229762973)
     @Test
@@ -120,4 +111,4 @@ class CloseAppHomeButtonTest(testSpec: FlickerTestParameter) : CloseAppTransitio
                 .getConfigNonRotationTests(repetitions = 3)
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/close/CloseAppTransition.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/close/CloseAppTransition.kt
index aaa2db7687927..ae002c9458f5f 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/close/CloseAppTransition.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/close/CloseAppTransition.kt
@@ -19,21 +19,22 @@ package com.android.server.wm.flicker.close
 import android.app.Instrumentation
 import android.platform.test.annotations.Presubmit
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.tapl.LauncherInstrumentation
 import com.android.server.wm.flicker.FlickerBuilderProvider
 import com.android.server.wm.flicker.FlickerTestParameter
-import com.android.server.wm.flicker.LAUNCHER_COMPONENT
 import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.helpers.SimpleAppHelper
 import com.android.server.wm.flicker.helpers.StandardAppHelper
 import com.android.server.wm.flicker.helpers.setRotation
 import com.android.server.wm.flicker.navBarLayerIsVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
 import com.android.server.wm.flicker.navBarWindowIsVisible
-import com.android.server.wm.flicker.entireScreenCovered
+import com.android.server.wm.flicker.replacesLayer
 import com.android.server.wm.flicker.statusBarLayerIsVisible
 import com.android.server.wm.flicker.statusBarLayerRotatesScales
 import com.android.server.wm.flicker.statusBarWindowIsVisible
-import com.android.server.wm.flicker.replacesLayer
+import com.android.server.wm.traces.common.FlickerComponentName.Companion.LAUNCHER
 import org.junit.Test
 
 /**
@@ -42,12 +43,16 @@ import org.junit.Test
 abstract class CloseAppTransition(protected val testSpec: FlickerTestParameter) {
     protected val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
     protected open val testApp: StandardAppHelper = SimpleAppHelper(instrumentation)
+    protected val tapl = LauncherInstrumentation()
 
     /**
      * Specification of the test transition to execute
      */
     protected open val transition: FlickerBuilder.() -> Unit = {
         setup {
+            test {
+                tapl.setExpectedRotation(testSpec.startRotation)
+            }
             eachRun {
                 testApp.launchViaIntent(wmHelper)
                 this.setRotation(testSpec.startRotation)
@@ -154,7 +159,7 @@ abstract class CloseAppTransition(protected val testSpec: FlickerTestParameter)
 
     /**
      * Checks that [testApp] is the top visible app window at the start of the transition and
-     * that it is replaced by [LAUNCHER_COMPONENT] during the transition
+     * that it is replaced by [LAUNCHER] during the transition
      */
     @Presubmit
     @Test
@@ -162,30 +167,30 @@ abstract class CloseAppTransition(protected val testSpec: FlickerTestParameter)
         testSpec.assertWm {
             this.isAppWindowOnTop(testApp.component)
                     .then()
-                    .isAppWindowOnTop(LAUNCHER_COMPONENT)
+                    .isAppWindowOnTop(LAUNCHER)
         }
     }
 
     /**
-     * Checks that [LAUNCHER_COMPONENT] is invisible at the start of the transition and that
+     * Checks that [LAUNCHER] is invisible at the start of the transition and that
      * it becomes visible during the transition
      */
     @Presubmit
     @Test
     open fun launcherWindowBecomesVisible() {
         testSpec.assertWm {
-            this.isAppWindowNotOnTop(LAUNCHER_COMPONENT)
+            this.isAppWindowNotOnTop(LAUNCHER)
                     .then()
-                    .isAppWindowOnTop(LAUNCHER_COMPONENT)
+                    .isAppWindowOnTop(LAUNCHER)
         }
     }
 
     /**
-     * Checks that [LAUNCHER_COMPONENT] layer becomes visible when [testApp] becomes invisible
+     * Checks that [LAUNCHER] layer becomes visible when [testApp] becomes invisible
      */
     @Presubmit
     @Test
     open fun launcherLayerReplacesApp() {
-        testSpec.replacesLayer(testApp.component, LAUNCHER_COMPONENT)
+        testSpec.replacesLayer(testApp.component, LAUNCHER)
     }
 }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeAppAutoFocusHelper.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeAppAutoFocusHelper.kt
index aacc17a49a242..cd184ac10c222 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeAppAutoFocusHelper.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeAppAutoFocusHelper.kt
@@ -16,19 +16,16 @@
 
 package com.android.server.wm.flicker.helpers
 
+import android.app.Instrumentation
 import android.view.WindowInsets.Type.ime
 import android.view.WindowInsets.Type.navigationBars
 import android.view.WindowInsets.Type.statusBars
-
-import android.app.Instrumentation
 import androidx.test.uiautomator.By
-import androidx.test.uiautomator.UiDevice
 import androidx.test.uiautomator.Until
 import com.android.server.wm.flicker.testapp.ActivityOptions
 import com.android.server.wm.traces.common.FlickerComponentName
 import com.android.server.wm.traces.parser.toFlickerComponent
 import com.android.server.wm.traces.parser.windowmanager.WindowManagerStateHelper
-
 import java.util.regex.Pattern
 
 class ImeAppAutoFocusHelper @JvmOverloads constructor(
@@ -39,12 +36,9 @@ class ImeAppAutoFocusHelper @JvmOverloads constructor(
     component: FlickerComponentName =
         ActivityOptions.IME_ACTIVITY_AUTO_FOCUS_COMPONENT_NAME.toFlickerComponent()
 ) : ImeAppHelper(instr, launcherName, component) {
-    override fun openIME(
-        device: UiDevice,
-        wmHelper: WindowManagerStateHelper?
-    ) {
+    override fun openIME(wmHelper: WindowManagerStateHelper) {
         // do nothing (the app is focused automatically)
-        waitIMEShown(device, wmHelper)
+        waitIMEShown(wmHelper)
     }
 
     override fun launchViaIntent(
@@ -54,7 +48,7 @@ class ImeAppAutoFocusHelper @JvmOverloads constructor(
         stringExtras: Map<String, String>
     ) {
         super.launchViaIntent(wmHelper, expectedWindowName, action, stringExtras)
-        waitIMEShown(uiDevice, wmHelper)
+        waitIMEShown(wmHelper)
     }
 
     override fun open() {
@@ -70,15 +64,15 @@ class ImeAppAutoFocusHelper @JvmOverloads constructor(
         val button = uiDevice.wait(Until.findObject(By.res(getPackage(),
                 "start_dialog_themed_activity_btn")), FIND_TIMEOUT)
 
-        require(button != null) {
+        requireNotNull(button) {
             "Button not found, this usually happens when the device " +
                     "was left in an unknown state (e.g. Screen turned off)"
         }
         button.click()
-        wmHelper.waitForAppTransitionIdle()
-        wmHelper.waitForFullScreenApp(
+        wmHelper.StateSyncBuilder()
+            .withFullScreenApp(
                 ActivityOptions.DIALOG_THEMED_ACTIVITY_COMPONENT_NAME.toFlickerComponent())
-        mInstrumentation.waitForIdleSync()
+            .waitForAndVerify()
     }
     fun dismissDialog(wmHelper: WindowManagerStateHelper) {
         val dialog = uiDevice.wait(
@@ -87,14 +81,16 @@ class ImeAppAutoFocusHelper @JvmOverloads constructor(
         // Pressing back key to dismiss the dialog
         if (dialog != null) {
             uiDevice.pressBack()
-            wmHelper.waitForAppTransitionIdle()
+            wmHelper.StateSyncBuilder()
+                .withAppTransitionIdle()
+                .waitForAndVerify()
         }
     }
     fun getInsetsVisibleFromDialog(type: Int): Boolean {
-        var insetsVisibilityTextView = uiDevice.wait(
+        val insetsVisibilityTextView = uiDevice.wait(
                 Until.findObject(By.res("android:id/text1")), FIND_TIMEOUT)
         if (insetsVisibilityTextView != null) {
-            var visibility = insetsVisibilityTextView.text.toString()
+            val visibility = insetsVisibilityTextView.text.toString()
             val matcher = when (type) {
                 ime() -> {
                     Pattern.compile("IME\\: (VISIBLE|INVISIBLE)").matcher(visibility)
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 5bd365c7eefd6..8ac9a75aff94d 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
@@ -20,7 +20,6 @@ import android.app.Instrumentation
 import android.support.test.launcherhelper.ILauncherStrategy
 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.flicker.testapp.ActivityOptions
 import com.android.server.wm.traces.common.FlickerComponentName
@@ -39,64 +38,49 @@ open class ImeAppHelper @JvmOverloads constructor(
     /**
      * Opens the IME and wait for it to be displayed
      *
-     * @param device UIDevice instance to interact with the device
      * @param wmHelper Helper used to wait for WindowManager states
      */
-    @JvmOverloads
-    open fun openIME(device: UiDevice, wmHelper: WindowManagerStateHelper? = null) {
-        val editText = device.wait(
+    open fun openIME(wmHelper: WindowManagerStateHelper) {
+        val editText = uiDevice.wait(
             Until.findObject(By.res(getPackage(), "plain_text_input")),
             FIND_TIMEOUT)
 
-        require(editText != null) {
+        requireNotNull(editText) {
             "Text field not found, this usually happens when the device " +
                 "was left in an unknown state (e.g. in split screen)"
         }
         editText.click()
-        waitIMEShown(device, wmHelper)
+        waitIMEShown(wmHelper)
     }
 
-    protected fun waitIMEShown(
-        device: UiDevice,
-        wmHelper: WindowManagerStateHelper? = null
-    ) {
-        if (wmHelper == null) {
-            device.waitForIdle()
-        } else {
-            wmHelper.waitImeShown()
-        }
+    protected fun waitIMEShown(wmHelper: WindowManagerStateHelper) {
+        wmHelper.StateSyncBuilder()
+            .withImeShown()
+            .waitForAndVerify()
     }
 
     /**
      * Opens the IME and wait for it to be gone
      *
-     * @param device UIDevice instance to interact with the device
      * @param wmHelper Helper used to wait for WindowManager states
      */
-    @JvmOverloads
-    open fun closeIME(device: UiDevice, wmHelper: WindowManagerStateHelper? = null) {
-        device.pressBack()
-        // Using only the AccessibilityInfo it is not possible to identify if the IME is active
-        if (wmHelper == null) {
-            device.waitForIdle()
-        } else {
-            wmHelper.waitImeGone()
-        }
+    open fun closeIME(wmHelper: WindowManagerStateHelper) {
+        uiDevice.pressBack()
+        wmHelper.StateSyncBuilder()
+            .withImeGone()
+            .waitForAndVerify()
     }
 
-    @JvmOverloads
-    open fun finishActivity(device: UiDevice, wmHelper: WindowManagerStateHelper? = null) {
-        val finishButton = device.wait(
+    open fun finishActivity(wmHelper: WindowManagerStateHelper) {
+        val finishButton = uiDevice.wait(
                 Until.findObject(By.res(getPackage(), "finish_activity_btn")),
                 FIND_TIMEOUT)
-        require(finishButton != null) {
-            "Finish activity button not found, probably IME activity is not on the screen ?"
+        requireNotNull(finishButton) {
+            "Finish activity button not found, probably IME activity is not on the screen?"
         }
         finishButton.click()
-        if (wmHelper == null) {
-            device.waitForIdle()
-        } else {
-            wmHelper.waitForActivityRemoved(component)
-        }
+        wmHelper.StateSyncBuilder()
+            .withActivityRemoved(component)
+            .waitForAndVerify()
     }
-}
\ No newline at end of file
+}
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeEditorPopupDialogAppHelper.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeEditorPopupDialogAppHelper.kt
index 172c4330c3c62..438aeeb2cc74d 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeEditorPopupDialogAppHelper.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeEditorPopupDialogAppHelper.kt
@@ -18,7 +18,6 @@ package com.android.server.wm.flicker.helpers
 
 import android.app.Instrumentation
 import androidx.test.uiautomator.By
-import androidx.test.uiautomator.UiDevice
 import androidx.test.uiautomator.Until
 import com.android.server.wm.flicker.testapp.ActivityOptions
 import com.android.server.wm.traces.common.FlickerComponentName
@@ -33,18 +32,15 @@ class ImeEditorPopupDialogAppHelper @JvmOverloads constructor(
     component: FlickerComponentName =
             ActivityOptions.EDITOR_POPUP_DIALOG_ACTIVITY_COMPONENT_NAME.toFlickerComponent()
 ) : ImeAppHelper(instr, launcherName, component) {
-    override fun openIME(
-        device: UiDevice,
-        wmHelper: WindowManagerStateHelper?
-    ) {
-        val editText = device.wait(Until.findObject(By.text("focused editText")), FIND_TIMEOUT)
+    override fun openIME(wmHelper: WindowManagerStateHelper) {
+        val editText = uiDevice.wait(Until.findObject(By.text("focused editText")), FIND_TIMEOUT)
 
         require(editText != null) {
             "Text field not found, this usually happens when the device " +
                     "was left in an unknown state (e.g. in split screen)"
         }
         editText.click()
-        waitIMEShown(device, wmHelper)
+        waitIMEShown(wmHelper)
     }
 
     fun dismissDialog(wmHelper: WindowManagerStateHelper) {
@@ -54,7 +50,9 @@ class ImeEditorPopupDialogAppHelper @JvmOverloads constructor(
         // Pressing back key to dismiss the dialog
         if (dismissButton != null) {
             dismissButton.click()
-            wmHelper.waitForAppTransitionIdle()
+            wmHelper.StateSyncBuilder()
+                .withAppTransitionIdle()
+                .waitForAndVerify()
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/NewTasksAppHelper.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/NewTasksAppHelper.kt
index be68704fc32d5..1a2543cb8f977 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/NewTasksAppHelper.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/NewTasksAppHelper.kt
@@ -41,12 +41,13 @@ class NewTasksAppHelper @JvmOverloads constructor(
             Until.findObject(By.res(getPackage(), "launch_new_task")),
             FIND_TIMEOUT)
 
-        require(button != null) {
+        requireNotNull(button) {
             "Button not found, this usually happens when the device " +
                     "was left in an unknown state (e.g. in split screen)"
         }
         button.click()
-        wmHelper.waitForAppTransitionIdle()
-        wmHelper.waitForFullScreenApp(component)
+        wmHelper.StateSyncBuilder()
+            .withFullScreenApp(component)
+            .waitForAndVerify()
     }
-}
\ No newline at end of file
+}
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/NotificationAppHelper.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/NotificationAppHelper.kt
index 4e360f98723e3..00b208bc8c23e 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/NotificationAppHelper.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/NotificationAppHelper.kt
@@ -20,7 +20,6 @@ import android.app.Instrumentation
 import android.support.test.launcherhelper.ILauncherStrategy
 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.flicker.testapp.ActivityOptions
 import com.android.server.wm.traces.common.FlickerComponentName
@@ -36,8 +35,8 @@ class NotificationAppHelper @JvmOverloads constructor(
             .getInstance(instr)
             .launcherStrategy
 ) : StandardAppHelper(instr, launcherName, component, launcherStrategy) {
-    fun postNotification(device: UiDevice, wmHelper: WindowManagerStateHelper) {
-        val button = device.wait(
+    fun postNotification(wmHelper: WindowManagerStateHelper) {
+        val button = uiDevice.wait(
                 Until.findObject(By.res(getPackage(), "post_notification")),
                 FIND_TIMEOUT)
 
@@ -47,7 +46,7 @@ class NotificationAppHelper @JvmOverloads constructor(
         }
         button.click()
 
-        device.wait(Until.findObject(By.text("Flicker Test Notification")), FIND_TIMEOUT)
+        uiDevice.wait(Until.findObject(By.text("Flicker Test Notification")), FIND_TIMEOUT)
                 ?: error("Flicker Notification not found")
     }
-}
\ No newline at end of file
+}
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/TwoActivitiesAppHelper.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/TwoActivitiesAppHelper.kt
index a135e0af067b6..8519da14aa2db 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/TwoActivitiesAppHelper.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/TwoActivitiesAppHelper.kt
@@ -19,13 +19,11 @@ package com.android.server.wm.flicker.helpers
 import android.app.Instrumentation
 import android.support.test.launcherhelper.ILauncherStrategy
 import android.support.test.launcherhelper.LauncherStrategyFactory
-import android.view.Display
 import androidx.test.uiautomator.By
 import androidx.test.uiautomator.UiDevice
 import androidx.test.uiautomator.Until
 import com.android.server.wm.flicker.testapp.ActivityOptions
 import com.android.server.wm.traces.common.FlickerComponentName
-import com.android.server.wm.traces.common.WindowManagerConditionsFactory
 import com.android.server.wm.traces.parser.toFlickerComponent
 import com.android.server.wm.traces.parser.windowmanager.WindowManagerStateHelper
 
@@ -46,21 +44,19 @@ class TwoActivitiesAppHelper @JvmOverloads constructor(
         val launchActivityButton = By.res(getPackage(), LAUNCH_SECOND_ACTIVITY)
         val button = device.wait(Until.findObject(launchActivityButton), FIND_TIMEOUT)
 
-        require(button != null) {
+        requireNotNull(button) {
             "Button not found, this usually happens when the device " +
                     "was left in an unknown state (e.g. in split screen)"
         }
         button.click()
 
         device.wait(Until.gone(launchActivityButton), FIND_TIMEOUT)
-        wmHelper.waitFor(
-            WindowManagerStateHelper.isAppFullScreen(secondActivityComponent),
-            WindowManagerConditionsFactory.isAppTransitionIdle(Display.DEFAULT_DISPLAY),
-            WindowManagerConditionsFactory.hasLayersAnimating().negate()
-        )
+        wmHelper.StateSyncBuilder()
+            .withFullScreenApp(secondActivityComponent)
+            .waitForAndVerify()
     }
 
     companion object {
         private const val LAUNCH_SECOND_ACTIVITY = "launch_second_activity"
     }
-}
\ 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 843ef232fdebd..67d2067d1b06b 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
@@ -17,10 +17,10 @@
 package com.android.server.wm.flicker.ime
 
 import android.app.Instrumentation
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.view.Surface
 import android.view.WindowManagerPolicyConstants
-import android.platform.test.annotations.FlakyTest
 import androidx.test.filters.RequiresDevice
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.server.wm.flicker.FlickerBuilderProvider
@@ -71,7 +71,6 @@ class CloseImeAutoOpenWindowToAppTest(private val testSpec: FlickerTestParameter
             setup {
                 eachRun {
                     testApp.launchViaIntent(wmHelper)
-                    testApp.openIME(device, wmHelper)
                 }
             }
             teardown {
@@ -80,7 +79,7 @@ class CloseImeAutoOpenWindowToAppTest(private val testSpec: FlickerTestParameter
                 }
             }
             transitions {
-                testApp.closeIME(device, wmHelper)
+                testApp.closeIME(wmHelper)
             }
         }
     }
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 13a49a1cb47c8..dcb5dad3e9313 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
@@ -17,10 +17,10 @@
 package com.android.server.wm.flicker.ime
 
 import android.app.Instrumentation
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.view.Surface
 import android.view.WindowManagerPolicyConstants
-import android.platform.test.annotations.FlakyTest
 import androidx.test.filters.RequiresDevice
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.server.wm.flicker.FlickerBuilderProvider
@@ -29,11 +29,11 @@ import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
 import com.android.server.wm.flicker.annotation.Group2
 import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.helpers.ImeAppAutoFocusHelper
 import com.android.server.wm.flicker.navBarLayerIsVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
 import com.android.server.wm.flicker.navBarWindowIsVisible
-import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.statusBarLayerIsVisible
 import com.android.server.wm.flicker.statusBarLayerRotatesScales
 import com.android.server.wm.flicker.statusBarWindowIsVisible
@@ -71,7 +71,6 @@ class CloseImeAutoOpenWindowToHomeTest(private val testSpec: FlickerTestParamete
             setup {
                 eachRun {
                     testApp.launchViaIntent(wmHelper)
-                    testApp.openIME(device, wmHelper)
                 }
             }
             teardown {
@@ -81,8 +80,10 @@ class CloseImeAutoOpenWindowToHomeTest(private val testSpec: FlickerTestParamete
             }
             transitions {
                 device.pressHome()
-                wmHelper.waitForHomeActivityVisible()
-                wmHelper.waitImeGone()
+                wmHelper.StateSyncBuilder()
+                    .withHomeActivityVisible()
+                    .withImeGone()
+                    .waitForAndVerify()
             }
         }
     }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeEditorPopupDialogTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeEditorPopupDialogTest.kt
index 2e29b3e314ca7..b979f3ec3442d 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeEditorPopupDialogTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeEditorPopupDialogTest.kt
@@ -23,9 +23,9 @@ import android.view.WindowManagerPolicyConstants
 import androidx.test.filters.RequiresDevice
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.server.wm.flicker.FlickerBuilderProvider
-import com.android.server.wm.flicker.FlickerTestParameterFactory
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
+import com.android.server.wm.flicker.FlickerTestParameterFactory
 import com.android.server.wm.flicker.annotation.Group4
 import com.android.server.wm.flicker.dsl.FlickerBuilder
 import com.android.server.wm.flicker.helpers.ImeEditorPopupDialogAppHelper
@@ -54,18 +54,22 @@ class CloseImeEditorPopupDialogTest(private val testSpec: FlickerTestParameter)
             setup {
                 eachRun {
                     imeTestApp.launchViaIntent(wmHelper)
-                    imeTestApp.openIME(device, wmHelper)
+                    imeTestApp.openIME(wmHelper)
                 }
             }
             transitions {
                 imeTestApp.dismissDialog(wmHelper)
-                wmHelper.waitImeGone()
+                wmHelper.StateSyncBuilder()
+                    .withImeGone()
+                    .waitForAndVerify()
             }
             teardown {
                 eachRun {
                     device.pressHome()
-                    wmHelper.waitForHomeActivityVisible()
-                    imeTestApp.exit()
+                    wmHelper.StateSyncBuilder()
+                        .withHomeActivityVisible()
+                        .waitForAndVerify()
+                    imeTestApp.exit(wmHelper)
                 }
             }
         }
@@ -133,4 +137,4 @@ class CloseImeEditorPopupDialogTest(private val testSpec: FlickerTestParameter)
                     )
         }
     }
-}
\ No newline at end of file
+}
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 1ce1e1fb05f6b..928bb5346671e 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
@@ -17,8 +17,8 @@
 package com.android.server.wm.flicker.ime
 
 import android.app.Instrumentation
-import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.FlakyTest
+import android.platform.test.annotations.Presubmit
 import androidx.test.filters.RequiresDevice
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.server.wm.flicker.FlickerBuilderProvider
@@ -27,11 +27,11 @@ import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
 import com.android.server.wm.flicker.annotation.Group2
 import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.helpers.ImeAppHelper
 import com.android.server.wm.flicker.navBarLayerIsVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
 import com.android.server.wm.flicker.navBarWindowIsVisible
-import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.statusBarLayerRotatesScales
 import com.android.server.wm.flicker.statusBarWindowIsVisible
 import com.android.server.wm.traces.common.FlickerComponentName
@@ -64,7 +64,7 @@ class CloseImeWindowToAppTest(private val testSpec: FlickerTestParameter) {
                     testApp.launchViaIntent(wmHelper)
                 }
                 eachRun {
-                    testApp.openIME(device, wmHelper)
+                    testApp.openIME(wmHelper)
                 }
             }
             teardown {
@@ -73,7 +73,7 @@ class CloseImeWindowToAppTest(private val testSpec: FlickerTestParameter) {
                 }
             }
             transitions {
-                testApp.closeIME(device, wmHelper)
+                testApp.closeIME(wmHelper)
             }
         }
     }
@@ -163,4 +163,4 @@ class CloseImeWindowToAppTest(private val testSpec: FlickerTestParameter) {
                 .getConfigNonRotationTests(repetitions = 3)
         }
     }
-}
\ No newline at end of file
+}
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 c77ea12a58509..3ff809d383301 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
@@ -17,10 +17,10 @@
 package com.android.server.wm.flicker.ime
 
 import android.app.Instrumentation
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.view.Surface
 import android.view.WindowManagerPolicyConstants
-import android.platform.test.annotations.FlakyTest
 import androidx.test.filters.RequiresDevice
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.server.wm.flicker.FlickerBuilderProvider
@@ -29,11 +29,11 @@ import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
 import com.android.server.wm.flicker.annotation.Group2
 import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.helpers.ImeAppHelper
 import com.android.server.wm.flicker.navBarLayerIsVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
 import com.android.server.wm.flicker.navBarWindowIsVisible
-import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.statusBarLayerRotatesScales
 import com.android.server.wm.flicker.statusBarWindowIsVisible
 import com.android.server.wm.traces.common.FlickerComponentName
@@ -42,6 +42,7 @@ import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.MethodSorters
 import org.junit.runners.Parameterized
+
 /**
  * Test IME window closing to home transitions.
  * To run this test: `atest FlickerTests:CloseImeWindowToHomeTest`
@@ -61,21 +62,19 @@ class CloseImeWindowToHomeTest(private val testSpec: FlickerTestParameter) {
             setup {
                 eachRun {
                     testApp.launchViaIntent(wmHelper)
-                    testApp.openIME(device, wmHelper)
+                    testApp.openIME(wmHelper)
                 }
             }
             transitions {
                 device.pressHome()
-                wmHelper.waitForHomeActivityVisible()
-                wmHelper.waitImeGone()
+                wmHelper.StateSyncBuilder()
+                    .withHomeActivityVisible()
+                    .withImeGone()
+                    .waitForAndVerify()
             }
             teardown {
-                eachRun {
-                    device.pressHome()
-                    wmHelper.waitForHomeActivityVisible()
-                }
                 test {
-                    testApp.exit()
+                    testApp.exit(wmHelper)
                 }
             }
         }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/LaunchAppShowImeAndDialogThemeAppTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/LaunchAppShowImeAndDialogThemeAppTest.kt
index b298df2f8ffe3..d1761fca899bc 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/LaunchAppShowImeAndDialogThemeAppTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/LaunchAppShowImeAndDialogThemeAppTest.kt
@@ -16,15 +16,14 @@
 
 package com.android.server.wm.flicker.ime
 
+import android.app.Instrumentation
+import android.platform.test.annotations.FlakyTest
+import android.platform.test.annotations.Presubmit
+import android.view.Surface
 import android.view.WindowInsets.Type.ime
 import android.view.WindowInsets.Type.navigationBars
 import android.view.WindowInsets.Type.statusBars
-
-import android.app.Instrumentation
-import android.platform.test.annotations.Presubmit
-import android.view.Surface
 import android.view.WindowManagerPolicyConstants
-import android.platform.test.annotations.FlakyTest
 import androidx.test.filters.RequiresDevice
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.server.wm.flicker.FlickerBuilderProvider
@@ -34,13 +33,13 @@ import com.android.server.wm.flicker.FlickerTestParameterFactory
 import com.android.server.wm.flicker.dsl.FlickerBuilder
 import com.android.server.wm.flicker.helpers.ImeAppAutoFocusHelper
 import com.android.server.wm.traces.common.FlickerComponentName
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertTrue
 import org.junit.FixMethodOrder
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.MethodSorters
 import org.junit.runners.Parameterized
-import org.junit.Assert.assertFalse
-import org.junit.Assert.assertTrue
 
 /**
  * Test IME snapshot mechanism won't apply when transitioning from non-IME focused dialog activity.
@@ -60,7 +59,9 @@ class LaunchAppShowImeAndDialogThemeAppTest(private val testSpec: FlickerTestPar
             setup {
                 eachRun {
                     testApp.launchViaIntent(wmHelper)
-                    wmHelper.waitImeShown()
+                    wmHelper.StateSyncBuilder()
+                        .withImeShown()
+                        .waitForAndVerify()
                     testApp.startDialogThemedActivity(wmHelper)
                     // Verify IME insets isn't visible on dialog since it's non-IME focusable window
                     assertFalse(testApp.getInsetsVisibleFromDialog(ime()))
@@ -70,7 +71,7 @@ class LaunchAppShowImeAndDialogThemeAppTest(private val testSpec: FlickerTestPar
             }
             teardown {
                 eachRun {
-                    testApp.exit()
+                    testApp.exit(wmHelper)
                 }
             }
             transitions {
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/LaunchAppShowImeOnStartTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/LaunchAppShowImeOnStartTest.kt
index b897ca2a9c153..d393ace3254dd 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/LaunchAppShowImeOnStartTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/LaunchAppShowImeOnStartTest.kt
@@ -79,19 +79,21 @@ class LaunchAppShowImeOnStartTest(private val testSpec: FlickerTestParameter) {
         return FlickerBuilder(instrumentation).apply {
             setup {
                 eachRun {
-                    initializeApp.launchViaIntent()
+                    initializeApp.launchViaIntent(wmHelper)
                     this.setRotation(testSpec.startRotation)
                 }
             }
             teardown {
                 eachRun {
-                    initializeApp.exit()
-                    testApp.exit()
+                    initializeApp.exit(wmHelper)
+                    testApp.exit(wmHelper)
                 }
             }
             transitions {
                 testApp.launchViaIntent(wmHelper)
-                wmHelper.waitImeShown()
+                wmHelper.StateSyncBuilder()
+                    .withImeShown()
+                    .waitForAndVerify()
             }
         }
     }
@@ -153,4 +155,4 @@ class LaunchAppShowImeOnStartTest(private val testSpec: FlickerTestParameter) {
                 )
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowAndCloseTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowAndCloseTest.kt
index 972918e28fa72..1854d6e5506e5 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowAndCloseTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowAndCloseTest.kt
@@ -22,11 +22,19 @@ import android.view.Surface
 import android.view.WindowManagerPolicyConstants
 import androidx.test.filters.RequiresDevice
 import androidx.test.platform.app.InstrumentationRegistry
-import com.android.server.wm.flicker.*
+import com.android.server.wm.flicker.FlickerBuilderProvider
+import com.android.server.wm.flicker.FlickerParametersRunnerFactory
+import com.android.server.wm.flicker.FlickerTestParameter
+import com.android.server.wm.flicker.FlickerTestParameterFactory
 import com.android.server.wm.flicker.annotation.Group2
 import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.helpers.ImeAppHelper
 import com.android.server.wm.flicker.helpers.SimpleAppHelper
+import com.android.server.wm.flicker.navBarLayerIsVisible
+import com.android.server.wm.flicker.navBarWindowIsVisible
+import com.android.server.wm.flicker.statusBarLayerIsVisible
+import com.android.server.wm.flicker.statusBarWindowIsVisible
 import org.junit.FixMethodOrder
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -56,15 +64,15 @@ class OpenImeWindowAndCloseTest(private val testSpec: FlickerTestParameter) {
                 eachRun {
                     simpleApp.launchViaIntent(wmHelper)
                     testApp.launchViaIntent(wmHelper)
-                    testApp.openIME(device, wmHelper)
+                    testApp.openIME(wmHelper)
                 }
             }
             transitions {
-                testApp.finishActivity(device, wmHelper)
+                testApp.finishActivity(wmHelper)
             }
             teardown {
                 test {
-                    simpleApp.exit()
+                    simpleApp.exit(wmHelper)
                 }
             }
         }
@@ -128,4 +136,4 @@ class OpenImeWindowAndCloseTest(private val testSpec: FlickerTestParameter) {
                     )
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowFromFixedOrientationAppTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowFromFixedOrientationAppTest.kt
index fdaeacb5e8a28..4f26fc2841988 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowFromFixedOrientationAppTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowFromFixedOrientationAppTest.kt
@@ -17,12 +17,12 @@
 package com.android.server.wm.flicker.ime
 
 import android.app.Instrumentation
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Postsubmit
 import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.RequiresDevice
 import android.view.Surface
 import android.view.WindowManagerPolicyConstants
-import android.platform.test.annotations.FlakyTest
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.launcher3.tapl.LauncherInstrumentation
 import com.android.server.wm.flicker.FlickerBuilderProvider
@@ -70,7 +70,9 @@ class OpenImeWindowFromFixedOrientationAppTest(private val testSpec: FlickerTest
                 eachRun {
                     // Swiping out the IME activity to home.
                     taplInstrumentation.goHome()
-                    wmHelper.waitForHomeActivityVisible()
+                    wmHelper.StateSyncBuilder()
+                        .withHomeActivityVisible()
+                        .waitForAndVerify()
                 }
             }
             transitions {
@@ -135,4 +137,4 @@ class OpenImeWindowFromFixedOrientationAppTest(private val testSpec: FlickerTest
                     )
         }
     }
-}
\ No newline at end of file
+}
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 94a6639e75acc..dbec3643c39f1 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
@@ -17,10 +17,10 @@
 package com.android.server.wm.flicker.ime
 
 import android.app.Instrumentation
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.view.Surface
 import android.view.WindowManagerPolicyConstants
-import android.platform.test.annotations.FlakyTest
 import androidx.test.filters.RequiresDevice
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.server.wm.flicker.FlickerBuilderProvider
@@ -28,12 +28,12 @@ import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
 import com.android.server.wm.flicker.annotation.Group2
+import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.helpers.ImeAppHelper
 import com.android.server.wm.flicker.navBarLayerIsVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
 import com.android.server.wm.flicker.navBarWindowIsVisible
-import com.android.server.wm.flicker.entireScreenCovered
-import com.android.server.wm.flicker.dsl.FlickerBuilder
 import com.android.server.wm.flicker.statusBarLayerIsVisible
 import com.android.server.wm.flicker.statusBarLayerRotatesScales
 import com.android.server.wm.flicker.statusBarWindowIsVisible
@@ -65,14 +65,14 @@ class OpenImeWindowTest(private val testSpec: FlickerTestParameter) {
                 }
             }
             transitions {
-                testApp.openIME(device, wmHelper)
+                testApp.openIME(wmHelper)
             }
             teardown {
                 eachRun {
-                    testApp.closeIME(device, wmHelper)
+                    testApp.closeIME(wmHelper)
                 }
                 test {
-                    testApp.exit()
+                    testApp.exit(wmHelper)
                 }
             }
         }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowToOverViewTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowToOverViewTest.kt
index e899a750a5df8..eedbd5d3dba8b 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowToOverViewTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/OpenImeWindowToOverViewTest.kt
@@ -17,12 +17,12 @@
 package com.android.server.wm.flicker.ime
 
 import android.app.Instrumentation
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Postsubmit
 import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.RequiresDevice
 import android.view.Surface
 import android.view.WindowManagerPolicyConstants
-import android.platform.test.annotations.FlakyTest
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.server.wm.flicker.FlickerBuilderProvider
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
@@ -30,8 +30,8 @@ import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
 import com.android.server.wm.flicker.annotation.Group4
 import com.android.server.wm.flicker.dsl.FlickerBuilder
-import com.android.server.wm.flicker.helpers.isShellTransitionsEnabled
 import com.android.server.wm.flicker.helpers.ImeAppAutoFocusHelper
+import com.android.server.wm.flicker.helpers.isShellTransitionsEnabled
 import com.android.server.wm.flicker.navBarLayerIsVisible
 import com.android.server.wm.flicker.navBarWindowIsVisible
 import com.android.server.wm.flicker.statusBarLayerIsVisible
@@ -39,8 +39,8 @@ import com.android.server.wm.flicker.statusBarWindowIsVisible
 import com.android.server.wm.traces.common.FlickerComponentName
 import com.android.server.wm.traces.common.WindowManagerConditionsFactory
 import com.android.server.wm.traces.parser.windowmanager.WindowManagerStateHelper
-import org.junit.Assume.assumeTrue
 import org.junit.Assume.assumeFalse
+import org.junit.Assume.assumeTrue
 import org.junit.FixMethodOrder
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -72,12 +72,17 @@ class OpenImeWindowToOverViewTest(private val testSpec: FlickerTestParameter) {
             }
             transitions {
                 device.pressRecentApps()
-                waitForRecentsActivityVisible(wmHelper)
-                waitNavStatusBarVisibility(wmHelper)
+                val builder = wmHelper.StateSyncBuilder()
+                    .withRecentsActivityVisible()
+                waitNavStatusBarVisibility(builder)
+                builder.waitForAndVerify()
             }
             teardown {
                 test {
                     device.pressHome()
+                    wmHelper.StateSyncBuilder()
+                        .withHomeActivityVisible()
+                        .waitForAndVerify()
                     imeTestApp.exit(wmHelper)
                 }
             }
@@ -97,12 +102,12 @@ class OpenImeWindowToOverViewTest(private val testSpec: FlickerTestParameter) {
      *
      * b/227189877
      */
-    private fun waitNavStatusBarVisibility(wmHelper: WindowManagerStateHelper) {
+    private fun waitNavStatusBarVisibility(stateSync: WindowManagerStateHelper.StateSyncBuilder) {
         when {
-            testSpec.isLandscapeOrSeascapeAtStart && !testSpec.isGesturalNavigation ->
-                wmHelper.waitFor(statusBarInvisible)
             testSpec.isLandscapeOrSeascapeAtStart ->
-                wmHelper.waitFor(statusBarInvisible, navBarInvisible)
+                stateSync.add(statusBarInvisible)
+            else ->
+                stateSync.withNavBarStatusBarVisible()
         }
     }
 
@@ -195,25 +200,6 @@ class OpenImeWindowToOverViewTest(private val testSpec: FlickerTestParameter) {
         }
     }
 
-    private fun waitForRecentsActivityVisible(
-        wmHelper: WindowManagerStateHelper
-    ) {
-        val waitMsg = "state of Recents activity to be visible"
-        require(
-                wmHelper.waitFor(waitMsg) {
-                    it.wmState.homeActivity?.let { act ->
-                        it.wmState.isActivityVisible(act.name)
-                    } == true ||
-                            it.wmState.recentsActivity?.let { act ->
-                                it.wmState.isActivityVisible(act.name)
-                            } == true
-                }
-        ) { "Recents activity should be visible" }
-        wmHelper.waitForAppTransitionIdle()
-        // Ensure WindowManagerService wait until all animations have completed
-        instrumentation.uiAutomation.syncInputTransactions()
-    }
-
     companion object {
         /**
          * Creates the test configurations.
@@ -235,4 +221,4 @@ class OpenImeWindowToOverViewTest(private val testSpec: FlickerTestParameter) {
                     )
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/ReOpenImeWindowTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/ReOpenImeWindowTest.kt
index 2b23d53e8a851..7d4724c48bbbc 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/ReOpenImeWindowTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/ReOpenImeWindowTest.kt
@@ -17,33 +17,30 @@
 package com.android.server.wm.flicker.ime
 
 import android.app.Instrumentation
-import android.platform.test.annotations.Presubmit
-import android.view.Display
-import android.view.Surface
 import android.platform.test.annotations.FlakyTest
+import android.platform.test.annotations.Presubmit
+import android.view.Surface
 import androidx.test.filters.RequiresDevice
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.server.wm.flicker.FlickerBuilderProvider
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
-import com.android.server.wm.flicker.LAUNCHER_COMPONENT
 import com.android.server.wm.flicker.annotation.Group2
+import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.helpers.ImeAppAutoFocusHelper
+import com.android.server.wm.flicker.helpers.isShellTransitionsEnabled
 import com.android.server.wm.flicker.helpers.reopenAppFromOverview
 import com.android.server.wm.flicker.helpers.setRotation
 import com.android.server.wm.flicker.navBarLayerIsVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
 import com.android.server.wm.flicker.navBarWindowIsVisible
-import com.android.server.wm.flicker.dsl.FlickerBuilder
-import com.android.server.wm.flicker.entireScreenCovered
-import com.android.server.wm.flicker.helpers.isShellTransitionsEnabled
 import com.android.server.wm.flicker.statusBarLayerIsVisible
 import com.android.server.wm.flicker.statusBarLayerRotatesScales
 import com.android.server.wm.flicker.statusBarWindowIsVisible
-import com.android.server.wm.traces.common.ConditionList
 import com.android.server.wm.traces.common.FlickerComponentName
-import com.android.server.wm.traces.common.WindowManagerConditionsFactory
+import com.android.server.wm.traces.common.FlickerComponentName.Companion.LAUNCHER
 import org.junit.Assume.assumeFalse
 import org.junit.Assume.assumeTrue
 import org.junit.Before
@@ -66,12 +63,6 @@ open class ReOpenImeWindowTest(private val testSpec: FlickerTestParameter) {
     private val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
     private val testApp = ImeAppAutoFocusHelper(instrumentation, testSpec.startRotation)
 
-    private val waitConditionSetup = ConditionList(listOf(
-        WindowManagerConditionsFactory.isAppTransitionIdle(Display.DEFAULT_DISPLAY),
-        WindowManagerConditionsFactory.hasLayersAnimating().negate(),
-        WindowManagerConditionsFactory.isHomeActivityVisible()
-    ))
-
     @Before
     open fun before() {
         assumeFalse(isShellTransitionsEnabled)
@@ -83,21 +74,25 @@ open class ReOpenImeWindowTest(private val testSpec: FlickerTestParameter) {
             setup {
                 test {
                     testApp.launchViaIntent(wmHelper)
-                    testApp.openIME(device, wmHelper)
+                    testApp.openIME(wmHelper)
                 }
                 eachRun {
-                    device.pressRecentApps()
-                    wmHelper.waitFor(waitConditionSetup)
                     this.setRotation(testSpec.startRotation)
+                    device.pressRecentApps()
+                    wmHelper.StateSyncBuilder()
+                        .withRecentsActivityVisible()
+                        .waitForAndVerify()
                 }
             }
             transitions {
                 device.reopenAppFromOverview(wmHelper)
-                wmHelper.waitImeShown()
+                wmHelper.StateSyncBuilder()
+                    .withImeShown()
+                    .waitForAndVerify()
             }
             teardown {
                 test {
-                    testApp.exit()
+                    testApp.exit(wmHelper)
                 }
             }
         }
@@ -128,9 +123,9 @@ open class ReOpenImeWindowTest(private val testSpec: FlickerTestParameter) {
     @Test
     fun launcherWindowBecomesInvisible() {
         testSpec.assertWm {
-            this.isAppWindowVisible(LAUNCHER_COMPONENT)
+            this.isAppWindowVisible(LAUNCHER)
                     .then()
-                    .isAppWindowInvisible(LAUNCHER_COMPONENT)
+                    .isAppWindowInvisible(LAUNCHER)
         }
     }
 
@@ -208,7 +203,7 @@ open class ReOpenImeWindowTest(private val testSpec: FlickerTestParameter) {
     @Test
     fun appLayerReplacesLauncher() {
         testSpec.assertLayers {
-            this.isVisible(LAUNCHER_COMPONENT)
+            this.isVisible(FlickerComponentName.LAUNCHER)
                 .then()
                 .isVisible(FlickerComponentName.SNAPSHOT, isOptional = true)
                 .then()
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 4b268a871fa0b..d1188467cc018 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
@@ -18,17 +18,16 @@ package com.android.server.wm.flicker.ime
 
 import android.app.Instrumentation
 import android.platform.test.annotations.Presubmit
-import android.view.Display
 import android.view.Surface
 import android.view.WindowManagerPolicyConstants
 import androidx.test.filters.RequiresDevice
 import androidx.test.platform.app.InstrumentationRegistry
-import com.android.server.wm.flicker.dsl.FlickerBuilder
 import com.android.server.wm.flicker.FlickerBuilderProvider
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
 import com.android.server.wm.flicker.annotation.Group4
+import com.android.server.wm.flicker.dsl.FlickerBuilder
 import com.android.server.wm.flicker.helpers.ImeAppAutoFocusHelper
 import com.android.server.wm.flicker.helpers.SimpleAppHelper
 import com.android.server.wm.flicker.helpers.WindowUtils
@@ -37,11 +36,8 @@ import com.android.server.wm.flicker.helpers.setRotation
 import com.android.server.wm.flicker.navBarWindowIsVisible
 import com.android.server.wm.flicker.statusBarWindowIsVisible
 import com.android.server.wm.traces.common.FlickerComponentName
-import com.android.server.wm.traces.common.WindowManagerConditionsFactory
-import com.android.server.wm.traces.parser.windowmanager.WindowManagerStateHelper
 import org.junit.Assume
 import org.junit.Before
-
 import org.junit.FixMethodOrder
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -75,32 +71,26 @@ open class SwitchImeWindowsFromGestureNavTest(private val testSpec: FlickerTestP
                 eachRun {
                     this.setRotation(testSpec.startRotation)
                     testApp.launchViaIntent(wmHelper)
-                    val testAppVisible = wmHelper.waitFor(
-                        WindowManagerStateHelper.isAppFullScreen(testApp.component),
-                        WindowManagerConditionsFactory.isAppTransitionIdle(
-                            Display.DEFAULT_DISPLAY))
-                    require(testAppVisible) {
-                        "Expected ${testApp.component.toWindowName()} to be visible"
-                    }
+                    wmHelper.StateSyncBuilder()
+                        .withFullScreenApp(testApp.component)
+                        .waitForAndVerify()
 
                     imeTestApp.launchViaIntent(wmHelper)
-                    val imeAppVisible = wmHelper.waitFor(
-                        WindowManagerStateHelper.isAppFullScreen(imeTestApp.component),
-                        WindowManagerConditionsFactory.isAppTransitionIdle(
-                            Display.DEFAULT_DISPLAY))
-                    require(imeAppVisible) {
-                        "Expected ${imeTestApp.component.toWindowName()} to be visible"
-                    }
+                    wmHelper.StateSyncBuilder()
+                        .withFullScreenApp(imeTestApp.component)
+                        .waitForAndVerify()
 
-                    imeTestApp.openIME(device, wmHelper)
+                    imeTestApp.openIME(wmHelper)
                 }
             }
             teardown {
                 eachRun {
                     device.pressHome()
-                    wmHelper.waitForHomeActivityVisible()
-                    testApp.exit()
-                    imeTestApp.exit()
+                    wmHelper.StateSyncBuilder()
+                        .withHomeActivityVisible()
+                        .waitForAndVerify()
+                    testApp.exit(wmHelper)
+                    imeTestApp.exit(wmHelper)
                 }
             }
             transitions {
@@ -110,8 +100,9 @@ open class SwitchImeWindowsFromGestureNavTest(private val testSpec: FlickerTestP
                 device.swipe(0, displayBounds.bounds.height,
                         displayBounds.bounds.width, displayBounds.bounds.height, 50)
 
-                wmHelper.waitForFullScreenApp(testApp.component)
-                wmHelper.waitForAppTransitionIdle()
+                wmHelper.StateSyncBuilder()
+                    .withFullScreenApp(testApp.component)
+                    .waitForAndVerify()
                 createTag(TAG_IME_INVISIBLE)
             }
             transitions {
@@ -119,7 +110,9 @@ open class SwitchImeWindowsFromGestureNavTest(private val testSpec: FlickerTestP
                 val displayBounds = WindowUtils.getDisplayBounds(testSpec.startRotation)
                 device.swipe(displayBounds.bounds.width, displayBounds.bounds.height,
                         0, displayBounds.bounds.height, 50)
-                wmHelper.waitForFullScreenApp(imeTestApp.component)
+                wmHelper.StateSyncBuilder()
+                    .withFullScreenApp(imeTestApp.component)
+                    .waitForAndVerify()
             }
         }
     }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/ActivitiesTransitionTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/ActivitiesTransitionTest.kt
index cc808a0ce871c..a33f0ea5a074a 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/ActivitiesTransitionTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/ActivitiesTransitionTest.kt
@@ -18,22 +18,20 @@ package com.android.server.wm.flicker.launch
 
 import android.app.Instrumentation
 import android.platform.test.annotations.Presubmit
-import android.view.Display
 import androidx.test.filters.RequiresDevice
 import androidx.test.platform.app.InstrumentationRegistry
-import com.android.server.wm.flicker.entireScreenCovered
+import com.android.launcher3.tapl.LauncherInstrumentation
 import com.android.server.wm.flicker.FlickerBuilderProvider
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
-import com.android.server.wm.flicker.LAUNCHER_COMPONENT
 import com.android.server.wm.flicker.annotation.Group4
 import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.helpers.TwoActivitiesAppHelper
 import com.android.server.wm.flicker.testapp.ActivityOptions
-import com.android.server.wm.traces.common.WindowManagerConditionsFactory
+import com.android.server.wm.traces.common.FlickerComponentName
 import com.android.server.wm.traces.parser.toFlickerComponent
-import com.android.server.wm.traces.parser.windowmanager.WindowManagerStateHelper
 import org.junit.FixMethodOrder
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -64,6 +62,7 @@ import org.junit.runners.Parameterized
 class ActivitiesTransitionTest(val testSpec: FlickerTestParameter) {
     private val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
     private val testApp: TwoActivitiesAppHelper = TwoActivitiesAppHelper(instrumentation)
+    private val tapl = LauncherInstrumentation()
 
     /**
      * Entry point for the test runner. It will use this method to initialize and cache
@@ -74,8 +73,8 @@ class ActivitiesTransitionTest(val testSpec: FlickerTestParameter) {
         return FlickerBuilder(instrumentation).apply {
             setup {
                 test {
+                    tapl.setExpectedRotation(testSpec.startRotation)
                     testApp.launchViaIntent(wmHelper)
-                    wmHelper.waitForFullScreenApp(testApp.component)
                 }
             }
             teardown {
@@ -85,11 +84,10 @@ class ActivitiesTransitionTest(val testSpec: FlickerTestParameter) {
             }
             transitions {
                 testApp.openSecondActivity(device, wmHelper)
-                device.pressBack()
-                val firstActivityVisible = wmHelper.waitFor(
-                    WindowManagerConditionsFactory.isAppTransitionIdle(Display.DEFAULT_DISPLAY),
-                    WindowManagerStateHelper.isAppFullScreen(testApp.component))
-                require(firstActivityVisible) { "Expected ${testApp.component} to be visible" }
+                tapl.pressBack()
+                wmHelper.StateSyncBuilder()
+                    .withFullScreenApp(testApp.component)
+                    .waitForAndVerify()
             }
         }
     }
@@ -125,7 +123,7 @@ class ActivitiesTransitionTest(val testSpec: FlickerTestParameter) {
     fun entireScreenCovered() = testSpec.entireScreenCovered()
 
     /**
-     * Checks that the [LAUNCHER_COMPONENT] window is not on top. The launcher cannot be
+     * Checks that the [FlickerComponentName.LAUNCHER] window is not on top. The launcher cannot be
      * asserted with `isAppWindowVisible` because it contains 2 windows with the exact same name,
      * and both are never simultaneously visible
      */
@@ -133,17 +131,17 @@ class ActivitiesTransitionTest(val testSpec: FlickerTestParameter) {
     @Test
     fun launcherWindowNotOnTop() {
         testSpec.assertWm {
-            this.isAppWindowNotOnTop(LAUNCHER_COMPONENT)
+            this.isAppWindowNotOnTop(FlickerComponentName.LAUNCHER)
         }
     }
 
     /**
-     * Checks that the [LAUNCHER_COMPONENT] layer is never visible during the transition
+     * Checks that the [FlickerComponentName.LAUNCHER] layer is never visible during the transition
      */
     @Presubmit
     @Test
     fun launcherLayerNotVisible() {
-        testSpec.assertLayers { this.isInvisible(LAUNCHER_COMPONENT) }
+        testSpec.assertLayers { this.isInvisible(FlickerComponentName.LAUNCHER) }
     }
 
     companion object {
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdFromIcon.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdFromIcon.kt
new file mode 100644
index 0000000000000..c92704464817b
--- /dev/null
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdFromIcon.kt
@@ -0,0 +1,179 @@
+/*
+ * 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.server.wm.flicker.launch
+
+import android.platform.test.annotations.Postsubmit
+import android.platform.test.annotations.RequiresDevice
+import android.view.Surface
+import com.android.server.wm.flicker.FlickerParametersRunnerFactory
+import com.android.server.wm.flicker.FlickerTestParameter
+import com.android.server.wm.flicker.FlickerTestParameterFactory
+import com.android.server.wm.flicker.annotation.Group1
+import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.flicker.helpers.setRotation
+import com.android.server.wm.flicker.rules.RemoveAllTasksButHomeRule
+import org.junit.FixMethodOrder
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.MethodSorters
+import org.junit.runners.Parameterized
+
+/**
+ * Test cold launching an app from launcher
+ *
+ * To run this test: `atest FlickerTests:OpenAppColdFromIcon`
+ *
+ * Actions:
+ *     Make sure no apps are running on the device
+ *     Launch an app [testApp] by clicking it's icon on all apps and wait animation to complete
+ *
+ * Notes:
+ *     1. Some default assertions (e.g., nav bar, status bar and screen covered)
+ *        are inherited [OpenAppTransition]
+ *     2. Part of the test setup occurs automatically via
+ *        [com.android.server.wm.flicker.TransitionRunnerWithRules],
+ *        including configuring navigation mode, initial orientation and ensuring no
+ *        apps are running before setup
+ */
+@RequiresDevice
+@RunWith(Parameterized::class)
+@Parameterized.UseParametersRunnerFactory(FlickerParametersRunnerFactory::class)
+@FixMethodOrder(MethodSorters.NAME_ASCENDING)
+@Group1
+class OpenAppColdFromIcon(testSpec: FlickerTestParameter) :
+    OpenAppFromLauncherTransition(testSpec) {
+    /**
+     * Defines the transition used to run the test
+     */
+    override val transition: FlickerBuilder.() -> Unit
+        get() = {
+            super.transition(this)
+            setup {
+                eachRun {
+                    tapl.setExpectedRotation(Surface.ROTATION_0)
+                    RemoveAllTasksButHomeRule.removeAllTasksButHome()
+                    this.setRotation(testSpec.startRotation)
+                }
+            }
+            teardown {
+                eachRun {
+                    testApp.exit(wmHelper)
+                }
+            }
+            transitions {
+                tapl.goHome()
+                    .switchToAllApps()
+                    .getAppIcon(testApp.launcherName)
+                    .launch(testApp.`package`)
+            }
+        }
+
+    @Postsubmit
+    @Test
+    override fun appWindowReplacesLauncherAsTopWindow() =
+        super.appWindowReplacesLauncherAsTopWindow()
+
+    @Postsubmit
+    @Test
+    override fun appLayerBecomesVisible() =
+        super.appLayerBecomesVisible()
+
+    @Postsubmit
+    @Test
+    override fun appLayerReplacesLauncher() =
+        super.appLayerReplacesLauncher()
+
+    @Postsubmit
+    @Test
+    override fun appWindowBecomesTopWindow() =
+        super.appWindowBecomesTopWindow()
+
+    @Postsubmit
+    @Test
+    override fun appWindowBecomesVisible() =
+        super.appWindowBecomesVisible()
+
+    @Postsubmit
+    @Test
+    override fun entireScreenCovered() =
+        super.entireScreenCovered()
+
+    @Postsubmit
+    @Test
+    override fun focusChanges() =
+        super.focusChanges()
+
+    @Postsubmit
+    @Test
+    override fun navBarLayerIsVisible() =
+        super.navBarLayerIsVisible()
+
+    @Postsubmit
+    @Test
+    override fun navBarLayerRotatesAndScales() =
+        super.navBarLayerRotatesAndScales()
+
+    @Postsubmit
+    @Test
+    override fun navBarWindowIsVisible() =
+        super.navBarWindowIsVisible()
+
+    @Postsubmit
+    @Test
+    override fun statusBarLayerRotatesScales() =
+        super.statusBarLayerRotatesScales()
+
+    @Postsubmit
+    @Test
+    override fun statusBarLayerIsVisible() =
+        super.statusBarLayerIsVisible()
+
+    @Postsubmit
+    @Test
+    override fun statusBarWindowIsVisible() =
+        super.statusBarWindowIsVisible()
+
+    @Postsubmit
+    @Test
+    override fun visibleLayersShownMoreThanOneConsecutiveEntry() =
+        super.visibleLayersShownMoreThanOneConsecutiveEntry()
+
+    @Postsubmit
+    @Test
+    override fun visibleWindowsShownMoreThanOneConsecutiveEntry() =
+        super.visibleWindowsShownMoreThanOneConsecutiveEntry()
+
+    @Postsubmit
+    @Test
+    override fun appWindowIsTopWindowAtEnd() =
+        super.appWindowIsTopWindowAtEnd()
+
+    companion object {
+        /**
+         * Creates the test configurations.
+         *
+         * See [FlickerTestParameterFactory.getConfigNonRotationTests] for configuring
+         * repetitions, screen orientation and navigation modes.
+         */
+        @Parameterized.Parameters(name = "{0}")
+        @JvmStatic
+        fun getParams(): Collection<FlickerTestParameter> {
+            return FlickerTestParameterFactory.getInstance()
+                .getConfigNonRotationTests()
+        }
+    }
+}
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdTest.kt
index a25ccfb219489..e52b6c326b739 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdTest.kt
@@ -54,8 +54,8 @@ import org.junit.runners.Parameterized
 @Parameterized.UseParametersRunnerFactory(FlickerParametersRunnerFactory::class)
 @FixMethodOrder(MethodSorters.NAME_ASCENDING)
 @Group1
-open class OpenAppColdTest(testSpec: FlickerTestParameter)
-    : OpenAppFromLauncherTransition(testSpec) {
+open class OpenAppColdTest(testSpec: FlickerTestParameter) :
+    OpenAppFromLauncherTransition(testSpec) {
     /**
      * Defines the transition used to run the test
      */
@@ -75,7 +75,6 @@ open class OpenAppColdTest(testSpec: FlickerTestParameter)
             }
             transitions {
                 testApp.launchViaIntent(wmHelper)
-                wmHelper.waitForFullScreenApp(testApp.component)
             }
         }
 
@@ -85,7 +84,7 @@ open class OpenAppColdTest(testSpec: FlickerTestParameter)
     override fun statusBarLayerRotatesScales() = super.statusBarLayerRotatesScales()
 
     /** {@inheritDoc} */
-    @FlakyTest
+    @FlakyTest(bugId = 206753786)
     @Test
     override fun navBarLayerRotatesAndScales() {
         super.navBarLayerRotatesAndScales()
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLauncherTransition.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLauncherTransition.kt
index c6e92adce8c7f..269d82d7fa5af 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLauncherTransition.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLauncherTransition.kt
@@ -18,7 +18,6 @@ package com.android.server.wm.flicker.launch
 
 import android.platform.test.annotations.Presubmit
 import com.android.server.wm.flicker.FlickerTestParameter
-import com.android.server.wm.flicker.LAUNCHER_COMPONENT
 import com.android.server.wm.flicker.replacesLayer
 import com.android.server.wm.traces.common.FlickerComponentName
 import org.junit.Test
@@ -26,8 +25,8 @@ import org.junit.Test
 /**
  * Base class for app launch tests
  */
-abstract class OpenAppFromLauncherTransition(testSpec: FlickerTestParameter)
-    : OpenAppTransition(testSpec) {
+abstract class OpenAppFromLauncherTransition(testSpec: FlickerTestParameter) :
+    OpenAppTransition(testSpec) {
 
     /**
      * Checks that the focus changes from the launcher to [testApp]
@@ -41,25 +40,25 @@ abstract class OpenAppFromLauncherTransition(testSpec: FlickerTestParameter)
     }
 
     /**
-     * Checks that [LAUNCHER_COMPONENT] layer is visible at the start of the transition, and
-     * is replaced by [testApp], which remains visible until the end
+     * Checks that [FlickerComponentName.LAUNCHER] layer is visible at the start of the transition,
+     * and is replaced by [testApp], which remains visible until the end
      */
     open fun appLayerReplacesLauncher() {
-        testSpec.replacesLayer(LAUNCHER_COMPONENT, testApp.component,
+        testSpec.replacesLayer(FlickerComponentName.LAUNCHER, testApp.component,
                 ignoreEntriesWithRotationLayer = true, ignoreSnapshot = true,
                 ignoreSplashscreen = true)
     }
 
     /**
-     * Checks that [LAUNCHER_COMPONENT] window is visible at the start of the transition, and
-     * is replaced by a snapshot or splash screen (optional), and finally, is replaced by
-     * [testApp], which remains visible until the end
+     * Checks that [FlickerComponentName.LAUNCHER] window is visible at the start of the
+     * transition, and is replaced by a snapshot or splash screen (optional), and finally, is
+     * replaced by [testApp], which remains visible until the end
      */
     @Presubmit
     @Test
     open fun appWindowReplacesLauncherAsTopWindow() {
         testSpec.assertWm {
-            this.isAppWindowOnTop(LAUNCHER_COMPONENT)
+            this.isAppWindowOnTop(FlickerComponentName.LAUNCHER)
                     .then()
                     .isAppWindowOnTop(FlickerComponentName.SNAPSHOT, isOptional = true)
                     .then()
@@ -68,4 +67,4 @@ abstract class OpenAppFromLauncherTransition(testSpec: FlickerTestParameter)
                     .isAppWindowOnTop(testApp.component)
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockNotificationCold.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockNotificationCold.kt
index eee7cf3acbe66..cea97ad0302a8 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockNotificationCold.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockNotificationCold.kt
@@ -64,9 +64,9 @@ open class OpenAppFromLockNotificationCold(testSpec: FlickerTestParameter) :
             setup {
                 eachRun {
                     device.sleep()
-                    wmHelper.waitFor("noAppWindowsOnTop") {
-                        it.wmState.topVisibleAppWindow.isEmpty()
-                    }
+                    wmHelper.StateSyncBuilder()
+                        .withoutTopVisibleAppWindows()
+                        .waitForAndVerify()
                 }
             }
         }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockNotificationWarm.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockNotificationWarm.kt
index b156eb1a0933a..3d5f31709a667 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockNotificationWarm.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockNotificationWarm.kt
@@ -65,9 +65,9 @@ open class OpenAppFromLockNotificationWarm(testSpec: FlickerTestParameter) :
             setup {
                 eachRun {
                     device.sleep()
-                    wmHelper.waitFor("noAppWindowsOnTop") {
-                        it.wmState.topVisibleAppWindow.isEmpty()
-                    }
+                    wmHelper.StateSyncBuilder()
+                        .withoutTopVisibleAppWindows()
+                        .waitForAndVerify()
                 }
             }
         }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockNotificationWithLockOverlayApp.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockNotificationWithLockOverlayApp.kt
index 5c86a4278b8c1..446faa7434505 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockNotificationWithLockOverlayApp.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockNotificationWithLockOverlayApp.kt
@@ -66,12 +66,14 @@ class OpenAppFromLockNotificationWithLockOverlayApp(testSpec: FlickerTestParamet
 
                     // Launch an activity that is shown when the device is locked
                     showWhenLockedApp.launchViaIntent(wmHelper)
-                    wmHelper.waitForFullScreenApp(showWhenLockedApp.component)
+                    wmHelper.StateSyncBuilder()
+                        .withFullScreenApp(showWhenLockedApp.component)
+                        .waitForAndVerify()
 
                     device.sleep()
-                    wmHelper.waitFor("noAppWindowsOnTop") {
-                        it.wmState.topVisibleAppWindow.isEmpty()
-                    }
+                    wmHelper.StateSyncBuilder()
+                        .withoutTopVisibleAppWindows()
+                        .waitForAndVerify()
                 }
             }
 
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockTransition.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockTransition.kt
index 9840520d5dbb3..85024e744dea0 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockTransition.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromLockTransition.kt
@@ -16,19 +16,21 @@
 
 package com.android.server.wm.flicker.launch
 
-import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.FlakyTest
+import android.platform.test.annotations.Postsubmit
+import android.platform.test.annotations.Presubmit
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.dsl.FlickerBuilder
 import com.android.server.wm.flicker.navBarLayerPositionEnd
+import com.android.server.wm.flicker.statusBarLayerPositionEnd
 import com.android.server.wm.traces.common.FlickerComponentName
 import org.junit.Test
 
 /**
  * Base class for app launch tests from lock screen
  */
-abstract class OpenAppFromLockTransition(testSpec: FlickerTestParameter)
-    : OpenAppTransition(testSpec) {
+abstract class OpenAppFromLockTransition(testSpec: FlickerTestParameter) :
+    OpenAppTransition(testSpec) {
 
     /**
      * Defines the transition used to run the test
@@ -39,9 +41,9 @@ abstract class OpenAppFromLockTransition(testSpec: FlickerTestParameter)
             setup {
                 eachRun {
                     device.sleep()
-                    wmHelper.waitFor("noAppWindowsOnTop") {
-                        it.wmState.topVisibleAppWindow.isEmpty()
-                    }
+                    wmHelper.StateSyncBuilder()
+                        .withoutTopVisibleAppWindows()
+                        .waitForAndVerify()
                 }
             }
             teardown {
@@ -51,7 +53,6 @@ abstract class OpenAppFromLockTransition(testSpec: FlickerTestParameter)
             }
             transitions {
                 testApp.launchViaIntent(wmHelper)
-                wmHelper.waitForFullScreenApp(testApp.component)
             }
         }
 
@@ -111,6 +112,16 @@ abstract class OpenAppFromLockTransition(testSpec: FlickerTestParameter)
     @Test
     override fun navBarLayerRotatesAndScales() = testSpec.navBarLayerPositionEnd()
 
+    /**
+     * Checks the position of the status bar at the start and end of the transition
+     *
+     * Differently from the normal usage of this assertion, check only the final state of the
+     * transition because the display is off at the start and the NavBar is never visible
+     */
+    @Postsubmit
+    @Test
+    override fun statusBarLayerRotatesScales() = testSpec.statusBarLayerPositionEnd()
+
     /**
      * Checks that the status bar layer is visible at the end of the trace
      *
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromNotificationCold.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromNotificationCold.kt
index 5022dd8f9bffd..68e8e467b950e 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromNotificationCold.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromNotificationCold.kt
@@ -42,8 +42,8 @@ import org.junit.runners.Parameterized
 @FixMethodOrder(MethodSorters.NAME_ASCENDING)
 @Group1
 @Postsubmit
-open class OpenAppFromNotificationCold(testSpec: FlickerTestParameter)
-    : OpenAppFromNotificationWarm(testSpec) {
+open class OpenAppFromNotificationCold(testSpec: FlickerTestParameter) :
+    OpenAppFromNotificationWarm(testSpec) {
     override val transition: FlickerBuilder.() -> Unit
         get() = {
             super.transition(this)
@@ -52,9 +52,9 @@ open class OpenAppFromNotificationCold(testSpec: FlickerTestParameter)
                 eachRun {
                     // Close the app that posted the notification to trigger a cold start next time
                     // it is open - can't just kill it because that would remove the notification.
-                    taplInstrumentation.goHome()
-                    taplInstrumentation.workspace.switchToOverview()
-                    taplInstrumentation.overview.dismissAllTasks()
+                    tapl.goHome()
+                    tapl.workspace.switchToOverview()
+                    tapl.overview.dismissAllTasks()
                 }
             }
         }
@@ -77,8 +77,8 @@ open class OpenAppFromNotificationCold(testSpec: FlickerTestParameter)
         @Parameterized.Parameters(name = "{0}")
         @JvmStatic
         fun getParams(): Collection<FlickerTestParameter> {
-            return com.android.server.wm.flicker.FlickerTestParameterFactory.getInstance()
+            return FlickerTestParameterFactory.getInstance()
                     .getConfigNonRotationTests(repetitions = 3)
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromNotificationWarm.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromNotificationWarm.kt
index 844311259a065..6d5911152dc03 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromNotificationWarm.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromNotificationWarm.kt
@@ -23,7 +23,6 @@ import android.view.WindowInsets
 import android.view.WindowManager
 import androidx.test.uiautomator.By
 import androidx.test.uiautomator.Until
-import com.android.launcher3.tapl.LauncherInstrumentation
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
@@ -55,8 +54,6 @@ import org.junit.runners.Parameterized
 @Postsubmit
 open class OpenAppFromNotificationWarm(testSpec: FlickerTestParameter) :
     OpenAppTransition(testSpec) {
-    protected val taplInstrumentation = LauncherInstrumentation()
-
     override val testApp: NotificationAppHelper = NotificationAppHelper(instrumentation)
 
     open val openingNotificationsFromLockScreen = false
@@ -70,10 +67,14 @@ open class OpenAppFromNotificationWarm(testSpec: FlickerTestParameter) :
                 }
                 eachRun {
                     testApp.launchViaIntent(wmHelper)
-                    wmHelper.waitForFullScreenApp(testApp.component)
-                    testApp.postNotification(device, wmHelper)
-                    device.pressHome()
-                    wmHelper.waitForAppTransitionIdle()
+                    wmHelper.StateSyncBuilder()
+                        .withFullScreenApp(testApp.component)
+                        .waitForAndVerify()
+                    testApp.postNotification(wmHelper)
+                    tapl.goHome()
+                    wmHelper.StateSyncBuilder()
+                        .withHomeActivityVisible()
+                        .waitForAndVerify()
                 }
             }
 
@@ -106,7 +107,9 @@ open class OpenAppFromNotificationWarm(testSpec: FlickerTestParameter) :
                 instrumentation.uiAutomation.syncInputTransactions()
 
                 // Wait for the app to launch
-                wmHelper.waitForFullScreenApp(testApp.component)
+                wmHelper.StateSyncBuilder()
+                    .withFullScreenApp(testApp.component)
+                    .waitForAndVerify()
             }
 
             teardown {
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromOverviewTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromOverviewTest.kt
index f7532b201b75e..c7883d65856d6 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromOverviewTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromOverviewTest.kt
@@ -19,17 +19,14 @@ package com.android.server.wm.flicker.launch
 import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.RequiresDevice
-import android.view.Display
+import android.view.Surface
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
-import com.android.server.wm.flicker.LAUNCHER_COMPONENT
 import com.android.server.wm.flicker.annotation.Group1
 import com.android.server.wm.flicker.dsl.FlickerBuilder
 import com.android.server.wm.flicker.helpers.isShellTransitionsEnabled
-import com.android.server.wm.flicker.helpers.reopenAppFromOverview
 import com.android.server.wm.flicker.helpers.setRotation
-import com.android.server.wm.traces.common.WindowManagerConditionsFactory
 import org.junit.Assume
 import org.junit.FixMethodOrder
 import org.junit.Test
@@ -61,8 +58,8 @@ import org.junit.runners.Parameterized
 @Parameterized.UseParametersRunnerFactory(FlickerParametersRunnerFactory::class)
 @FixMethodOrder(MethodSorters.NAME_ASCENDING)
 @Group1
-open class OpenAppFromOverviewTest(testSpec: FlickerTestParameter)
-    : OpenAppFromLauncherTransition(testSpec) {
+open class OpenAppFromOverviewTest(testSpec: FlickerTestParameter) :
+    OpenAppFromLauncherTransition(testSpec) {
 
     /**
      * Defines the transition used to run the test
@@ -75,27 +72,25 @@ open class OpenAppFromOverviewTest(testSpec: FlickerTestParameter)
                     testApp.launchViaIntent(wmHelper)
                 }
                 eachRun {
+                    // Can't use tapl.goHome() because of b/235841947
                     device.pressHome()
-                    wmHelper.waitForAppTransitionIdle()
-                    device.pressRecentApps()
-                    wmHelper.waitFor(
-                        WindowManagerConditionsFactory
-                            .isAppTransitionIdle(Display.DEFAULT_DISPLAY),
-                        WindowManagerConditionsFactory.isActivityVisible(LAUNCHER_COMPONENT),
-                        WindowManagerConditionsFactory.hasLayersAnimating().negate()
-                    )
+                    wmHelper.StateSyncBuilder()
+                        .withHomeActivityVisible()
+                        .waitForAndVerify()
+                    // Launcher is always ROTATION_0
+                    tapl.setExpectedRotation(Surface.ROTATION_0)
+                    tapl.workspace.switchToOverview()
+                    wmHelper.StateSyncBuilder()
+                        .withRecentsActivityVisible()
+                        .waitForAndVerify()
                     this.setRotation(testSpec.startRotation)
                 }
             }
             transitions {
-                device.reopenAppFromOverview(wmHelper)
-                wmHelper.waitFor(
-                    WindowManagerConditionsFactory.hasLayersAnimating().negate(),
-                    WindowManagerConditionsFactory.isWMStateComplete(),
-                    WindowManagerConditionsFactory.isLayerVisible(LAUNCHER_COMPONENT).negate(),
-                    WindowManagerConditionsFactory.isActivityVisible(LAUNCHER_COMPONENT).negate()
-                )
-                wmHelper.waitForFullScreenApp(testApp.component)
+                tapl.overview.currentTask.open()
+                wmHelper.StateSyncBuilder()
+                    .withFullScreenApp(testApp.component)
+                    .waitForAndVerify()
             }
         }
 
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppNonResizeableTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppNonResizeableTest.kt
index fd74ea5dd9731..04523398046f6 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppNonResizeableTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppNonResizeableTest.kt
@@ -16,8 +16,8 @@
 
 package com.android.server.wm.flicker.launch
 
-import android.platform.test.annotations.Postsubmit
 import android.platform.test.annotations.FlakyTest
+import android.platform.test.annotations.Postsubmit
 import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.RequiresDevice
 import android.view.Surface
@@ -60,10 +60,9 @@ import org.junit.runners.Parameterized
 @Parameterized.UseParametersRunnerFactory(FlickerParametersRunnerFactory::class)
 @FixMethodOrder(MethodSorters.NAME_ASCENDING)
 @Group1
-open class OpenAppNonResizeableTest(testSpec: FlickerTestParameter)
-    : OpenAppFromLockTransition(testSpec) {
+open class OpenAppNonResizeableTest(testSpec: FlickerTestParameter) :
+    OpenAppFromLockTransition(testSpec) {
     override val testApp = NonResizeableAppHelper(instrumentation)
-    private val colorFadComponent = FlickerComponentName("", "ColorFade BLAST#")
 
     /**
      * Checks that the nav bar layer starts invisible, becomes visible during unlocking animation
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppTransition.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppTransition.kt
index 20e6d0222854a..9c97970b4b915 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppTransition.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppTransition.kt
@@ -19,14 +19,15 @@ package com.android.server.wm.flicker.launch
 import android.app.Instrumentation
 import android.platform.test.annotations.Presubmit
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.tapl.LauncherInstrumentation
 import com.android.server.wm.flicker.FlickerBuilderProvider
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.helpers.SimpleAppHelper
 import com.android.server.wm.flicker.helpers.StandardAppHelper
 import com.android.server.wm.flicker.helpers.setRotation
 import com.android.server.wm.flicker.helpers.wakeUpAndGoToHomeScreen
-import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.navBarLayerIsVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
 import com.android.server.wm.flicker.navBarWindowIsVisible
@@ -42,6 +43,7 @@ import org.junit.Test
 abstract class OpenAppTransition(protected val testSpec: FlickerTestParameter) {
     protected val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
     protected open val testApp: StandardAppHelper = SimpleAppHelper(instrumentation)
+    protected val tapl = LauncherInstrumentation()
 
     /**
      * Defines the transition used to run the test
@@ -49,6 +51,7 @@ abstract class OpenAppTransition(protected val testSpec: FlickerTestParameter) {
     protected open val transition: FlickerBuilder.() -> Unit = {
         setup {
             test {
+                tapl.setExpectedRotation(testSpec.startRotation)
                 device.wakeUpAndGoToHomeScreen()
                 this.setRotation(testSpec.startRotation)
             }
@@ -225,11 +228,23 @@ abstract class OpenAppTransition(protected val testSpec: FlickerTestParameter) {
         testSpec.assertWm {
             this.isAppWindowNotOnTop(testApp.component)
                     .then()
-                    .isAppWindowOnTop(FlickerComponentName.SNAPSHOT, isOptional = true)
-                    .then()
-                    .isAppWindowOnTop(FlickerComponentName.SPLASH_SCREEN, isOptional = true)
-                    .then()
-                    .isAppWindowOnTop(testApp.component)
+                    .isAppWindowOnTop(
+                        testApp.component
+                            .or(FlickerComponentName.SNAPSHOT)
+                            .or(FlickerComponentName.SPLASH_SCREEN)
+                    )
+        }
+    }
+
+    /**
+     * Checks that [testApp] window is not on top at the start of the transition, and then becomes
+     * the top visible window until the end of the transition.
+     */
+    @Presubmit
+    @Test
+    open fun appWindowIsTopWindowAtEnd() {
+        testSpec.assertWmEnd {
+            this.isAppWindowOnTop(testApp.component)
         }
     }
 }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppWarmTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppWarmTest.kt
index 45f339979b15b..88e2b2d8e7b7c 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppWarmTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppWarmTest.kt
@@ -16,9 +16,9 @@
 
 package com.android.server.wm.flicker.launch
 
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.RequiresDevice
-import android.platform.test.annotations.FlakyTest
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
@@ -54,8 +54,8 @@ import org.junit.runners.Parameterized
 @Parameterized.UseParametersRunnerFactory(FlickerParametersRunnerFactory::class)
 @FixMethodOrder(MethodSorters.NAME_ASCENDING)
 @Group1
-open class OpenAppWarmTest(testSpec: FlickerTestParameter)
-    : OpenAppFromLauncherTransition(testSpec) {
+open class OpenAppWarmTest(testSpec: FlickerTestParameter) :
+    OpenAppFromLauncherTransition(testSpec) {
     /**
      * Defines the transition used to run the test
      */
@@ -68,18 +68,19 @@ open class OpenAppWarmTest(testSpec: FlickerTestParameter)
                 }
                 eachRun {
                     device.pressHome()
-                    wmHelper.waitForHomeActivityVisible()
+                    wmHelper.StateSyncBuilder()
+                        .withHomeActivityVisible()
+                        .waitForAndVerify()
                     this.setRotation(testSpec.startRotation)
                 }
             }
             teardown {
-                eachRun {
+                test {
                     testApp.exit(wmHelper)
                 }
             }
             transitions {
                 testApp.launchViaIntent(wmHelper)
-                wmHelper.waitForFullScreenApp(testApp.component)
             }
         }
 
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/TaskTransitionTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/TaskTransitionTest.kt
index 2f546b56f1457..99df1f4b5798f 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/TaskTransitionTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/TaskTransitionTest.kt
@@ -21,11 +21,11 @@ import android.app.WallpaperManager
 import android.platform.test.annotations.Postsubmit
 import androidx.test.filters.RequiresDevice
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.tapl.LauncherInstrumentation
 import com.android.server.wm.flicker.FlickerBuilderProvider
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
-import com.android.server.wm.flicker.LAUNCHER_COMPONENT
 import com.android.server.wm.flicker.annotation.Group4
 import com.android.server.wm.flicker.dsl.FlickerBuilder
 import com.android.server.wm.flicker.entireScreenCovered
@@ -64,10 +64,10 @@ import org.junit.runners.Parameterized
 @Group4
 class TaskTransitionTest(val testSpec: FlickerTestParameter) {
     private val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
+    private val tapl = LauncherInstrumentation()
     private val mTestApp: NewTasksAppHelper = NewTasksAppHelper(instrumentation)
     private val mWallpaper by lazy {
-        getWallpaperPackage(InstrumentationRegistry.getInstrumentation())
-            ?: error("Unable to obtain wallpaper")
+        getWallpaperPackage(instrumentation) ?: error("Unable to obtain wallpaper")
     }
 
     @FlickerBuilderProvider
@@ -76,19 +76,19 @@ class TaskTransitionTest(val testSpec: FlickerTestParameter) {
             setup {
                 eachRun {
                     mTestApp.launchViaIntent(wmHelper)
-                    wmHelper.waitForFullScreenApp(mTestApp.component)
                 }
             }
             teardown {
                 test {
-                    mTestApp.exit()
+                    mTestApp.exit(wmHelper)
                 }
             }
             transitions {
                 mTestApp.openNewTask(device, wmHelper)
-                device.pressBack()
-                wmHelper.waitForAppTransitionIdle()
-                wmHelper.waitForFullScreenApp(mTestApp.component)
+                tapl.pressBack()
+                wmHelper.StateSyncBuilder()
+                    .withFullScreenApp(mTestApp.component)
+                    .waitForAndVerify()
             }
         }
     }
@@ -126,7 +126,7 @@ class TaskTransitionTest(val testSpec: FlickerTestParameter) {
     @Test
     fun launcherWindowIsNeverVisible() {
         testSpec.assertWm {
-            this.isAppWindowInvisible(LAUNCHER_COMPONENT)
+            this.isAppWindowInvisible(FlickerComponentName.LAUNCHER)
         }
     }
 
@@ -138,7 +138,7 @@ class TaskTransitionTest(val testSpec: FlickerTestParameter) {
     @Test
     fun launcherLayerIsNeverVisible() {
         testSpec.assertLayers {
-            this.isInvisible(LAUNCHER_COMPONENT)
+            this.isInvisible(FlickerComponentName.LAUNCHER)
         }
     }
 
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsBackTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsBackTest.kt
index c89e6a44ab6c6..1a712842fec0e 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsBackTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsBackTest.kt
@@ -17,6 +17,7 @@
 package com.android.server.wm.flicker.quickswitch
 
 import android.app.Instrumentation
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.RequiresDevice
 import android.view.Surface
@@ -27,12 +28,10 @@ import com.android.server.wm.flicker.FlickerBuilderProvider
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
-import com.android.server.wm.flicker.LAUNCHER_COMPONENT
 import com.android.server.wm.flicker.annotation.Group1
 import com.android.server.wm.flicker.dsl.FlickerBuilder
 import com.android.server.wm.flicker.helpers.NonResizeableAppHelper
 import com.android.server.wm.flicker.helpers.SimpleAppHelper
-import com.android.server.wm.flicker.helpers.WindowUtils
 import com.android.server.wm.flicker.helpers.isShellTransitionsEnabled
 import com.android.server.wm.flicker.navBarLayerIsVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
@@ -40,6 +39,7 @@ import com.android.server.wm.flicker.navBarWindowIsVisible
 import com.android.server.wm.flicker.statusBarLayerIsVisible
 import com.android.server.wm.flicker.statusBarWindowIsVisible
 import com.android.server.wm.traces.common.FlickerComponentName
+import com.android.server.wm.traces.common.Rect
 import org.junit.Assume
 import org.junit.Before
 import org.junit.FixMethodOrder
@@ -66,13 +66,11 @@ import org.junit.runners.Parameterized
 @Group1
 open class QuickSwitchBetweenTwoAppsBackTest(private val testSpec: FlickerTestParameter) {
     private val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
-    private val taplInstrumentation = LauncherInstrumentation()
+    private val tapl = LauncherInstrumentation()
 
     private val testApp1 = SimpleAppHelper(instrumentation)
     private val testApp2 = NonResizeableAppHelper(instrumentation)
 
-    private val startDisplayBounds = WindowUtils.getDisplayBounds(testSpec.startRotation)
-
     @Before
     open fun before() {
         Assume.assumeFalse(isShellTransitionsEnabled)
@@ -83,28 +81,26 @@ open class QuickSwitchBetweenTwoAppsBackTest(private val testSpec: FlickerTestPa
         return FlickerBuilder(instrumentation).apply {
             setup {
                 test {
-                    taplInstrumentation.setExpectedRotation(testSpec.startRotation)
+                    tapl.setExpectedRotation(testSpec.startRotation)
                 }
-
                 eachRun {
                     testApp1.launchViaIntent(wmHelper)
-                    wmHelper.waitForFullScreenApp(testApp1.component)
-
                     testApp2.launchViaIntent(wmHelper)
-                    wmHelper.waitForFullScreenApp(testApp2.component)
+                    startDisplayBounds = wmHelper.currentState.layerState
+                        .physicalDisplayBounds ?: error("Display not found")
                 }
             }
             transitions {
-                taplInstrumentation.launchedAppState.quickSwitchToPreviousApp()
-                wmHelper.waitForFullScreenApp(testApp1.component)
-                wmHelper.waitForAppTransitionIdle()
-                wmHelper.waitForNavBarStatusBarVisible()
+                tapl.launchedAppState.quickSwitchToPreviousApp()
+                wmHelper.StateSyncBuilder()
+                    .withNavBarStatusBarVisible()
+                    .waitForAndVerify()
             }
 
             teardown {
                 test {
-                    testApp1.exit()
-                    testApp2.exit()
+                    testApp1.exit(wmHelper)
+                    testApp2.exit(wmHelper)
                 }
             }
         }
@@ -251,7 +247,7 @@ open class QuickSwitchBetweenTwoAppsBackTest(private val testSpec: FlickerTestPa
             this.isAppWindowVisible(testApp2.component)
                     .then()
                     // TODO: Do we actually want to test this? Seems too implementation specific...
-                    .isAppWindowVisible(LAUNCHER_COMPONENT, isOptional = true)
+                    .isAppWindowVisible(FlickerComponentName.LAUNCHER, isOptional = true)
                     .then()
                     .isAppWindowVisible(FlickerComponentName.SNAPSHOT, isOptional = true)
                     .then()
@@ -270,7 +266,7 @@ open class QuickSwitchBetweenTwoAppsBackTest(private val testSpec: FlickerTestPa
         testSpec.assertLayers {
             this.isVisible(testApp2.component)
                     .then()
-                    .isVisible(LAUNCHER_COMPONENT, isOptional = true)
+                    .isVisible(FlickerComponentName.LAUNCHER, isOptional = true)
                     .then()
                     .isVisible(FlickerComponentName.SNAPSHOT, isOptional = true)
                     .then()
@@ -297,7 +293,7 @@ open class QuickSwitchBetweenTwoAppsBackTest(private val testSpec: FlickerTestPa
      *
      * NOTE: This doesn't check that the navbar is visible or not.
      */
-    @Presubmit
+    @FlakyTest
     @Test
     fun navbarIsAlwaysInRightPosition() = testSpec.navBarLayerRotatesAndScales()
 
@@ -316,6 +312,8 @@ open class QuickSwitchBetweenTwoAppsBackTest(private val testSpec: FlickerTestPa
     fun statusBarLayerIsAlwaysVisible() = testSpec.statusBarLayerIsVisible()
 
     companion object {
+        private var startDisplayBounds = Rect.EMPTY
+
         @Parameterized.Parameters(name = "{0}")
         @JvmStatic
         fun getParams(): Collection<FlickerTestParameter> {
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsForwardTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsForwardTest.kt
index 725d2c3d818c4..9e43a97a62429 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsForwardTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsForwardTest.kt
@@ -17,6 +17,7 @@
 package com.android.server.wm.flicker.quickswitch
 
 import android.app.Instrumentation
+import android.platform.test.annotations.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.RequiresDevice
 import android.view.Surface
@@ -27,7 +28,6 @@ import com.android.server.wm.flicker.FlickerBuilderProvider
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
-import com.android.server.wm.flicker.LAUNCHER_COMPONENT
 import com.android.server.wm.flicker.annotation.Group1
 import com.android.server.wm.flicker.dsl.FlickerBuilder
 import com.android.server.wm.flicker.helpers.NonResizeableAppHelper
@@ -66,7 +66,7 @@ import org.junit.runners.Parameterized
 @Group1
 open class QuickSwitchBetweenTwoAppsForwardTest(private val testSpec: FlickerTestParameter) {
     private val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
-    private val taplInstrumentation = LauncherInstrumentation()
+    private val tapl = LauncherInstrumentation()
 
     private val testApp1 = SimpleAppHelper(instrumentation)
     private val testApp2 = NonResizeableAppHelper(instrumentation)
@@ -81,33 +81,24 @@ open class QuickSwitchBetweenTwoAppsForwardTest(private val testSpec: FlickerTes
         return FlickerBuilder(instrumentation).apply {
             setup {
                 test {
-                    taplInstrumentation.setExpectedRotation(testSpec.startRotation)
+                    tapl.setExpectedRotation(testSpec.startRotation)
                 }
-
                 eachRun {
                     testApp1.launchViaIntent(wmHelper)
-                    wmHelper.waitForFullScreenApp(testApp1.component)
-
                     testApp2.launchViaIntent(wmHelper)
-                    wmHelper.waitForFullScreenApp(testApp2.component)
-
+                    tapl.launchedAppState.quickSwitchToPreviousApp()
+                    wmHelper.StateSyncBuilder()
+                        .withNavBarStatusBarVisible()
+                        .waitForAndVerify()
                     startDisplayBounds = wmHelper.currentState.layerState
-                        .displays.firstOrNull { !it.isVirtual }
-                        ?.layerStackSpace
-                        ?: error("Display not found")
-
-                    taplInstrumentation.launchedAppState.quickSwitchToPreviousApp()
-
-                    wmHelper.waitForFullScreenApp(testApp1.component)
-                    wmHelper.waitForAppTransitionIdle()
+                        .physicalDisplayBounds ?: error("Display not found")
                 }
             }
             transitions {
-                taplInstrumentation.launchedAppState.quickSwitchToPreviousAppSwipeLeft()
-
-                wmHelper.waitForFullScreenApp(testApp2.component)
-                wmHelper.waitForAppTransitionIdle()
-                wmHelper.waitForNavBarStatusBarVisible()
+                tapl.launchedAppState.quickSwitchToPreviousAppSwipeLeft()
+                wmHelper.StateSyncBuilder()
+                    .withNavBarStatusBarVisible()
+                    .waitForAndVerify()
             }
 
             teardown {
@@ -261,7 +252,7 @@ open class QuickSwitchBetweenTwoAppsForwardTest(private val testSpec: FlickerTes
         testSpec.assertWm {
             this.isAppWindowVisible(testApp1.component)
                     .then()
-                    .isAppWindowVisible(LAUNCHER_COMPONENT, isOptional = true)
+                    .isAppWindowVisible(FlickerComponentName.LAUNCHER, isOptional = true)
                     .then()
                     .isAppWindowVisible(FlickerComponentName.SNAPSHOT, isOptional = true)
                     .then()
@@ -280,7 +271,7 @@ open class QuickSwitchBetweenTwoAppsForwardTest(private val testSpec: FlickerTes
         testSpec.assertLayers {
             this.isVisible(testApp1.component)
                     .then()
-                    .isVisible(LAUNCHER_COMPONENT, isOptional = true)
+                    .isVisible(FlickerComponentName.LAUNCHER, isOptional = true)
                     .then()
                     .isVisible(FlickerComponentName.SNAPSHOT, isOptional = true)
                     .then()
@@ -311,7 +302,7 @@ open class QuickSwitchBetweenTwoAppsForwardTest(private val testSpec: FlickerTes
      *
      * NOTE: This doesn't check that the navbar is visible or not.
      */
-    @Presubmit
+    @FlakyTest
     @Test
     open fun navbarIsAlwaysInRightPosition() {
         testSpec.navBarLayerRotatesAndScales()
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchFromLauncherTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchFromLauncherTest.kt
index cc4a4b2d38aa4..4ea1f1c18b522 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchFromLauncherTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchFromLauncherTest.kt
@@ -27,18 +27,17 @@ import com.android.server.wm.flicker.FlickerBuilderProvider
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
-import com.android.server.wm.flicker.LAUNCHER_COMPONENT
 import com.android.server.wm.flicker.annotation.Group1
 import com.android.server.wm.flicker.dsl.FlickerBuilder
 import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.helpers.SimpleAppHelper
-import com.android.server.wm.flicker.helpers.WindowUtils
 import com.android.server.wm.flicker.navBarLayerIsVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
 import com.android.server.wm.flicker.navBarWindowIsVisible
 import com.android.server.wm.flicker.statusBarLayerIsVisible
 import com.android.server.wm.flicker.statusBarWindowIsVisible
 import com.android.server.wm.traces.common.FlickerComponentName
+import com.android.server.wm.traces.common.Rect
 import org.junit.FixMethodOrder
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -67,8 +66,6 @@ class QuickSwitchFromLauncherTest(private val testSpec: FlickerTestParameter) {
 
     private val testApp = SimpleAppHelper(instrumentation)
 
-    private val startDisplayBounds = WindowUtils.getDisplayBounds(testSpec.startRotation)
-
     @FlickerBuilderProvider
     fun buildFlicker(): FlickerBuilder {
         return FlickerBuilder(instrumentation).apply {
@@ -80,20 +77,26 @@ class QuickSwitchFromLauncherTest(private val testSpec: FlickerTestParameter) {
                 eachRun {
                     testApp.launchViaIntent(wmHelper)
                     device.pressHome()
-                    wmHelper.waitForHomeActivityVisible()
-                    wmHelper.waitForWindowSurfaceDisappeared(testApp.component)
+                    wmHelper.StateSyncBuilder()
+                        .withHomeActivityVisible()
+                        .withWindowSurfaceDisappeared(testApp.component)
+                        .waitForAndVerify()
+
+                    startDisplayBounds = wmHelper.currentState.layerState
+                        .physicalDisplayBounds ?: error("Display not found")
                 }
             }
             transitions {
                 taplInstrumentation.workspace.quickSwitchToPreviousApp()
-                wmHelper.waitForFullScreenApp(testApp.component)
-                wmHelper.waitForAppTransitionIdle()
-                wmHelper.waitForNavBarStatusBarVisible()
+                wmHelper.StateSyncBuilder()
+                    .withFullScreenApp(testApp.component)
+                    .withNavBarStatusBarVisible()
+                    .waitForAndVerify()
             }
 
             teardown {
                 eachRun {
-                    testApp.exit()
+                    testApp.exit(wmHelper)
                 }
             }
         }
@@ -154,7 +157,7 @@ class QuickSwitchFromLauncherTest(private val testSpec: FlickerTestParameter) {
     @Test
     fun startsWithLauncherWindowsCoverFullScreen() {
         testSpec.assertWmStart {
-            this.frameRegion(LAUNCHER_COMPONENT).coversExactly(startDisplayBounds)
+            this.frameRegion(FlickerComponentName.LAUNCHER).coversExactly(startDisplayBounds)
         }
     }
 
@@ -166,7 +169,7 @@ class QuickSwitchFromLauncherTest(private val testSpec: FlickerTestParameter) {
     @Test
     fun startsWithLauncherLayersCoverFullScreen() {
         testSpec.assertLayersStart {
-            this.visibleRegion(LAUNCHER_COMPONENT).coversExactly(startDisplayBounds)
+            this.visibleRegion(FlickerComponentName.LAUNCHER).coversExactly(startDisplayBounds)
         }
     }
 
@@ -177,7 +180,7 @@ class QuickSwitchFromLauncherTest(private val testSpec: FlickerTestParameter) {
     @Test
     fun startsWithLauncherBeingOnTop() {
         testSpec.assertWmStart {
-            this.isAppWindowOnTop(LAUNCHER_COMPONENT)
+            this.isAppWindowOnTop(FlickerComponentName.LAUNCHER)
         }
     }
 
@@ -229,9 +232,9 @@ class QuickSwitchFromLauncherTest(private val testSpec: FlickerTestParameter) {
     @Test
     fun launcherWindowBecomesAndStaysInvisible() {
         testSpec.assertWm {
-            this.isAppWindowOnTop(LAUNCHER_COMPONENT)
+            this.isAppWindowOnTop(FlickerComponentName.LAUNCHER)
                     .then()
-                    .isAppWindowNotOnTop(LAUNCHER_COMPONENT)
+                    .isAppWindowNotOnTop(FlickerComponentName.LAUNCHER)
         }
     }
 
@@ -243,9 +246,9 @@ class QuickSwitchFromLauncherTest(private val testSpec: FlickerTestParameter) {
     @Test
     fun launcherLayerBecomesAndStaysInvisible() {
         testSpec.assertLayers {
-            this.isVisible(LAUNCHER_COMPONENT)
+            this.isVisible(FlickerComponentName.LAUNCHER)
                     .then()
-                    .isInvisible(LAUNCHER_COMPONENT)
+                    .isInvisible(FlickerComponentName.LAUNCHER)
         }
     }
 
@@ -257,7 +260,7 @@ class QuickSwitchFromLauncherTest(private val testSpec: FlickerTestParameter) {
     @Test
     fun appWindowIsVisibleOnceLauncherWindowIsInvisible() {
         testSpec.assertWm {
-            this.isAppWindowOnTop(LAUNCHER_COMPONENT)
+            this.isAppWindowOnTop(FlickerComponentName.LAUNCHER)
                     .then()
                     .isAppWindowVisible(FlickerComponentName.SNAPSHOT, isOptional = true)
                     .then()
@@ -273,7 +276,7 @@ class QuickSwitchFromLauncherTest(private val testSpec: FlickerTestParameter) {
     @Test
     fun appLayerIsVisibleOnceLauncherLayerIsInvisible() {
         testSpec.assertLayers {
-            this.isVisible(LAUNCHER_COMPONENT)
+            this.isVisible(FlickerComponentName.LAUNCHER)
                     .then()
                     .isVisible(FlickerComponentName.SNAPSHOT, isOptional = true)
                     .then()
@@ -326,6 +329,8 @@ class QuickSwitchFromLauncherTest(private val testSpec: FlickerTestParameter) {
     fun screenIsAlwaysFilled() = testSpec.entireScreenCovered()
 
     companion object {
+        private var startDisplayBounds = Rect.EMPTY
+
         @Parameterized.Parameters(name = "{0}")
         @JvmStatic
         fun getParams(): Collection<FlickerTestParameter> {
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/RotationTransition.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/RotationTransition.kt
index 0becadf630e16..586be06d1a298 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/RotationTransition.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/RotationTransition.kt
@@ -22,12 +22,12 @@ import androidx.test.platform.app.InstrumentationRegistry
 import com.android.server.wm.flicker.FlickerBuilderProvider
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.flicker.helpers.StandardAppHelper
 import com.android.server.wm.flicker.helpers.setRotation
 import com.android.server.wm.flicker.navBarLayerIsVisible
 import com.android.server.wm.flicker.navBarLayerRotatesAndScales
 import com.android.server.wm.flicker.navBarWindowIsVisible
-import com.android.server.wm.flicker.entireScreenCovered
 import com.android.server.wm.traces.common.FlickerComponentName
 import org.junit.Test
 
@@ -47,7 +47,7 @@ abstract class RotationTransition(protected val testSpec: FlickerTestParameter)
         }
         teardown {
             test {
-                testApp.exit()
+                testApp.exit(wmHelper)
             }
         }
         transitions {