Merge "Register DesktopUserRepos with ShellController" into main

9 files changed, 37 insertions, 9 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index 0cd0f4a97bbf..9d2da41725b1 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -1025,12 +1025,14 @@ public abstract class WMShellModule {
     static DesktopUserRepositories provideDesktopUserRepositories(
             Context context,
             ShellInit shellInit,
+            ShellController shellController,
             DesktopPersistentRepository desktopPersistentRepository,
             DesktopRepositoryInitializer desktopRepositoryInitializer,
             @ShellMainThread CoroutineScope mainScope,
             UserManager userManager
     ) {
-        return new DesktopUserRepositories(context, shellInit, desktopPersistentRepository,
+        return new DesktopUserRepositories(context, shellInit, shellController,
+                desktopPersistentRepository,
                 desktopRepositoryInitializer,
                 mainScope, userManager);
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopUserRepositories.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopUserRepositories.kt
index e5f52839d9f4..8b5d1c502bc9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopUserRepositories.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopUserRepositories.kt
@@ -28,6 +28,7 @@ import com.android.wm.shell.desktopmode.persistence.DesktopRepositoryInitializer
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
 import com.android.wm.shell.shared.annotations.ShellMainThread
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
+import com.android.wm.shell.sysui.ShellController
 import com.android.wm.shell.sysui.ShellInit
 import com.android.wm.shell.sysui.UserChangeListener
 import kotlinx.coroutines.CoroutineScope
@@ -36,6 +37,7 @@ import kotlinx.coroutines.CoroutineScope
 class DesktopUserRepositories(
     context: Context,
     shellInit: ShellInit,
+    private val shellController: ShellController,
     private val persistentRepository: DesktopPersistentRepository,
     private val repositoryInitializer: DesktopRepositoryInitializer,
     @ShellMainThread private val mainCoroutineScope: CoroutineScope,
@@ -61,15 +63,16 @@ class DesktopUserRepositories(
     init {
         userId = ActivityManager.getCurrentUser()
         if (DesktopModeStatus.canEnterDesktopMode(context)) {
-            shellInit.addInitCallback(::initRepoFromPersistentStorage, this)
+            shellInit.addInitCallback(::onInit, this)
         }
         if (Flags.enableDesktopWindowingHsum()) {
             userIdToProfileIdsMap[userId] = userManager.getProfiles(userId).map { it.id }
         }
     }
 
-    private fun initRepoFromPersistentStorage() {
+    private fun onInit() {
         repositoryInitializer.initialize(this)
+        shellController.addUserChangeListener(this)
     }
 
     /** Returns [DesktopRepository] for the parent user id. */
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandlerTest.kt
index 41a594a3347a..4cc641cd1d81 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandlerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandlerTest.kt
@@ -36,6 +36,7 @@ import com.android.dx.mockito.inline.extended.ExtendedMockito.never
 import com.android.dx.mockito.inline.extended.StaticMockitoSession
 import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE
 import com.android.window.flags.Flags.FLAG_RESPECT_ORIENTATION_CHANGE_FOR_UNRESIZEABLE
+import com.android.wm.shell.sysui.ShellController
 import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.ShellTestCase
 import com.android.wm.shell.common.ShellExecutor
@@ -98,6 +99,7 @@ class DesktopActivityOrientationChangeHandlerTest : ShellTestCase() {
     @Mock lateinit var persistentRepository: DesktopPersistentRepository
     @Mock lateinit var repositoryInitializer: DesktopRepositoryInitializer
     @Mock lateinit var userManager: UserManager
+    @Mock lateinit var shellController: ShellController
 
     private lateinit var mockitoSession: StaticMockitoSession
     private lateinit var handler: DesktopActivityOrientationChangeHandler
@@ -123,6 +125,7 @@ class DesktopActivityOrientationChangeHandlerTest : ShellTestCase() {
             DesktopUserRepositories(
                 context,
                 shellInit,
+                shellController,
                 persistentRepository,
                 repositoryInitializer,
                 testScope,
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopImmersiveControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopImmersiveControllerTest.kt
index 447da8799e35..b87f20023796 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopImmersiveControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopImmersiveControllerTest.kt
@@ -89,7 +89,7 @@ class DesktopImmersiveControllerTest : ShellTestCase() {
     @Before
     fun setUp() {
         userRepositories = DesktopUserRepositories(
-            context, ShellInit(TestShellExecutor()), mock(), mock(), mock(), mock()
+            context, ShellInit(TestShellExecutor()), mock(), mock(), mock(), mock(), mock()
         )
         whenever(mockDisplayController.getDisplayLayout(DEFAULT_DISPLAY))
             .thenReturn(mockDisplayLayout)
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 33c14516d48c..cf16e42a729b 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
@@ -276,6 +276,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
       DesktopUserRepositories(
         context,
         shellInit,
+        shellController,
         persistentRepository,
         repositoryInitializer,
         testScope,
@@ -3907,7 +3908,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
 
   @Test
   fun shellController_registersUserChangeListener() {
-      verify(shellController, times(1)).addUserChangeListener(any())
+      verify(shellController, times(2)).addUserChangeListener(any())
   }
 
   @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt
index 0712d58166bb..39178cb2cd25 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt
@@ -42,6 +42,7 @@ import com.android.internal.jank.InteractionJankMonitor
 import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION
 import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.ShellTestCase
+import com.android.wm.shell.sysui.ShellController
 import com.android.wm.shell.common.ShellExecutor
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.createFreeformTask
 import com.android.wm.shell.desktopmode.persistence.DesktopPersistentRepository
@@ -96,6 +97,7 @@ class DesktopTasksLimiterTest : ShellTestCase() {
     @Mock lateinit var persistentRepository: DesktopPersistentRepository
     @Mock lateinit var repositoryInitializer: DesktopRepositoryInitializer
     @Mock lateinit var userManager: UserManager
+    @Mock lateinit var shellController: ShellController
 
     private lateinit var mockitoSession: StaticMockitoSession
     private lateinit var desktopTasksLimiter: DesktopTasksLimiter
@@ -117,6 +119,7 @@ class DesktopTasksLimiterTest : ShellTestCase() {
             DesktopUserRepositories(
                 context,
                 shellInit,
+                shellController,
                 persistentRepository,
                 repositoryInitializer,
                 testScope,
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopUserRepositoriesTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopUserRepositoriesTest.kt
index 5767df4c5a8e..a2e939d86adb 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopUserRepositoriesTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopUserRepositoriesTest.kt
@@ -24,14 +24,15 @@ import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
-import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
 import com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn
+import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
 import com.android.dx.mockito.inline.extended.StaticMockitoSession
 import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_HSUM
 import com.android.wm.shell.ShellTestCase
 import com.android.wm.shell.common.ShellExecutor
 import com.android.wm.shell.desktopmode.persistence.DesktopPersistentRepository
 import com.android.wm.shell.desktopmode.persistence.DesktopRepositoryInitializer
+import com.android.wm.shell.sysui.ShellController
 import com.android.wm.shell.sysui.ShellInit
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
@@ -66,6 +67,7 @@ class DesktopUserRepositoriesTest : ShellTestCase() {
     private val persistentRepository = mock<DesktopPersistentRepository>()
     private val repositoryInitializer = mock<DesktopRepositoryInitializer>()
     private val userManager = mock<UserManager>()
+    private val shellController = mock<ShellController>()
 
     @Before
     fun setUp() {
@@ -86,8 +88,14 @@ class DesktopUserRepositoriesTest : ShellTestCase() {
         whenever(userManager.getProfiles(USER_ID_1)).thenReturn(profiles)
 
         userRepositories = DesktopUserRepositories(
-            context, shellInit, persistentRepository, repositoryInitializer, datastoreScope,
-                userManager)
+            context,
+            shellInit,
+            shellController,
+            persistentRepository,
+            repositoryInitializer,
+            datastoreScope,
+            userManager
+        )
     }
 
     @After
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializerTest.kt
index 1c88a290d677..cdf064b075a1 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializerTest.kt
@@ -27,6 +27,7 @@ import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE
 import com.android.wm.shell.ShellTestCase
 import com.android.wm.shell.common.ShellExecutor
 import com.android.wm.shell.desktopmode.DesktopUserRepositories
+import com.android.wm.shell.sysui.ShellController
 import com.android.wm.shell.sysui.ShellInit
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
@@ -64,6 +65,7 @@ class DesktopRepositoryInitializerTest : ShellTestCase() {
     private val persistentRepository = mock<DesktopPersistentRepository>()
     private val userManager = mock<UserManager>()
     private val testExecutor = mock<ShellExecutor>()
+    private val shellController = mock<ShellController>()
 
     @Before
     fun setUp() {
@@ -74,7 +76,12 @@ class DesktopRepositoryInitializerTest : ShellTestCase() {
             DesktopRepositoryInitializerImpl(context, persistentRepository, datastoreScope)
         desktopUserRepositories =
             DesktopUserRepositories(
-                context, shellInit, persistentRepository, repositoryInitializer, datastoreScope,
+                context,
+                shellInit,
+                shellController,
+                persistentRepository,
+                repositoryInitializer,
+                datastoreScope,
                 userManager
             )
     }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopHeaderManageWindowsMenuTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopHeaderManageWindowsMenuTest.kt
index e871711fd25e..cf6c3a5e03a0 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopHeaderManageWindowsMenuTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopHeaderManageWindowsMenuTest.kt
@@ -63,6 +63,7 @@ class DesktopHeaderManageWindowsMenuTest : ShellTestCase() {
         userRepositories = DesktopUserRepositories(
             context = context,
             shellInit = ShellInit(TestShellExecutor()),
+            shellController = mock(),
             persistentRepository = mock(),
             repositoryInitializer = mock(),
             mainCoroutineScope = mock(),