From 604bd2b7e60c3e739a746be54520a95a7fb4f5b8 Mon Sep 17 00:00:00 2001
From: Pragya Bajoria <pragyabajoria@google.com>
Date: Tue, 30 Jul 2024 15:16:02 +0000
Subject: Handle state management of repository state to remove freeform
 windows without requiring the caller to know the right order of methods to
 call.

Bug: 332682201
Change-Id: I04923eae519cb787270b56b4692cae06f24aaa80
Flag: EXEMPT (no-op for functionality)
Test: atest DesktopModeTaskRepositoryTest
---
 .../shell/desktopmode/DesktopModeTaskRepository.kt |  6 +++-
 .../wm/shell/freeform/FreeformTaskListener.java    |  4 ---
 .../desktopmode/DesktopModeTaskRepositoryTest.kt   | 37 ++++++++++++++++++++++
 3 files changed, 42 insertions(+), 5 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 4060900c63ce..83752945e9f9 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
@@ -300,13 +300,17 @@ class DesktopModeTaskRepository {
         }
     }
 
-    /** Removes given task from a valid [displayId]. */
+    /** Removes given task from a valid [displayId] and updates the repository state. */
     private fun removeTaskFromDisplay(displayId: Int, taskId: Int) {
         logD("Removes freeform task: taskId=%d, displayId=%d", taskId, displayId)
         desktopTaskDataByDisplayId[displayId]?.freeformTasksInZOrder?.remove(taskId)
         boundsBeforeMaximizeByTaskId.remove(taskId)
         logD("Remaining freeform tasks: %s",
             desktopTaskDataByDisplayId[displayId]?.freeformTasksInZOrder?.toDumpString())
+        // Remove task from unminimized task if it is minimized.
+        unminimizeTask(displayId, taskId)
+        removeActiveTask(taskId)
+        updateTaskVisibility(displayId, taskId, visible = false);
     }
 
     /**
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
index 640f87289f27..456767a1c9af 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
@@ -117,10 +117,6 @@ public class FreeformTaskListener implements ShellTaskOrganizer.TaskListener,
         if (DesktopModeStatus.canEnterDesktopMode(mContext)) {
             mDesktopModeTaskRepository.ifPresent(repository -> {
                 repository.removeFreeformTask(taskInfo.displayId, taskInfo.taskId);
-                repository.unminimizeTask(taskInfo.displayId, taskInfo.taskId);
-                repository.removeActiveTask(taskInfo.taskId, /* excludedDisplayId= */ null);
-                repository.updateTaskVisibility(
-                    taskInfo.displayId, taskInfo.taskId, /* visible= */ false);
             });
         }
         mWindowDecorationViewModel.onTaskVanished(taskInfo);
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 7acee783f0ee..d3404f7bd261 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
@@ -534,6 +534,43 @@ class DesktopModeTaskRepositoryTest : ShellTestCase() {
         assertThat(repo.removeBoundsBeforeMaximize(taskId)).isNull()
     }
 
+    @Test
+    fun removeFreeformTask_removesActiveTask() {
+        val taskId = 1
+        val listener = TestListener()
+        repo.addActiveTaskListener(listener)
+        repo.addActiveTask(DEFAULT_DISPLAY, taskId)
+        repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, taskId)
+
+        repo.removeFreeformTask(THIRD_DISPLAY, taskId)
+
+        assertThat(repo.isActiveTask(taskId)).isFalse()
+        assertThat(listener.activeChangesOnDefaultDisplay).isEqualTo(2)
+    }
+
+    @Test
+    fun removeFreeformTask_unminimizesTask() {
+        val taskId = 1
+        repo.addActiveTask(DEFAULT_DISPLAY, taskId)
+        repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, taskId)
+        repo.minimizeTask(DEFAULT_DISPLAY, taskId)
+
+        repo.removeFreeformTask(DEFAULT_DISPLAY, taskId)
+
+        assertThat(repo.isMinimizedTask(taskId)).isFalse()
+    }
+
+    @Test
+    fun removeFreeformTask_updatesTaskVisibility() {
+        val taskId = 1
+        repo.addActiveTask(DEFAULT_DISPLAY, taskId)
+        repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, taskId)
+
+        repo.removeFreeformTask(THIRD_DISPLAY, taskId)
+
+        assertThat(repo.isVisibleTask(taskId)).isFalse()
+    }
+
     @Test
     fun saveBoundsBeforeMaximize_boundsSavedByTaskId() {
         val taskId = 1
-- 
cgit v1.2.3-59-g8ed1b