Fix work tile setup when adding work profile

Test: Auto + manual
 - Add work profile -> QS tile is added
 - Remove work profile -> QS tile is removed
 - Reboot the device -> QS tile is in place
 - Reorder tiles, reboot the device -> order is persisted
Fixes: 234639083
Change-Id: I4481a3f28f9d1f0c30618afaf78179e0033f50fb
(cherry picked from commit a240e7eebe6fc81f123150a362c493a6f93b9d2c)
Merged-In: I4481a3f28f9d1f0c30618afaf78179e0033f50fb

8 files changed, 159 insertions, 50 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSHost.java b/packages/SystemUI/src/com/android/systemui/qs/QSHost.java
index 7cf63f678c1d..1da30ade951b 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSHost.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSHost.java
@@ -36,7 +36,6 @@ public interface QSHost {
     void removeCallback(Callback callback);
     void removeTile(String tileSpec);
     void removeTiles(Collection<String> specs);
-    void unmarkTileAsAutoAdded(String tileSpec);
 
     int indexOf(String tileSpec);
 
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSTileHost.java b/packages/SystemUI/src/com/android/systemui/qs/QSTileHost.java
index cad296b671b3..100853caa2d7 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSTileHost.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSTileHost.java
@@ -427,11 +427,6 @@ public class QSTileHost implements QSHost, Tunable, PluginListener<QSFactory>, P
         mMainExecutor.execute(() -> changeTileSpecs(tileSpecs -> tileSpecs.removeAll(specs)));
     }
 
-    @Override
-    public void unmarkTileAsAutoAdded(String spec) {
-        if (mAutoTiles != null) mAutoTiles.unmarkTileAsAutoAdded(spec);
-    }
-
     /**
      * Add a tile to the end
      *
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/WorkModeTile.java b/packages/SystemUI/src/com/android/systemui/qs/tiles/WorkModeTile.java
index a6c7781d891c..72c6bfe371ce 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/WorkModeTile.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/WorkModeTile.java
@@ -101,7 +101,6 @@ public class WorkModeTile extends QSTileImpl<BooleanState> implements
     @MainThread
     public void onManagedProfileRemoved() {
         mHost.removeTile(getTileSpec());
-        mHost.unmarkTileAsAutoAdded(getTileSpec());
     }
 
     @Override
diff --git a/packages/SystemUI/src/com/android/systemui/settings/UserTrackerImpl.kt b/packages/SystemUI/src/com/android/systemui/settings/UserTrackerImpl.kt
index 28da38b701bc..61390c582fd6 100644
--- a/packages/SystemUI/src/com/android/systemui/settings/UserTrackerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/settings/UserTrackerImpl.kt
@@ -112,7 +112,7 @@ class UserTrackerImpl internal constructor(
             // These get called when a managed profile goes in or out of quiet mode.
             addAction(Intent.ACTION_MANAGED_PROFILE_AVAILABLE)
             addAction(Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE)
-
+            addAction(Intent.ACTION_MANAGED_PROFILE_ADDED)
             addAction(Intent.ACTION_MANAGED_PROFILE_REMOVED)
             addAction(Intent.ACTION_MANAGED_PROFILE_UNLOCKED)
         }
@@ -129,6 +129,7 @@ class UserTrackerImpl internal constructor(
             Intent.ACTION_USER_INFO_CHANGED,
             Intent.ACTION_MANAGED_PROFILE_AVAILABLE,
             Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE,
+            Intent.ACTION_MANAGED_PROFILE_ADDED,
             Intent.ACTION_MANAGED_PROFILE_REMOVED,
             Intent.ACTION_MANAGED_PROFILE_UNLOCKED -> {
                 handleProfilesChanged()
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/AutoTileManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/AutoTileManager.java
index 9070eadd9944..149ec545dfa7 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/AutoTileManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/AutoTileManager.java
@@ -154,9 +154,7 @@ public class AutoTileManager implements UserAwareController {
         if (!mAutoTracker.isAdded(SAVER)) {
             mDataSaverController.addCallback(mDataSaverListener);
         }
-        if (!mAutoTracker.isAdded(WORK)) {
-            mManagedProfileController.addCallback(mProfileCallback);
-        }
+        mManagedProfileController.addCallback(mProfileCallback);
         if (!mAutoTracker.isAdded(NIGHT)
                 && ColorDisplayManager.isNightDisplayAvailable(mContext)) {
             mNightDisplayListener.setCallback(mNightDisplayCallback);
@@ -275,18 +273,18 @@ public class AutoTileManager implements UserAwareController {
         return mCurrentUser.getIdentifier();
     }
 
-    public void unmarkTileAsAutoAdded(String tabSpec) {
-        mAutoTracker.setTileRemoved(tabSpec);
-    }
-
     private final ManagedProfileController.Callback mProfileCallback =
             new ManagedProfileController.Callback() {
                 @Override
                 public void onManagedProfileChanged() {
-                    if (mAutoTracker.isAdded(WORK)) return;
                     if (mManagedProfileController.hasActiveProfile()) {
+                        if (mAutoTracker.isAdded(WORK)) return;
                         mHost.addTile(WORK);
                         mAutoTracker.setTileAdded(WORK);
+                    } else {
+                        if (!mAutoTracker.isAdded(WORK)) return;
+                        mHost.removeTile(WORK);
+                        mAutoTracker.setTileRemoved(WORK);
                     }
                 }
 
@@ -429,7 +427,7 @@ public class AutoTileManager implements UserAwareController {
                 initSafetyTile();
             } else if (!isSafetyCenterEnabled && mAutoTracker.isAdded(mSafetySpec)) {
                 mHost.removeTile(mSafetySpec);
-                mHost.unmarkTileAsAutoAdded(mSafetySpec);
+                mAutoTracker.setTileRemoved(mSafetySpec);
             }
         }
     };
diff --git a/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplReceiveTest.kt b/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplReceiveTest.kt
new file mode 100644
index 000000000000..3710281499b3
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplReceiveTest.kt
@@ -0,0 +1,100 @@
+package com.android.systemui.settings
+
+import android.content.Context
+import android.content.Intent
+import android.content.pm.UserInfo
+import android.os.Handler
+import android.os.UserHandle
+import android.os.UserManager
+import androidx.concurrent.futures.DirectExecutor
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.dump.DumpManager
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.capture
+import com.google.common.truth.Truth.assertThat
+import java.util.concurrent.Executor
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
+import org.mockito.ArgumentCaptor
+import org.mockito.ArgumentMatchers
+import org.mockito.ArgumentMatchers.anyInt
+import org.mockito.Captor
+import org.mockito.Mock
+import org.mockito.Mockito.times
+import org.mockito.Mockito.verify
+import org.mockito.Mockito.`when`
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+@RunWith(Parameterized::class)
+class UserTrackerImplReceiveTest : SysuiTestCase() {
+
+    companion object {
+
+        @JvmStatic
+        @Parameterized.Parameters
+        fun data(): Iterable<String> =
+            listOf(
+                Intent.ACTION_USER_INFO_CHANGED,
+                Intent.ACTION_MANAGED_PROFILE_AVAILABLE,
+                Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE,
+                Intent.ACTION_MANAGED_PROFILE_ADDED,
+                Intent.ACTION_MANAGED_PROFILE_REMOVED,
+                Intent.ACTION_MANAGED_PROFILE_UNLOCKED
+            )
+    }
+
+    private val executor: Executor = DirectExecutor.INSTANCE
+
+    @Mock private lateinit var context: Context
+    @Mock private lateinit var userManager: UserManager
+    @Mock(stubOnly = true) private lateinit var dumpManager: DumpManager
+    @Mock(stubOnly = true) private lateinit var handler: Handler
+
+    @Parameterized.Parameter lateinit var intentAction: String
+    @Mock private lateinit var callback: UserTracker.Callback
+    @Captor private lateinit var captor: ArgumentCaptor<List<UserInfo>>
+
+    private lateinit var tracker: UserTrackerImpl
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+
+        `when`(context.user).thenReturn(UserHandle.SYSTEM)
+        `when`(context.createContextAsUser(ArgumentMatchers.any(), anyInt())).thenReturn(context)
+
+        tracker = UserTrackerImpl(context, userManager, dumpManager, handler)
+    }
+
+    @Test
+    fun `calls callback and updates profiles when an intent received`() {
+        tracker.initialize(0)
+        tracker.addCallback(callback, executor)
+        val profileID = tracker.userId + 10
+
+        `when`(userManager.getProfiles(anyInt())).thenAnswer { invocation ->
+            val id = invocation.getArgument<Int>(0)
+            val info = UserInfo(id, "", UserInfo.FLAG_FULL)
+            val infoProfile =
+                UserInfo(
+                    id + 10,
+                    "",
+                    "",
+                    UserInfo.FLAG_MANAGED_PROFILE,
+                    UserManager.USER_TYPE_PROFILE_MANAGED
+                )
+            infoProfile.profileGroupId = id
+            listOf(info, infoProfile)
+        }
+
+        tracker.onReceive(context, Intent(intentAction))
+
+        verify(callback, times(0)).onUserChanged(anyInt(), any())
+        verify(callback, times(1)).onProfilesChanged(capture(captor))
+        assertThat(captor.value.map { it.id }).containsExactly(0, profileID)
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplTest.kt
index 52462c7186d4..e65bbb1bea08 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplTest.kt
@@ -124,6 +124,16 @@ class UserTrackerImplTest : SysuiTestCase() {
 
         verify(context).registerReceiverForAllUsers(
                 eq(tracker), capture(captor), isNull(), eq(handler))
+        with(captor.value) {
+            assertThat(countActions()).isEqualTo(7)
+            assertThat(hasAction(Intent.ACTION_USER_SWITCHED)).isTrue()
+            assertThat(hasAction(Intent.ACTION_USER_INFO_CHANGED)).isTrue()
+            assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_AVAILABLE)).isTrue()
+            assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE)).isTrue()
+            assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_ADDED)).isTrue()
+            assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_REMOVED)).isTrue()
+            assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_UNLOCKED)).isTrue()
+        }
     }
 
     @Test
@@ -280,37 +290,6 @@ class UserTrackerImplTest : SysuiTestCase() {
     }
 
     @Test
-    fun testCallbackCalledOnProfileChanged() {
-        tracker.initialize(0)
-        val callback = TestCallback()
-        tracker.addCallback(callback, executor)
-        val profileID = tracker.userId + 10
-
-        `when`(userManager.getProfiles(anyInt())).thenAnswer { invocation ->
-            val id = invocation.getArgument<Int>(0)
-            val info = UserInfo(id, "", UserInfo.FLAG_FULL)
-            val infoProfile = UserInfo(
-                    id + 10,
-                    "",
-                    "",
-                    UserInfo.FLAG_MANAGED_PROFILE,
-                    UserManager.USER_TYPE_PROFILE_MANAGED
-            )
-            infoProfile.profileGroupId = id
-            listOf(info, infoProfile)
-        }
-
-        val intent = Intent(Intent.ACTION_MANAGED_PROFILE_AVAILABLE)
-                .putExtra(Intent.EXTRA_USER, UserHandle.of(profileID))
-
-        tracker.onReceive(context, intent)
-
-        assertThat(callback.calledOnUserChanged).isEqualTo(0)
-        assertThat(callback.calledOnProfilesChanged).isEqualTo(1)
-        assertThat(callback.lastUserProfiles.map { it.id }).containsExactly(0, profileID)
-    }
-
-    @Test
     fun testCallbackCalledOnUserInfoChanged() {
         tracker.initialize(0)
         val callback = TestCallback()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/AutoTileManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/AutoTileManagerTest.java
index 4ccbc6d45e63..091bb5455d93 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/AutoTileManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/AutoTileManagerTest.java
@@ -24,6 +24,7 @@ import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.isNotNull;
 import static org.mockito.ArgumentMatchers.isNull;
 import static org.mockito.Mockito.doReturn;
@@ -74,6 +75,7 @@ import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 import org.mockito.Spy;
+import org.mockito.stubbing.Answer;
 
 import java.util.Collections;
 import java.util.List;
@@ -115,8 +117,10 @@ public class AutoTileManagerTest extends SysuiTestCase {
     @Spy private PackageManager mPackageManager;
     private final boolean mIsReduceBrightColorsAvailable = true;
 
-    private AutoTileManager mAutoTileManager;
+    private AutoTileManager mAutoTileManager; // under test
+
     private SecureSettings mSecureSettings;
+    private ManagedProfileController.Callback mManagedProfileCallback;
 
     @Before
     public void setUp() throws Exception {
@@ -303,7 +307,7 @@ public class AutoTileManagerTest extends SysuiTestCase {
 
         InOrder inOrderManagedProfile = inOrder(mManagedProfileController);
         inOrderManagedProfile.verify(mManagedProfileController).removeCallback(any());
-        inOrderManagedProfile.verify(mManagedProfileController, never()).addCallback(any());
+        inOrderManagedProfile.verify(mManagedProfileController).addCallback(any());
 
         if (ColorDisplayManager.isNightDisplayAvailable(mContext)) {
             InOrder inOrderNightDisplay = inOrder(mNightDisplayListener);
@@ -504,6 +508,40 @@ public class AutoTileManagerTest extends SysuiTestCase {
     }
 
     @Test
+    public void managedProfileAdded_tileAdded() {
+        when(mAutoAddTracker.isAdded(eq("work"))).thenReturn(false);
+        mAutoTileManager = createAutoTileManager(mContext);
+        Mockito.doAnswer((Answer<Object>) invocation -> {
+            mManagedProfileCallback = invocation.getArgument(0);
+            return null;
+        }).when(mManagedProfileController).addCallback(any());
+        mAutoTileManager.init();
+        when(mManagedProfileController.hasActiveProfile()).thenReturn(true);
+
+        mManagedProfileCallback.onManagedProfileChanged();
+
+        verify(mQsTileHost, times(1)).addTile(eq("work"));
+        verify(mAutoAddTracker, times(1)).setTileAdded(eq("work"));
+    }
+
+    @Test
+    public void managedProfileRemoved_tileRemoved() {
+        when(mAutoAddTracker.isAdded(eq("work"))).thenReturn(true);
+        mAutoTileManager = createAutoTileManager(mContext);
+        Mockito.doAnswer((Answer<Object>) invocation -> {
+            mManagedProfileCallback = invocation.getArgument(0);
+            return null;
+        }).when(mManagedProfileController).addCallback(any());
+        mAutoTileManager.init();
+        when(mManagedProfileController.hasActiveProfile()).thenReturn(false);
+
+        mManagedProfileCallback.onManagedProfileChanged();
+
+        verify(mQsTileHost, times(1)).removeTile(eq("work"));
+        verify(mAutoAddTracker, times(1)).setTileRemoved(eq("work"));
+    }
+
+    @Test
     public void testEmptyArray_doesNotCrash() {
         mContext.getOrCreateTestableResources().addOverride(
                 R.array.config_quickSettingsAutoAdd, new String[0]);