Use task's displayId in moveToBackground flows

Moving a background task to desktop involves "activating" the desk it'll
end up in, meaning, bringing up the home/wallpaper in front of other
non-desktop tasks, bringing up inactive tasks that already existed in
that desk and applying task-limit checks.

This function previously used DEFAULT_DISPLAY, so this CL changes that
to use the displayId available in the TaskInfo of the task.

Flag: com.android.window.flags.enable_desktop_windowing_mode
Flag: com.android.window.flags.enable_move_to_next_display_shortcut
Flag: com.android.window.flags.enable_multiple_desktops_backend
Bug: 342378842
Test: m WMShellUnitTests
Change-Id: Ia7db80d6608f2ccb9807d8695a19d14fb54ea4c0
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index bedea2f..73cba57 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -468,15 +468,14 @@
         remoteTransition: RemoteTransition? = null,
         callback: IMoveToDesktopCallback? = null,
     ): Boolean {
-        val runningTask = shellTaskOrganizer.getRunningTaskInfo(taskId)
-        val backgroundTask = recentTasksController?.findTaskInBackground(taskId)
-        if (runningTask == null && backgroundTask == null) {
+        val task =
+            shellTaskOrganizer.getRunningTaskInfo(taskId)
+                ?: recentTasksController?.findTaskInBackground(taskId)
+        if (task == null) {
             logW("moveTaskToDefaultDeskAndActivate taskId=%d not found", taskId)
             return false
         }
-        // TODO(342378842): Instead of using default display, support multiple displays
-        val displayId = runningTask?.displayId ?: DEFAULT_DISPLAY
-        val deskId = getDefaultDeskId(displayId)
+        val deskId = getDefaultDeskId(task.displayId)
         return moveTaskToDesk(
             taskId = taskId,
             deskId = deskId,
@@ -528,14 +527,14 @@
         remoteTransition: RemoteTransition? = null,
         callback: IMoveToDesktopCallback? = null,
     ): Boolean {
-        if (recentTasksController?.findTaskInBackground(taskId) == null) {
+        val task = recentTasksController?.findTaskInBackground(taskId)
+        if (task == null) {
             logW("moveBackgroundTaskToDesktop taskId=%d not found", taskId)
             return false
         }
         logV("moveBackgroundTaskToDesktop with taskId=%d", taskId)
-        // TODO(342378842): Instead of using default display, support multiple displays
         val taskIdToMinimize =
-            bringDesktopAppsToFrontBeforeShowingNewTask(DEFAULT_DISPLAY, wct, taskId)
+            bringDesktopAppsToFrontBeforeShowingNewTask(task.displayId, wct, taskId)
         val exitResult =
             desktopImmersiveController.exitImmersiveIfApplicable(
                 wct = wct,
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index 8ce31a5..25c7b44 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -1611,7 +1611,7 @@
     @Test
     @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
     fun moveTaskToDesktop_desktopWallpaperDisabled_nonRunningTask_launchesInFreeform() {
-        val task = createTaskInfo(1)
+        val task = createRecentTaskInfo(1)
         whenever(shellTaskOrganizer.getRunningTaskInfo(anyInt())).thenReturn(null)
         whenever(recentTasksController.findTaskInBackground(anyInt())).thenReturn(task)
 
@@ -1626,7 +1626,7 @@
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
     fun moveTaskToDesktop_desktopWallpaperEnabled_nonRunningTask_launchesInFreeform() {
         whenever(desktopWallpaperActivityTokenProvider.getToken()).thenReturn(null)
-        val task = createTaskInfo(1)
+        val task = createRecentTaskInfo(1)
         whenever(shellTaskOrganizer.getRunningTaskInfo(anyInt())).thenReturn(null)
         whenever(recentTasksController.findTaskInBackground(anyInt())).thenReturn(task)
 
@@ -1798,7 +1798,7 @@
         whenever(transitions.startTransition(anyInt(), any(), transitionHandlerArgCaptor.capture()))
             .thenReturn(Binder())
 
-        val task = createTaskInfo(1)
+        val task = createRecentTaskInfo(1)
         whenever(shellTaskOrganizer.getRunningTaskInfo(anyInt())).thenReturn(null)
         whenever(recentTasksController.findTaskInBackground(anyInt())).thenReturn(task)
         controller.moveTaskToDefaultDeskAndActivate(
@@ -1813,6 +1813,34 @@
     }
 
     @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY,
+        Flags.FLAG_ENABLE_PER_DISPLAY_DESKTOP_WALLPAPER_ACTIVITY,
+        Flags.FLAG_ENABLE_DESKTOP_WALLPAPER_ACTIVITY_FOR_SYSTEM_USER,
+    )
+    fun moveBackgroundTaskToDesktop_nonDefaultDisplay_reordersHomeAndWallpaperOfNonDefaultDisplay() {
+        val homeTask = setUpHomeTask(displayId = SECOND_DISPLAY)
+        val wallpaperToken = MockToken().token()
+        whenever(desktopWallpaperActivityTokenProvider.getToken(SECOND_DISPLAY))
+            .thenReturn(wallpaperToken)
+        val task = setUpFreeformTask(displayId = SECOND_DISPLAY, deskId = 2, background = true)
+
+        controller.moveTaskToDefaultDeskAndActivate(
+            taskId = task.taskId,
+            transitionSource = UNKNOWN,
+            remoteTransition = RemoteTransition(spy(TestRemoteTransition())),
+        )
+
+        val wct = getLatestTransition()
+        val homeReorderIndex = wct.indexOfReorder(homeTask, toTop = true)
+        val wallpaperReorderIndex = wct.indexOfReorder(wallpaperToken, toTop = true)
+        assertThat(homeReorderIndex).isNotEqualTo(-1)
+        assertThat(wallpaperReorderIndex).isNotEqualTo(-1)
+        // Wallpaper last, to be in front of Home.
+        assertThat(wallpaperReorderIndex).isGreaterThan(homeReorderIndex)
+    }
+
+    @Test
     fun moveRunningTaskToDesktop_remoteTransition_usesOneShotHandler() {
         val transitionHandlerArgCaptor = argumentCaptor<TransitionHandler>()
         whenever(transitions.startTransition(anyInt(), any(), transitionHandlerArgCaptor.capture()))
@@ -2463,7 +2491,7 @@
 
     @Test
     fun moveTaskToFront_backgroundTask_launchesTask() {
-        val task = createTaskInfo(1)
+        val task = createRecentTaskInfo(1)
         whenever(shellTaskOrganizer.getRunningTaskInfo(anyInt())).thenReturn(null)
         whenever(
                 desktopMixedTransitionHandler.startLaunchTransition(
@@ -2485,7 +2513,7 @@
     @Test
     fun moveTaskToFront_backgroundTaskBringsTasksOverLimit_minimizesBackTask() {
         val freeformTasks = (1..MAX_TASK_LIMIT).map { _ -> setUpFreeformTask() }
-        val task = createTaskInfo(1001)
+        val task = createRecentTaskInfo(1001)
         whenever(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(null)
         whenever(
                 desktopMixedTransitionHandler.startLaunchTransition(
@@ -6711,7 +6739,7 @@
         if (background) {
             whenever(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(null)
             whenever(recentTasksController.findTaskInBackground(task.taskId))
-                .thenReturn(createTaskInfo(task.taskId))
+                .thenReturn(createRecentTaskInfo(taskId = task.taskId, displayId = displayId))
         } else {
             whenever(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(task)
         }
@@ -7146,8 +7174,9 @@
     } ?: false
 }
 
-private fun createTaskInfo(id: Int) =
+private fun createRecentTaskInfo(taskId: Int, displayId: Int = DEFAULT_DISPLAY) =
     RecentTaskInfo().apply {
-        taskId = id
+        this.taskId = taskId
+        this.displayId = displayId
         token = WindowContainerToken(mock(IWindowContainerToken::class.java))
     }