From 2490b439d0e5bf08d04b02382cf43fd9007383ac Mon Sep 17 00:00:00 2001
From: Chris Li <lihongyu@google.com>
Date: Tue, 4 Jan 2022 14:59:45 +0800
Subject: [PATCH] Fix ExitPipWithDismissButtonTest with Shell transition

Fix: 210965692
Test: atest WMShellFlickerTests:ExitPipWithDismissButtonTest
Change-Id: I828095f55e7400b661513b853372a3bef5466663
---
 .../wm/shell/flicker/pip/ExitPipTransition.kt | 28 +++++++++++++++----
 .../pip/ExitPipWithDismissButtonTest.kt       | 15 +---------
 2 files changed, 24 insertions(+), 19 deletions(-)

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 6c9fed9dc19b5..9a220070db01c 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
@@ -22,6 +22,7 @@ import androidx.test.filters.FlakyTest
 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 org.junit.Test
 
@@ -51,11 +52,28 @@ abstract class ExitPipTransition(testSpec: FlickerTestParameter) : PipTransition
     @Presubmit
     @Test
     open fun pipWindowBecomesInvisible() {
-        testSpec.assertWm {
-            this.invoke("hasPipWindow") {
-                it.isPinned(pipApp.component).isAppWindowVisible(pipApp.component)
-            }.then().invoke("!hasPipWindow") {
-                it.isNotPinned(pipApp.component).isAppWindowInvisible(pipApp.component)
+        if (isShellTransitionsEnabled) {
+            // When Shell transition is enabled, we change the windowing mode at start, but
+            // update the visibility after the transition is finished, so we can't check isNotPinned
+            // and isAppWindowInvisible in the same assertion block.
+            testSpec.assertWm {
+                this.invoke("hasPipWindow") {
+                    it.isPinned(pipApp.component)
+                            .isAppWindowVisible(pipApp.component)
+                            .isAppWindowOnTop(pipApp.component)
+                }.then().invoke("!hasPipWindow") {
+                    it.isNotPinned(pipApp.component)
+                            .isAppWindowNotOnTop(pipApp.component)
+                }
+            }
+            testSpec.assertWmEnd { isAppWindowInvisible(pipApp.component) }
+        } else {
+            testSpec.assertWm {
+                this.invoke("hasPipWindow") {
+                    it.isPinned(pipApp.component).isAppWindowVisible(pipApp.component)
+                }.then().invoke("!hasPipWindow") {
+                    it.isNotPinned(pipApp.component).isAppWindowInvisible(pipApp.component)
+                }
             }
         }
     }
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipWithDismissButtonTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipWithDismissButtonTest.kt
index c75076d00574d..8adebb8f28c97 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipWithDismissButtonTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExitPipWithDismissButtonTest.kt
@@ -24,10 +24,7 @@ import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
 import com.android.server.wm.flicker.annotation.Group3
 import com.android.server.wm.flicker.dsl.FlickerBuilder
-import com.android.server.wm.flicker.helpers.isShellTransitionsEnabled
 import com.android.server.wm.flicker.rules.WMFlickerServiceRuleForTestSpec
-import org.junit.Assume.assumeFalse
-import org.junit.Before
 import org.junit.FixMethodOrder
 import org.junit.Rule
 import org.junit.Test
@@ -88,20 +85,10 @@ class ExitPipWithDismissButtonTest(testSpec: FlickerTestParameter) : ExitPipTran
         flickerRule.checkFlakyAssertions()
     }
 
-    @Before
-    fun onBefore() {
-        // This CUJ don't work in shell transitions because of b/204570898 b/204562589
-        assumeFalse(isShellTransitionsEnabled)
-    }
-
     /** {@inheritDoc}  */
     @FlakyTest(bugId = 206753786)
     @Test
-    override fun statusBarLayerRotatesScales() {
-        // This test doesn't work in shell transitions because of b/206753786
-        assumeFalse(isShellTransitionsEnabled)
-        super.statusBarLayerRotatesScales()
-    }
+    override fun statusBarLayerRotatesScales() = super.statusBarLayerRotatesScales()
 
     companion object {
         /**