Wait on aspect ratio change for flicker test

Flag: EXEMPT bugfix
Bug: 358278071
Test: atest WMShellFlickerTestsPip3:PipAspectRatioChangeTest
Change-Id: I601e763b10c291ca40ba806f2f2d6af4000c4494
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipAspectRatioChangeTest.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipAspectRatioChangeTest.kt
index 5b7521a..429774f 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipAspectRatioChangeTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipAspectRatioChangeTest.kt
@@ -22,7 +22,6 @@
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
-import androidx.test.filters.FlakyTest
 import com.android.wm.shell.flicker.pip.common.PipTransition
 import org.junit.FixMethodOrder
 import org.junit.Test
@@ -36,7 +35,7 @@
 @FixMethodOrder(MethodSorters.NAME_ASCENDING)
 class PipAspectRatioChangeTest(flicker: LegacyFlickerTest) : PipTransition(flicker) {
     override val thisTransition: FlickerBuilder.() -> Unit = {
-        transitions { pipApp.changeAspectRatio() }
+        transitions { pipApp.changeAspectRatio(wmHelper) }
     }
 
     @Presubmit
@@ -46,22 +45,6 @@
         flicker.assertLayersEnd { this.visibleRegion(pipApp).isSameAspectRatio(1, 2) }
     }
 
-    @FlakyTest(bugId = 358278071)
-    override fun hasAtMostOnePipDismissOverlayWindow() =
-        super.hasAtMostOnePipDismissOverlayWindow()
-
-    @FlakyTest(bugId = 358278071)
-    override fun statusBarLayerPositionAtStartAndEnd() =
-        super.statusBarLayerPositionAtStartAndEnd()
-
-    @FlakyTest(bugId = 358278071)
-    override fun taskBarWindowIsAlwaysVisible() =
-        super.taskBarWindowIsAlwaysVisible()
-
-    @FlakyTest(bugId = 358278071)
-    override fun visibleWindowsShownMoreThanOneConsecutiveEntry() =
-        super.visibleWindowsShownMoreThanOneConsecutiveEntry()
-
     companion object {
         /**
          * Creates the test configurations.
diff --git a/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/PipAppHelper.kt b/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/PipAppHelper.kt
index 43fd57b..931e4f8 100644
--- a/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/PipAppHelper.kt
+++ b/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/PipAppHelper.kt
@@ -269,9 +269,23 @@
     /** Expand the PIP window back to full screen via intent and wait until the app is visible */
     fun exitPipToFullScreenViaIntent(wmHelper: WindowManagerStateHelper) = launchViaIntent(wmHelper)
 
-    fun changeAspectRatio() {
+    fun changeAspectRatio(wmHelper: WindowManagerStateHelper) {
         val intent = Intent("com.android.wm.shell.flicker.testapp.ASPECT_RATIO")
         context.sendBroadcast(intent)
+        // Wait on WMHelper on size change upon aspect ratio change
+        val windowRect = getWindowRect(wmHelper)
+        wmHelper
+            .StateSyncBuilder()
+            .add("pipAspectRatioChanged") {
+                val pipAppWindow =
+                    it.wmState.visibleWindows.firstOrNull { window ->
+                        this.windowMatchesAnyOf(window)
+                    }
+                        ?: return@add false
+                val pipRegion = pipAppWindow.frameRegion
+                return@add pipRegion != Region(windowRect)
+            }
+            .waitForAndVerify()
     }
 
     fun clickEnterPipButton(wmHelper: WindowManagerStateHelper) {