Merge "Find and remove visible task onTaskVanished." into main

2 files changed, 32 insertions, 4 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
index 7c8fcbb16711..99a00b8d33c4 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
@@ -20,6 +20,7 @@ import android.graphics.Region
 import android.util.ArrayMap
 import android.util.ArraySet
 import android.util.SparseArray
+import android.view.Display.INVALID_DISPLAY
 import androidx.core.util.forEach
 import androidx.core.util.keyIterator
 import androidx.core.util.valueIterator
@@ -226,6 +227,14 @@ class DesktopModeTaskRepository {
                         displayData[otherDisplayId].visibleTasks.size)
                 }
             }
+        } else if (displayId == INVALID_DISPLAY) {
+            // Task has vanished. Check which display to remove the task from.
+            displayData.forEach { displayId, data ->
+                if (data.visibleTasks.remove(taskId)) {
+                    notifyVisibleTaskListeners(displayId, data.visibleTasks.size)
+                }
+            }
+            return
         }
 
         val prevCount = getVisibleTaskCount(displayId)
@@ -236,6 +245,7 @@ class DesktopModeTaskRepository {
         }
         val newCount = getVisibleTaskCount(displayId)
 
+        // Check if count changed
         if (prevCount != newCount) {
             KtProtoLog.d(
                 WM_SHELL_DESKTOP_MODE,
@@ -244,10 +254,6 @@ class DesktopModeTaskRepository {
                 visible,
                 displayId
             )
-        }
-
-        // Check if count changed
-        if (prevCount != newCount) {
             KtProtoLog.d(
                 WM_SHELL_DESKTOP_MODE,
                 "DesktopTaskRepo: visibleTaskCount has changed from %d to %d",
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
index 445f74a52b0d..9f3a4d9c135d 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
@@ -18,6 +18,7 @@ package com.android.wm.shell.desktopmode
 
 import android.testing.AndroidTestingRunner
 import android.view.Display.DEFAULT_DISPLAY
+import android.view.Display.INVALID_DISPLAY
 import androidx.test.filters.SmallTest
 import com.android.wm.shell.ShellTestCase
 import com.android.wm.shell.TestShellExecutor
@@ -237,6 +238,27 @@ class DesktopModeTaskRepositoryTest : ShellTestCase() {
         assertThat(listener.visibleChangesOnDefaultDisplay).isEqualTo(4)
     }
 
+    /**
+     * When a task vanishes, the displayId of the task is set to INVALID_DISPLAY.
+     * This tests that task is removed from the last parent display when it vanishes.
+     */
+    @Test
+    fun updateVisibleFreeformTasks_removeVisibleTasksRemovesTaskWithInvalidDisplay() {
+        val listener = TestVisibilityListener()
+        val executor = TestShellExecutor()
+        repo.addVisibleTasksListener(listener, executor)
+        repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = true)
+        repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 2, visible = true)
+        executor.flushAll()
+
+        assertThat(listener.visibleTasksCountOnDefaultDisplay).isEqualTo(2)
+        repo.updateVisibleFreeformTasks(INVALID_DISPLAY, taskId = 1, visible = false)
+        executor.flushAll()
+
+        assertThat(listener.visibleChangesOnDefaultDisplay).isEqualTo(3)
+        assertThat(listener.visibleTasksCountOnDefaultDisplay).isEqualTo(1)
+    }
+
     @Test
     fun getVisibleTaskCount() {
         // No tasks, count is 0