diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/PipDragThenSnapTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/PipDragThenSnapTest.kt
index 5cd66abf5012f..53ce3936fbe47 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/PipDragThenSnapTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/PipDragThenSnapTest.kt
@@ -41,6 +41,9 @@ import org.junit.runners.Parameterized
 @Parameterized.UseParametersRunnerFactory(FlickerParametersRunnerFactory::class)
 @FixMethodOrder(MethodSorters.NAME_ASCENDING)
 class PipDragThenSnapTest(flicker: FlickerTest) : PipTransition(flicker){
+    // represents the direction in which the pip window should be snapping
+    private var willSnapRight: Boolean = true
+
     override val transition: FlickerBuilder.() -> Unit
         get() = {
             val stringExtras: Map<String, String> =
@@ -55,9 +58,16 @@ class PipDragThenSnapTest(flicker: FlickerTest) : PipTransition(flicker){
                 RemoveAllTasksButHomeRule.removeAllTasksButHome()
                 pipApp.launchViaIntentAndWaitForPip(wmHelper, stringExtras = stringExtras)
 
-                val initRegion = pipApp.dragPipWindowAwayFromEdge(wmHelper, 50)
+                // get the initial region bounds and cache them
+                val initRegion = pipApp.getWindowRect(wmHelper)
                 startBounds
-                    .set(initRegion.left, initRegion.top, initRegion.right, initRegion.bottom)
+                        .set(initRegion.left, initRegion.top, initRegion.right, initRegion.bottom)
+
+                // drag the pip window away from the edge
+                pipApp.dragPipWindowAwayFromEdge(wmHelper, 50)
+
+                // determine the direction in which the snapping should occur
+                willSnapRight = pipApp.isCloserToRightEdge(wmHelper)
             }
             transitions {
                 // continue the transition until the PIP snaps
@@ -65,14 +75,20 @@ class PipDragThenSnapTest(flicker: FlickerTest) : PipTransition(flicker){
             }
         }
 
-    /** Checks that the visible region area of [pipApp] always moves right during the animation. */
+    /**
+     * Checks that the visible region area of [pipApp] moves to closest edge during the animation.
+     */
     @Postsubmit
     @Test
-    fun pipLayerMovesRight() {
+    fun pipLayerMovesToClosestEdge() {
         flicker.assertLayers {
             val pipLayerList = layers { pipApp.layerMatchesAnyOf(it) && it.isVisible }
             pipLayerList.zipWithNext { previous, current ->
-                current.visibleRegion.isToTheRight(previous.visibleRegion.region)
+                if (willSnapRight) {
+                    current.visibleRegion.isToTheRight(previous.visibleRegion.region)
+                } else {
+                    previous.visibleRegion.isToTheRight(current.visibleRegion.region)
+                }
             }
         }
     }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/PipAppHelper.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/PipAppHelper.kt
index 6e4e937240550..de57d06a56195 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/PipAppHelper.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/PipAppHelper.kt
@@ -61,9 +61,8 @@ open class PipAppHelper(instrumentation: Instrumentation) :
      * Drags the PIP window away from the screen edge while not crossing the display center.
      *
      * @throws IllegalStateException if default display bounds are not available
-     * @return initial bounds of the PIP window
      */
-    fun dragPipWindowAwayFromEdge(wmHelper: WindowManagerStateHelper, steps: Int): Rect {
+    fun dragPipWindowAwayFromEdge(wmHelper: WindowManagerStateHelper, steps: Int) {
         val initWindowRect = getWindowRect(wmHelper).clone()
 
         // initial pointer at the center of the window
@@ -83,8 +82,6 @@ open class PipAppHelper(instrumentation: Instrumentation) :
 
         // drag the window to the left but not beyond the center of the display
         uiDevice.drag(startX, y, endX, y, steps)
-
-        return initWindowRect
     }
 
     /**
@@ -92,7 +89,7 @@ open class PipAppHelper(instrumentation: Instrumentation) :
      *
      * @throws IllegalStateException if default display bounds are not available
      */
-    private fun isCloserToRightEdge(wmHelper: WindowManagerStateHelper): Boolean {
+    fun isCloserToRightEdge(wmHelper: WindowManagerStateHelper): Boolean {
         val windowRect = getWindowRect(wmHelper)
 
         val displayRect = wmHelper.currentState.wmState.getDefaultDisplay()?.displayRect
@@ -267,7 +264,10 @@ open class PipAppHelper(instrumentation: Instrumentation) :
         closePipWindow(WindowManagerStateHelper(mInstrumentation))
     }
 
-    private fun getWindowRect(wmHelper: WindowManagerStateHelper): Rect {
+    /**
+     * Returns the pip window bounds.
+     */
+    fun getWindowRect(wmHelper: WindowManagerStateHelper): Rect {
         val windowRegion = wmHelper.getWindowRegion(this)
         require(!windowRegion.isEmpty) { "Unable to find a PIP window in the current state" }
         return windowRegion.bounds
@@ -354,7 +354,7 @@ open class PipAppHelper(instrumentation: Instrumentation) :
     /**
      * Waits until the PIP window snaps horizontally to the provided bounds.
      *
-     * @param finalRightX the final x coordinate of the right edge of the pip window
+     * @param finalBounds the bounds to wait for PIP window to snap to
      */
     fun waitForPipToSnapTo(wmHelper: WindowManagerStateHelper, finalBounds: android.graphics.Rect) {
         wmHelper
@@ -369,6 +369,7 @@ open class PipAppHelper(instrumentation: Instrumentation) :
                 return@add pipRegionBounds.left == finalBounds.left &&
                     pipRegionBounds.right == finalBounds.right
             }
+            .add(ConditionsFactory.isWMStateComplete())
             .waitForAndVerify()
     }