Add closeAllApps scenario and refactor

- Add new test to close multiple apps in desktop mode.
- Refactor flicker scenarios into their own class to avoid duplication.
- Replace DesktopModeUtils with DesktopModeAppHelper wrapper.

Test: atest com.android.wm.shell.flicker.service.desktopmode.flicker
Bug: 316318390
Change-Id: I60f37fb507850daddc3cb91529c272e9f0dbbb44

10 files changed, 432 insertions, 186 deletions

diff --git a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/CloseAllAppWithAppHeaderExitLandscape.kt b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/CloseAllAppWithAppHeaderExitLandscape.kt
new file mode 100644
index 000000000000..5563bb9fa934
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/CloseAllAppWithAppHeaderExitLandscape.kt
@@ -0,0 +1,47 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.flicker.service.desktopmode.flicker
+
+import android.tools.Rotation
+import android.tools.flicker.FlickerConfig
+import android.tools.flicker.annotation.ExpectedScenarios
+import android.tools.flicker.annotation.FlickerConfigProvider
+import android.tools.flicker.config.FlickerConfig
+import android.tools.flicker.config.FlickerServiceConfig
+import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
+import com.android.wm.shell.flicker.service.desktopmode.flicker.DesktopModeFlickerScenarios.Companion.CLOSE_APP
+import com.android.wm.shell.flicker.service.desktopmode.flicker.DesktopModeFlickerScenarios.Companion.CLOSE_LAST_APP
+import com.android.wm.shell.flicker.service.desktopmode.scenarios.CloseAllAppsWithAppHeaderExit
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(FlickerServiceJUnit4ClassRunner::class)
+class CloseAllAppWithAppHeaderExitLandscape : CloseAllAppsWithAppHeaderExit(Rotation.ROTATION_90) {
+    @ExpectedScenarios(["CLOSE_APP", "CLOSE_LAST_APP"])
+    @Test
+    override fun closeAllAppsInDesktop() = super.closeAllAppsInDesktop()
+
+    companion object {
+        @JvmStatic
+        @FlickerConfigProvider
+        fun flickerConfigProvider(): FlickerConfig =
+            FlickerConfig()
+                .use(FlickerServiceConfig.DEFAULT)
+                .use(CLOSE_APP)
+                .use(CLOSE_LAST_APP)
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/CloseAllAppWithAppHeaderExitPortrait.kt b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/CloseAllAppWithAppHeaderExitPortrait.kt
new file mode 100644
index 000000000000..3d16d2219c78
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/CloseAllAppWithAppHeaderExitPortrait.kt
@@ -0,0 +1,47 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.flicker.service.desktopmode.flicker
+
+import android.tools.Rotation
+import android.tools.flicker.FlickerConfig
+import android.tools.flicker.annotation.ExpectedScenarios
+import android.tools.flicker.annotation.FlickerConfigProvider
+import android.tools.flicker.config.FlickerConfig
+import android.tools.flicker.config.FlickerServiceConfig
+import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
+import com.android.wm.shell.flicker.service.desktopmode.flicker.DesktopModeFlickerScenarios.Companion.CLOSE_APP
+import com.android.wm.shell.flicker.service.desktopmode.flicker.DesktopModeFlickerScenarios.Companion.CLOSE_LAST_APP
+import com.android.wm.shell.flicker.service.desktopmode.scenarios.CloseAllAppsWithAppHeaderExit
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(FlickerServiceJUnit4ClassRunner::class)
+class CloseAllAppWithAppHeaderExitPortrait : CloseAllAppsWithAppHeaderExit(Rotation.ROTATION_0) {
+    @ExpectedScenarios(["CLOSE_APP", "CLOSE_LAST_APP"])
+    @Test
+    override fun closeAllAppsInDesktop() = super.closeAllAppsInDesktop()
+
+    companion object {
+        @JvmStatic
+        @FlickerConfigProvider
+        fun flickerConfigProvider(): FlickerConfig =
+            FlickerConfig()
+                .use(FlickerServiceConfig.DEFAULT)
+                .use(CLOSE_APP)
+                .use(CLOSE_LAST_APP)
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/DesktopModeFlickerScenarios.kt b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/DesktopModeFlickerScenarios.kt
new file mode 100644
index 000000000000..75dfeba3e662
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/DesktopModeFlickerScenarios.kt
@@ -0,0 +1,118 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.flicker.service.desktopmode.flicker
+
+import android.tools.flicker.AssertionInvocationGroup
+import android.tools.flicker.assertors.assertions.AppLayerIsInvisibleAtEnd
+import android.tools.flicker.assertors.assertions.AppLayerIsVisibleAlways
+import android.tools.flicker.assertors.assertions.AppLayerIsVisibleAtStart
+import android.tools.flicker.assertors.assertions.AppWindowHasDesktopModeInitialBoundsAtTheEnd
+import android.tools.flicker.assertors.assertions.AppWindowOnTopAtEnd
+import android.tools.flicker.assertors.assertions.AppWindowOnTopAtStart
+import android.tools.flicker.assertors.assertions.LauncherWindowMovesToTop
+import android.tools.flicker.config.AssertionTemplates
+import android.tools.flicker.config.FlickerConfigEntry
+import android.tools.flicker.config.ScenarioId
+import android.tools.flicker.config.desktopmode.Components
+import android.tools.flicker.extractors.ITransitionMatcher
+import android.tools.flicker.extractors.ShellTransitionScenarioExtractor
+import android.tools.traces.wm.Transition
+import android.tools.traces.wm.TransitionType
+
+class DesktopModeFlickerScenarios {
+    companion object {
+        val END_DRAG_TO_DESKTOP =
+            FlickerConfigEntry(
+                scenarioId = ScenarioId("END_DRAG_TO_DESKTOP"),
+                extractor =
+                    ShellTransitionScenarioExtractor(
+                        transitionMatcher =
+                            object : ITransitionMatcher {
+                                override fun findAll(
+                                    transitions: Collection<Transition>
+                                ): Collection<Transition> {
+                                    return transitions.filter {
+                                        it.type == TransitionType.DESKTOP_MODE_END_DRAG_TO_DESKTOP
+                                    }
+                                }
+                            }
+                    ),
+                assertions =
+                    AssertionTemplates.COMMON_ASSERTIONS +
+                        listOf(
+                                AppLayerIsVisibleAlways(Components.DESKTOP_MODE_APP),
+                                AppWindowOnTopAtEnd(Components.DESKTOP_MODE_APP),
+                                AppWindowHasDesktopModeInitialBoundsAtTheEnd(
+                                    Components.DESKTOP_MODE_APP
+                                )
+                            )
+                            .associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
+            )
+
+        val CLOSE_APP =
+            FlickerConfigEntry(
+                scenarioId = ScenarioId("CLOSE_APP"),
+                extractor =
+                    ShellTransitionScenarioExtractor(
+                        transitionMatcher =
+                            object : ITransitionMatcher {
+                                override fun findAll(
+                                    transitions: Collection<Transition>
+                                ): Collection<Transition> {
+                                    return transitions.filter { it.type == TransitionType.CLOSE }
+                                }
+                            }
+                    ),
+                assertions =
+                    AssertionTemplates.COMMON_ASSERTIONS +
+                        listOf(
+                                AppWindowOnTopAtStart(Components.DESKTOP_MODE_APP),
+                                AppLayerIsVisibleAtStart(Components.DESKTOP_MODE_APP),
+                                AppLayerIsInvisibleAtEnd(Components.DESKTOP_MODE_APP),
+                            )
+                            .associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
+            )
+
+        val CLOSE_LAST_APP =
+            FlickerConfigEntry(
+                scenarioId = ScenarioId("CLOSE_LAST_APP"),
+                extractor =
+                    ShellTransitionScenarioExtractor(
+                        transitionMatcher =
+                            object : ITransitionMatcher {
+                                override fun findAll(
+                                    transitions: Collection<Transition>
+                                ): Collection<Transition> {
+                                    val lastTransition =
+                                        transitions.findLast { it.type == TransitionType.CLOSE }
+                                    return if (lastTransition != null) listOf(lastTransition)
+                                    else emptyList()
+                                }
+                            }
+                    ),
+                assertions =
+                    AssertionTemplates.COMMON_ASSERTIONS +
+                        listOf(
+                                AppWindowOnTopAtStart(Components.DESKTOP_MODE_APP),
+                                AppLayerIsVisibleAtStart(Components.DESKTOP_MODE_APP),
+                                AppLayerIsInvisibleAtEnd(Components.DESKTOP_MODE_APP),
+                                LauncherWindowMovesToTop()
+                            )
+                            .associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
+            )
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/EnterDesktopWithDragLandscape.kt b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/EnterDesktopWithDragLandscape.kt
index 4c781d36acf6..9dfafe958b0b 100644
--- a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/EnterDesktopWithDragLandscape.kt
+++ b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/EnterDesktopWithDragLandscape.kt
@@ -17,58 +17,28 @@
 package com.android.wm.shell.flicker.service.desktopmode.flicker
 
 import android.tools.Rotation
-import android.tools.flicker.AssertionInvocationGroup
 import android.tools.flicker.FlickerConfig
 import android.tools.flicker.annotation.ExpectedScenarios
 import android.tools.flicker.annotation.FlickerConfigProvider
-import android.tools.flicker.assertors.assertions.AppLayerIsVisibleAlways
-import android.tools.flicker.assertors.assertions.AppWindowHasDesktopModeInitialBoundsAtTheEnd
-import android.tools.flicker.assertors.assertions.AppWindowOnTopAtEnd
-import android.tools.flicker.config.AssertionTemplates
 import android.tools.flicker.config.FlickerConfig
-import android.tools.flicker.config.FlickerConfigEntry
 import android.tools.flicker.config.FlickerServiceConfig
-import android.tools.flicker.config.ScenarioId
-import android.tools.flicker.config.desktopmode.Components
-import android.tools.flicker.extractors.ITransitionMatcher
-import android.tools.flicker.extractors.ShellTransitionScenarioExtractor
 import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
-import android.tools.traces.wm.Transition
-import android.tools.traces.wm.TransitionType
+import com.android.wm.shell.flicker.service.desktopmode.flicker.DesktopModeFlickerScenarios.Companion.END_DRAG_TO_DESKTOP
 import com.android.wm.shell.flicker.service.desktopmode.scenarios.EnterDesktopWithDrag
 import org.junit.Test
 import org.junit.runner.RunWith
 
 @RunWith(FlickerServiceJUnit4ClassRunner::class)
 class EnterDesktopWithDragLandscape : EnterDesktopWithDrag(Rotation.ROTATION_90) {
-    @ExpectedScenarios(["END_DRAG_TO_DESKTOP"]) @Test override fun enterDesktopWithDrag() =
-        super.enterDesktopWithDrag()
+    @ExpectedScenarios(["END_DRAG_TO_DESKTOP"])
+    @Test
+    override fun enterDesktopWithDrag() = super.enterDesktopWithDrag()
 
     companion object {
-        private val END_DRAG_TO_DESKTOP = FlickerConfigEntry(
-            scenarioId = ScenarioId("END_DRAG_TO_DESKTOP"),
-            extractor = ShellTransitionScenarioExtractor(
-                transitionMatcher = object : ITransitionMatcher {
-                    override fun findAll(
-                        transitions: Collection<Transition>
-                    ): Collection<Transition> {
-                        return transitions.filter {
-                            it.type == TransitionType.DESKTOP_MODE_END_DRAG_TO_DESKTOP}
-                    }
-                }),
-            assertions = AssertionTemplates.COMMON_ASSERTIONS +
-                    listOf(
-                        AppLayerIsVisibleAlways(Components.DESKTOP_MODE_APP),
-                        AppWindowOnTopAtEnd(Components.DESKTOP_MODE_APP),
-                        AppWindowHasDesktopModeInitialBoundsAtTheEnd(Components.DESKTOP_MODE_APP)
-                    ).associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
-        )
 
         @JvmStatic
         @FlickerConfigProvider
         fun flickerConfigProvider(): FlickerConfig =
-            FlickerConfig()
-                    .use(FlickerServiceConfig.DEFAULT)
-                    .use(END_DRAG_TO_DESKTOP)
+            FlickerConfig().use(FlickerServiceConfig.DEFAULT).use(END_DRAG_TO_DESKTOP)
     }
 }
diff --git a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/EnterDesktopWithDragPortrait.kt b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/EnterDesktopWithDragPortrait.kt
index d99d875fb126..1c7d6237eb8a 100644
--- a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/EnterDesktopWithDragPortrait.kt
+++ b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/flicker/EnterDesktopWithDragPortrait.kt
@@ -17,58 +17,27 @@
 package com.android.wm.shell.flicker.service.desktopmode.flicker
 
 import android.tools.Rotation
-import android.tools.flicker.AssertionInvocationGroup
 import android.tools.flicker.FlickerConfig
 import android.tools.flicker.annotation.ExpectedScenarios
 import android.tools.flicker.annotation.FlickerConfigProvider
-import android.tools.flicker.assertors.assertions.AppLayerIsVisibleAlways
-import android.tools.flicker.assertors.assertions.AppWindowHasDesktopModeInitialBoundsAtTheEnd
-import android.tools.flicker.assertors.assertions.AppWindowOnTopAtEnd
-import android.tools.flicker.config.AssertionTemplates
 import android.tools.flicker.config.FlickerConfig
-import android.tools.flicker.config.FlickerConfigEntry
 import android.tools.flicker.config.FlickerServiceConfig
-import android.tools.flicker.config.ScenarioId
-import android.tools.flicker.config.desktopmode.Components
-import android.tools.flicker.extractors.ITransitionMatcher
-import android.tools.flicker.extractors.ShellTransitionScenarioExtractor
 import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
-import android.tools.traces.wm.Transition
-import android.tools.traces.wm.TransitionType
+import com.android.wm.shell.flicker.service.desktopmode.flicker.DesktopModeFlickerScenarios.Companion.END_DRAG_TO_DESKTOP
 import com.android.wm.shell.flicker.service.desktopmode.scenarios.EnterDesktopWithDrag
 import org.junit.Test
 import org.junit.runner.RunWith
 
 @RunWith(FlickerServiceJUnit4ClassRunner::class)
 class EnterDesktopWithDragPortrait : EnterDesktopWithDrag(Rotation.ROTATION_0) {
-    @ExpectedScenarios(["END_DRAG_TO_DESKTOP"]) @Test override fun enterDesktopWithDrag() =
-        super.enterDesktopWithDrag()
+    @ExpectedScenarios(["END_DRAG_TO_DESKTOP"])
+    @Test
+    override fun enterDesktopWithDrag() = super.enterDesktopWithDrag()
 
     companion object {
-        private val END_DRAG_TO_DESKTOP = FlickerConfigEntry(
-            scenarioId = ScenarioId("END_DRAG_TO_DESKTOP"),
-            extractor = ShellTransitionScenarioExtractor(
-                transitionMatcher = object : ITransitionMatcher {
-                    override fun findAll(
-                        transitions: Collection<Transition>
-                    ): Collection<Transition> {
-                        return transitions.filter {
-                            it.type == TransitionType.DESKTOP_MODE_END_DRAG_TO_DESKTOP}
-                    }
-                }),
-            assertions = AssertionTemplates.COMMON_ASSERTIONS +
-                    listOf(
-                        AppLayerIsVisibleAlways(Components.DESKTOP_MODE_APP),
-                        AppWindowOnTopAtEnd(Components.DESKTOP_MODE_APP),
-                        AppWindowHasDesktopModeInitialBoundsAtTheEnd(Components.DESKTOP_MODE_APP)
-                    ).associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
-        )
-
         @JvmStatic
         @FlickerConfigProvider
         fun flickerConfigProvider(): FlickerConfig =
-            FlickerConfig()
-                    .use(FlickerServiceConfig.DEFAULT)
-                    .use(END_DRAG_TO_DESKTOP)
+            FlickerConfig().use(FlickerServiceConfig.DEFAULT).use(END_DRAG_TO_DESKTOP)
     }
 }
diff --git a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/scenarios/CloseAllAppsWithAppHeaderExit.kt b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/scenarios/CloseAllAppsWithAppHeaderExit.kt
new file mode 100644
index 000000000000..0c2b5015840d
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/scenarios/CloseAllAppsWithAppHeaderExit.kt
@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.flicker.service.desktopmode.scenarios
+
+import android.app.Instrumentation
+import android.tools.NavBar
+import android.tools.Rotation
+import android.tools.traces.parsers.WindowManagerStateHelper
+import androidx.test.platform.app.InstrumentationRegistry
+import androidx.test.uiautomator.UiDevice
+import com.android.launcher3.tapl.LauncherInstrumentation
+import com.android.server.wm.flicker.helpers.DesktopModeAppHelper
+import com.android.server.wm.flicker.helpers.MailAppHelper
+import com.android.server.wm.flicker.helpers.NonResizeableAppHelper
+import com.android.server.wm.flicker.helpers.SimpleAppHelper
+import com.android.window.flags.Flags
+import com.android.wm.shell.flicker.service.common.Utils
+import org.junit.After
+import org.junit.Assume
+import org.junit.Before
+import org.junit.Ignore
+import org.junit.Rule
+import org.junit.Test
+
+@Ignore("Base Test Class")
+abstract class CloseAllAppsWithAppHeaderExit
+@JvmOverloads
+constructor(val rotation: Rotation = Rotation.ROTATION_0) {
+
+    private val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
+    private val tapl = LauncherInstrumentation()
+    private val wmHelper = WindowManagerStateHelper(instrumentation)
+    private val device = UiDevice.getInstance(instrumentation)
+    private val testApp = DesktopModeAppHelper(SimpleAppHelper(instrumentation))
+    private val mailApp = DesktopModeAppHelper(MailAppHelper(instrumentation))
+    private val nonResizeableApp = DesktopModeAppHelper(NonResizeableAppHelper(instrumentation))
+
+
+
+    @Rule @JvmField val testSetupRule = Utils.testSetupRule(NavBar.MODE_GESTURAL, rotation)
+
+    @Before
+    fun setup() {
+        Assume.assumeTrue(Flags.enableDesktopWindowingMode())
+        tapl.setEnableRotation(true)
+        tapl.setExpectedRotation(rotation.value)
+        testApp.enterDesktopWithDrag(wmHelper, device)
+        mailApp.launchViaIntent(wmHelper)
+        nonResizeableApp.launchViaIntent(wmHelper)
+    }
+
+    @Test
+    open fun closeAllAppsInDesktop() {
+        nonResizeableApp.closeDesktopApp(wmHelper, device)
+        mailApp.closeDesktopApp(wmHelper, device)
+        testApp.closeDesktopApp(wmHelper, device)
+    }
+
+    @After
+    fun teardown() {
+        testApp.exit(wmHelper)
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/scenarios/EnterDesktopWithDrag.kt b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/scenarios/EnterDesktopWithDrag.kt
index 0403b4f64faf..9e9998ef7c2a 100644
--- a/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/scenarios/EnterDesktopWithDrag.kt
+++ b/libs/WindowManager/Shell/tests/flicker/service/src/com/android/wm/shell/flicker/service/desktopmode/scenarios/EnterDesktopWithDrag.kt
@@ -23,15 +23,18 @@ import android.tools.traces.parsers.WindowManagerStateHelper
 import androidx.test.platform.app.InstrumentationRegistry
 import androidx.test.uiautomator.UiDevice
 import com.android.launcher3.tapl.LauncherInstrumentation
+import com.android.server.wm.flicker.helpers.DesktopModeAppHelper
 import com.android.server.wm.flicker.helpers.SimpleAppHelper
+import com.android.window.flags.Flags
 import com.android.wm.shell.flicker.service.common.Utils
-import com.android.wm.shell.flicker.utils.DesktopModeUtils
 import org.junit.After
+import org.junit.Assume
 import org.junit.Before
 import org.junit.Ignore
 import org.junit.Rule
 import org.junit.Test
 
+
 @Ignore("Base Test Class")
 abstract class EnterDesktopWithDrag
 @JvmOverloads
@@ -41,19 +44,20 @@ constructor(val rotation: Rotation = Rotation.ROTATION_0) {
     private val tapl = LauncherInstrumentation()
     private val wmHelper = WindowManagerStateHelper(instrumentation)
     private val device = UiDevice.getInstance(instrumentation)
-    private val testApp = SimpleAppHelper(instrumentation)
+    private val testApp = DesktopModeAppHelper(SimpleAppHelper(instrumentation))
 
     @Rule @JvmField val testSetupRule = Utils.testSetupRule(NavBar.MODE_GESTURAL, rotation)
 
     @Before
     fun setup() {
+        Assume.assumeTrue(Flags.enableDesktopWindowingMode())
         tapl.setEnableRotation(true)
         tapl.setExpectedRotation(rotation.value)
     }
 
     @Test
     open fun enterDesktopWithDrag() {
-        DesktopModeUtils.enterDesktopWithDrag(wmHelper, device, testApp)
+        testApp.enterDesktopWithDrag(wmHelper, device)
     }
 
     @After
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/DesktopModeUtils.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/DesktopModeUtils.kt
deleted file mode 100644
index 345bc5ebb20e..000000000000
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/utils/DesktopModeUtils.kt
+++ /dev/null
@@ -1,112 +0,0 @@
-/*
- * Copyright (C) 2024 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.wm.shell.flicker.utils
-
-import android.tools.device.apphelpers.StandardAppHelper
-import android.tools.helpers.SYSTEMUI_PACKAGE
-import android.tools.traces.component.IComponentMatcher
-import android.tools.traces.parsers.WindowManagerStateHelper
-import android.tools.traces.wm.WindowingMode
-import androidx.test.uiautomator.By
-import androidx.test.uiautomator.BySelector
-import androidx.test.uiautomator.UiDevice
-import androidx.test.uiautomator.Until
-
-/**
- * Provides a collection of utility functions for desktop mode testing.
- */
-object DesktopModeUtils {
-    private const val TIMEOUT_MS = 3_000L
-    private const val CAPTION = "desktop_mode_caption"
-    private const val CAPTION_HANDLE = "caption_handle"
-    private const val MAXIMIZE_BUTTON = "maximize_button_view"
-
-    private val captionFullscreen: BySelector
-        get() = By.res(SYSTEMUI_PACKAGE, CAPTION)
-    private val captionHandle: BySelector
-        get() = By.res(SYSTEMUI_PACKAGE, CAPTION_HANDLE)
-    private val maximizeButton: BySelector
-        get() = By.res(SYSTEMUI_PACKAGE, MAXIMIZE_BUTTON)
-
-    /**
-     * Wait for an app moved to desktop to finish its transition.
-     */
-    private fun waitForAppToMoveToDesktop(
-        wmHelper: WindowManagerStateHelper,
-        currentApp: IComponentMatcher,
-    ) {
-        wmHelper
-            .StateSyncBuilder()
-            .withWindowSurfaceAppeared(currentApp)
-            .withFreeformApp(currentApp)
-            .withAppTransitionIdle()
-            .waitForAndVerify()
-    }
-
-    /**
-     * Click maximise button on the app header for the given app.
-     */
-    fun maximiseDesktopApp(
-        wmHelper: WindowManagerStateHelper,
-        device: UiDevice,
-        currentApp: StandardAppHelper
-    ) {
-        if (wmHelper.getWindow(currentApp)?.windowingMode
-            != WindowingMode.WINDOWING_MODE_FREEFORM.value)
-            error("expected a freeform window to maximise but window is not in freefrom mode")
-
-        val maximizeButton =
-            device.wait(Until.findObject(maximizeButton), TIMEOUT_MS)
-                ?: error("Unable to find view $maximizeButton\n")
-        maximizeButton.click()
-    }
-
-    /**
-     * Move an app to Desktop by dragging the app handle at the top.
-     */
-    fun enterDesktopWithDrag(
-        wmHelper: WindowManagerStateHelper,
-        device: UiDevice,
-        currentApp: StandardAppHelper,
-    ) {
-        currentApp.launchViaIntent(wmHelper)
-        dragToDesktop(wmHelper, currentApp, device)
-        waitForAppToMoveToDesktop(wmHelper, currentApp)
-    }
-
-    private fun dragToDesktop(
-        wmHelper: WindowManagerStateHelper,
-        currentApp: StandardAppHelper,
-        device: UiDevice
-    ) {
-        val windowRect = wmHelper.getWindowRegion(currentApp).bounds
-        val startX = windowRect.centerX()
-
-        // Start dragging a little under the top to prevent dragging the notification shade.
-        val startY = 10
-
-        val displayRect =
-            wmHelper.currentState.wmState.getDefaultDisplay()?.displayRect
-                ?: throw IllegalStateException("Default display is null")
-
-        // The position we want to drag to
-        val endY = displayRect.centerY() / 2
-
-        // drag the window to move to desktop
-        device.drag(startX, startY, startX, endY, 100)
-    }
-}
diff --git a/tests/FlickerTests/test-apps/app-helpers/OWNERS b/tests/FlickerTests/test-apps/app-helpers/OWNERS
new file mode 100644
index 000000000000..ab6253200f73
--- /dev/null
+++ b/tests/FlickerTests/test-apps/app-helpers/OWNERS
@@ -0,0 +1,2 @@
+uysalorhan@google.com
+pragyabajoria@google.com
\ No newline at end of file
diff --git a/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/DesktopModeAppHelper.kt b/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/DesktopModeAppHelper.kt
new file mode 100644
index 000000000000..461dfeca09e3
--- /dev/null
+++ b/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/DesktopModeAppHelper.kt
@@ -0,0 +1,124 @@
+/*
+ * Copyright (C) 2024 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.helpers
+
+import android.tools.device.apphelpers.IStandardAppHelper
+import android.tools.helpers.SYSTEMUI_PACKAGE
+import android.tools.traces.parsers.WindowManagerStateHelper
+import android.tools.traces.wm.WindowingMode
+import androidx.test.uiautomator.By
+import androidx.test.uiautomator.BySelector
+import androidx.test.uiautomator.UiDevice
+import androidx.test.uiautomator.UiObject2
+import androidx.test.uiautomator.Until
+
+/**
+ * Wrapper class around App helper classes. This class adds functionality to the apps that the
+ * desktop apps would have.
+ */
+open class DesktopModeAppHelper(private val innerHelper: IStandardAppHelper) :
+    IStandardAppHelper by innerHelper {
+    private val TIMEOUT_MS = 3_000L
+    private val CAPTION = "desktop_mode_caption"
+    private val CAPTION_HANDLE = "caption_handle"
+    private val MAXIMIZE_BUTTON = "maximize_window"
+    private val MAXIMIZE_BUTTON_VIEW = "maximize_button_view"
+    private val CLOSE_BUTTON = "close_window"
+
+    private val caption: BySelector
+        get() = By.res(SYSTEMUI_PACKAGE, CAPTION)
+
+    /** Wait for an app moved to desktop to finish its transition. */
+    private fun waitForAppToMoveToDesktop(wmHelper: WindowManagerStateHelper) {
+        wmHelper
+            .StateSyncBuilder()
+            .withWindowSurfaceAppeared(innerHelper)
+            .withFreeformApp(innerHelper)
+            .withAppTransitionIdle()
+            .waitForAndVerify()
+    }
+
+    /** Move an app to Desktop by dragging the app handle at the top. */
+    fun enterDesktopWithDrag(
+        wmHelper: WindowManagerStateHelper,
+        device: UiDevice,
+    ) {
+        innerHelper.launchViaIntent(wmHelper)
+        dragToDesktop(wmHelper, device)
+        waitForAppToMoveToDesktop(wmHelper)
+    }
+
+    private fun dragToDesktop(wmHelper: WindowManagerStateHelper, device: UiDevice) {
+        val windowRect = wmHelper.getWindowRegion(innerHelper).bounds
+        val startX = windowRect.centerX()
+
+        // Start dragging a little under the top to prevent dragging the notification shade.
+        val startY = 10
+
+        val displayRect =
+            wmHelper.currentState.wmState.getDefaultDisplay()?.displayRect
+                ?: throw IllegalStateException("Default display is null")
+
+        // The position we want to drag to
+        val endY = displayRect.centerY() / 2
+
+        // drag the window to move to desktop
+        device.drag(startX, startY, startX, endY, 100)
+    }
+
+    /** Click maximise button on the app header for the given app. */
+    fun maximiseDesktopApp(wmHelper: WindowManagerStateHelper, device: UiDevice) {
+        val caption = getCaptionForTheApp(wmHelper, device)
+        val maximizeButton =
+            caption
+                ?.children
+                ?.find { it.resourceName.endsWith(MAXIMIZE_BUTTON_VIEW) }
+                ?.children
+                ?.get(0)
+        maximizeButton?.click()
+        wmHelper.StateSyncBuilder().withAppTransitionIdle().waitForAndVerify()
+    }
+    /** Click close button on the app header for the given app. */
+    fun closeDesktopApp(wmHelper: WindowManagerStateHelper, device: UiDevice) {
+        val caption = getCaptionForTheApp(wmHelper, device)
+        val closeButton = caption?.children?.find { it.resourceName.endsWith(CLOSE_BUTTON) }
+        closeButton?.click()
+        wmHelper
+            .StateSyncBuilder()
+            .withAppTransitionIdle()
+            .withWindowSurfaceDisappeared(innerHelper)
+            .waitForAndVerify()
+    }
+
+    private fun getCaptionForTheApp(
+        wmHelper: WindowManagerStateHelper,
+        device: UiDevice
+    ): UiObject2? {
+        if (
+            wmHelper.getWindow(innerHelper)?.windowingMode !=
+                WindowingMode.WINDOWING_MODE_FREEFORM.value
+        )
+            error("expected a freeform window with caption but window is not in freeform mode")
+        val captions =
+            device.wait(Until.findObjects(caption), TIMEOUT_MS)
+                ?: error("Unable to find view $caption\n")
+
+        return captions.find {
+            wmHelper.getWindowRegion(innerHelper).bounds.contains(it.visibleBounds)
+        }
+    }
+}