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 5a96a7c8cbd97..f4dd7decb1b71 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
@@ -26,8 +26,6 @@ import com.android.server.wm.flicker.helpers.SYSTEMUI_PACKAGE
 import com.android.server.wm.traces.parser.windowmanager.WindowManagerStateHelper
 import com.android.wm.shell.flicker.pip.tv.closeTvPipWindow
 import com.android.wm.shell.flicker.pip.tv.isFocusedOrHasFocusedChild
-import com.android.wm.shell.flicker.pip.waitPipWindowGone
-import com.android.wm.shell.flicker.pip.waitPipWindowShown
 import com.android.wm.shell.flicker.testapp.Components
 
 class PipAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
@@ -64,7 +62,7 @@ class PipAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
         stringExtras: Map<String, String>
     ) {
         super.launchViaIntent(wmHelper, expectedWindowName, action, stringExtras)
-        wmHelper.waitPipWindowShown()
+        wmHelper.waitFor { it.wmState.hasPipWindow() }
     }
 
     private fun focusOnObject(selector: BySelector): Boolean {
@@ -86,7 +84,7 @@ class PipAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
         clickObject(ENTER_PIP_BUTTON_ID)
 
         // Wait on WMHelper or simply wait for 3 seconds
-        wmHelper?.waitPipWindowShown() ?: SystemClock.sleep(3_000)
+        wmHelper?.waitFor { it.wmState.hasPipWindow() } ?: SystemClock.sleep(3_000)
     }
 
     fun clickStartMediaSessionButton() {
@@ -139,7 +137,7 @@ class PipAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
         }
 
         // Wait for animation to complete.
-        wmHelper.waitPipWindowGone()
+        wmHelper.waitFor { !it.wmState.hasPipWindow() }
         wmHelper.waitForHomeActivityVisible()
     }
 
@@ -169,7 +167,7 @@ class PipAppHelper(instrumentation: Instrumentation) : BaseAppHelper(
         val windowRect = windowRegion.bounds
         uiDevice.click(windowRect.centerX(), windowRect.centerY())
         uiDevice.click(windowRect.centerX(), windowRect.centerY())
-        wmHelper.waitPipWindowGone()
+        wmHelper.waitFor { !it.wmState.hasPipWindow() }
         wmHelper.waitForAppTransitionIdle()
     }
 
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 dd6195c50f229..2beec2e8c1ed6 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
@@ -75,7 +75,7 @@ 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.waitPipWindowShown()
+                wmHelper.waitFor { it.wmState.hasPipWindow() }
                 wmHelper.waitForAppTransitionIdle()
             }
         }
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/Extensions.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/Extensions.kt
index 96eb66c3cc287..0037059e2c517 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/Extensions.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/Extensions.kt
@@ -16,62 +16,14 @@
 
 package com.android.wm.shell.flicker.pip
 
-import android.app.WindowConfiguration
 import android.content.ComponentName
-import com.android.server.wm.flicker.traces.windowmanager.WindowManagerStateSubject
 import com.android.server.wm.traces.common.windowmanager.WindowManagerState
 import com.android.server.wm.traces.parser.toWindowName
-import com.android.server.wm.traces.parser.windowmanager.WindowManagerStateHelper
-import com.google.common.truth.Truth
-
-inline val WindowManagerState.pinnedWindows
-    get() = visibleWindows
-        .filter { it.windowingMode == WindowConfiguration.WINDOWING_MODE_PINNED }
-
-/**
- * Checks if the state has any window in PIP mode
- */
-fun WindowManagerState.hasPipWindow(): Boolean = pinnedWindows.isNotEmpty()
 
 /**
  * Checks that an activity [activity] is in PIP mode
  */
 fun WindowManagerState.isInPipMode(activity: ComponentName): Boolean {
     val windowName = activity.toWindowName()
-    return pinnedWindows.any { it.title == windowName }
+    return isInPipMode(windowName)
 }
-
-/**
- * Asserts that an activity [activity] exists and is in PIP mode
- */
-fun WindowManagerStateSubject.isInPipMode(
-    activity: ComponentName
-): WindowManagerStateSubject = apply {
-    val windowName = activity.toWindowName()
-    contains(windowName)
-    val pinnedWindows = wmState.pinnedWindows
-        .map { it.title }
-    Truth.assertWithMessage("Window not in PIP mode")
-        .that(pinnedWindows)
-        .contains(windowName)
-}
-
-/**
- * Waits until the state has a window in PIP mode, i.e., with
- * windowingMode = WindowConfiguration.WINDOWING_MODE_PINNED
- */
-fun WindowManagerStateHelper.waitPipWindowShown(): Boolean =
-    waitFor("PIP window shown") {
-        val result = it.wmState.hasPipWindow()
-        result
-    }
-
-/**
- * Waits until the state doesn't have a window in PIP mode, i.e., with
- * windowingMode = WindowConfiguration.WINDOWING_MODE_PINNED
- */
-fun WindowManagerStateHelper.waitPipWindowGone(): Boolean =
-    waitFor("PIP window gone") {
-        val result = !it.wmState.hasPipWindow()
-        result
-    }
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 7dc7e7d25b798..f8856827ff425 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
@@ -131,13 +131,13 @@ abstract class PipTransition(protected val testSpec: FlickerTestParameter) {
                     removeAllTasksButHome()
                     if (!eachRun) {
                         pipApp.launchViaIntent(wmHelper, stringExtras = stringExtras)
-                        wmHelper.waitPipWindowShown()
+                        wmHelper.waitFor { it.wmState.hasPipWindow() }
                     }
                 }
                 eachRun {
                     if (eachRun) {
                         pipApp.launchViaIntent(wmHelper, stringExtras = stringExtras)
-                        wmHelper.waitPipWindowShown()
+                        wmHelper.waitFor { it.wmState.hasPipWindow() }
                     }
                 }
             }