Correctly implement multi-user theming support

WallpaperColors are now a per user map, and we use UserTracker to consult the userId. This avoids race conditions, given that all userIds are coming from the same controller. Test: manual Test: atest ThemeOverlayControllerTest Fixes: 206736090 Change-Id: I1735f483dc8d5913fd223fb5f13c5059e0ece2da (cherry picked from commit e4e2194e08)
2021-12-02 17:49:49 -08:00
parent cb1f78f054
commit ba48972396
2 changed files with 179 additions and 79 deletions
--- a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
+++ b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
@@ -48,6 +48,8 @@ import android.provider.Settings;
 import android.text.TextUtils;
 import android.util.ArrayMap;
 import android.util.Log;
+import android.util.SparseArray;
+import android.util.SparseIntArray;
 import android.util.TypedValue;

 import androidx.annotation.NonNull;
@@ -104,14 +106,15 @@ public class ThemeOverlayController extends SystemUI implements Dumpable {
    private final UserManager mUserManager;
    private final BroadcastDispatcher mBroadcastDispatcher;
    private final Executor mBgExecutor;
-    private SecureSettings mSecureSettings;
+    private final SecureSettings mSecureSettings;
    private final Executor mMainExecutor;
    private final Handler mBgHandler;
    private final boolean mIsMonetEnabled;
-    private UserTracker mUserTracker;
-    private DeviceProvisionedController mDeviceProvisionedController;
-    private WallpaperColors mCurrentColors;
-    private WallpaperManager mWallpaperManager;
+    private final UserTracker mUserTracker;
+    private final DeviceProvisionedController mDeviceProvisionedController;
+    // Current wallpaper colors associated to a user.
+    private final SparseArray<WallpaperColors> mCurrentColors = new SparseArray<>();
+    private final WallpaperManager mWallpaperManager;
    private ColorScheme mColorScheme;
    // If fabricated overlays were already created for the current theme.
    private boolean mNeedsOverlayCreation;
@@ -125,11 +128,11 @@ public class ThemeOverlayController extends SystemUI implements Dumpable {
    private FabricatedOverlay mNeutralOverlay;
    // If wallpaper color event will be accepted and change the UI colors.
    private boolean mAcceptColorEvents = true;
-    // If non-null, colors that were sent to the framework, and processing was deferred until
-    // the next time the screen is off.
-    private WallpaperColors mDeferredWallpaperColors;
-    private int mDeferredWallpaperColorsFlags;
-    private WakefulnessLifecycle mWakefulnessLifecycle;
+    // If non-null (per user), colors that were sent to the framework, and processing was deferred
+    // until the next time the screen is off.
+    private final SparseArray<WallpaperColors> mDeferredWallpaperColors = new SparseArray<>();
+    private final SparseIntArray mDeferredWallpaperColorsFlags = new SparseIntArray();
+    private final WakefulnessLifecycle mWakefulnessLifecycle;

    // Defers changing themes until Setup Wizard is done.
    private boolean mDeferredThemeEvaluation;
@@ -152,27 +155,53 @@ public class ThemeOverlayController extends SystemUI implements Dumpable {
                }
            };

-    private final OnColorsChangedListener mOnColorsChangedListener = (wallpaperColors, which) -> {
-        if (!mAcceptColorEvents && mWakefulnessLifecycle.getWakefulness() != WAKEFULNESS_ASLEEP) {
-            mDeferredWallpaperColors = wallpaperColors;
-            mDeferredWallpaperColorsFlags = which;
-            Log.i(TAG, "colors received; processing deferred until screen off: " + wallpaperColors);
-            return;
+    private final OnColorsChangedListener mOnColorsChangedListener = new OnColorsChangedListener() {
+        @Override
+        public void onColorsChanged(WallpaperColors wallpaperColors, int which) {
+            throw new IllegalStateException("This should never be invoked, all messages should "
+                    + "arrive on the overload that has a user id");
        }

-        if (wallpaperColors != null) {
-            mAcceptColorEvents = false;
-            // Any cache of colors deferred for process is now stale.
-            mDeferredWallpaperColors = null;
-            mDeferredWallpaperColorsFlags = 0;
-        }
+        @Override
+        public void onColorsChanged(WallpaperColors wallpaperColors, int which, int userId) {
+            boolean currentUser = userId == mUserTracker.getUserId();
+            if (currentUser && !mAcceptColorEvents
+                    && mWakefulnessLifecycle.getWakefulness() != WAKEFULNESS_ASLEEP) {
+                mDeferredWallpaperColors.put(userId, wallpaperColors);
+                mDeferredWallpaperColorsFlags.put(userId, which);
+                Log.i(TAG, "colors received; processing deferred until screen off: "
+                        + wallpaperColors + " user: " + userId);
+                return;
+            }

-        handleWallpaperColors(wallpaperColors, which);
+            if (currentUser && wallpaperColors != null) {
+                mAcceptColorEvents = false;
+                // Any cache of colors deferred for process is now stale.
+                mDeferredWallpaperColors.put(userId, null);
+                mDeferredWallpaperColorsFlags.put(userId, 0);
+            }
+
+            handleWallpaperColors(wallpaperColors, which, userId);
+        }
    };

-    private int getLatestWallpaperType() {
-        return mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_LOCK)
-                > mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_SYSTEM)
+    private final UserTracker.Callback mUserTrackerCallback = new UserTracker.Callback() {
+        @Override
+        public void onUserChanged(int newUser, @NonNull Context userContext) {
+            boolean isManagedProfile = mUserManager.isManagedProfile(newUser);
+            if (!mDeviceProvisionedController.isCurrentUserSetup() && isManagedProfile) {
+                Log.i(TAG, "User setup not finished when new user event was received. "
+                        + "Deferring... Managed profile? " + isManagedProfile);
+                return;
+            }
+            if (DEBUG) Log.d(TAG, "Updating overlays for user switch / profile added.");
+            reevaluateSystemTheme(true /* forceReload */);
+        }
+    };
+
+    private int getLatestWallpaperType(int userId) {
+        return mWallpaperManager.getWallpaperIdForUser(WallpaperManager.FLAG_LOCK, userId)
+                > mWallpaperManager.getWallpaperIdForUser(WallpaperManager.FLAG_SYSTEM, userId)
                ? WallpaperManager.FLAG_LOCK : WallpaperManager.FLAG_SYSTEM;
    }

@@ -204,14 +233,21 @@ public class ThemeOverlayController extends SystemUI implements Dumpable {
        return false;
    }

-    private void handleWallpaperColors(WallpaperColors wallpaperColors, int flags) {
-        final boolean hadWallpaperColors = mCurrentColors != null;
-        int latestWallpaperType = getLatestWallpaperType();
+    private void handleWallpaperColors(WallpaperColors wallpaperColors, int flags, int userId) {
+        final int currentUser = mUserTracker.getUserId();
+        final boolean hadWallpaperColors = mCurrentColors.get(userId) != null;
+        int latestWallpaperType = getLatestWallpaperType(userId);
        if ((flags & latestWallpaperType) != 0) {
-            mCurrentColors = wallpaperColors;
+            mCurrentColors.put(userId, wallpaperColors);
            if (DEBUG) Log.d(TAG, "got new colors: " + wallpaperColors + " where: " + flags);
        }

+        if (userId != currentUser) {
+            Log.d(TAG, "Colors " + wallpaperColors + " for user " + userId + ". "
+                    + "Not for current user: " + currentUser);
+            return;
+        }
+
        if (mDeviceProvisionedController != null
                && !mDeviceProvisionedController.isCurrentUserSetup()) {
            if (hadWallpaperColors) {
@@ -226,13 +262,12 @@ public class ThemeOverlayController extends SystemUI implements Dumpable {
            } else {
                if (DEBUG) {
                    Log.i(TAG, "During user setup, but allowing first color event: had? "
-                            + hadWallpaperColors + " has? " + (mCurrentColors != null));
+                            + hadWallpaperColors + " has? " + (mCurrentColors.get(userId) != null));
                }
            }
        }
        // Check if we need to reset to default colors (if a color override was set that is sourced
        // from the wallpaper)
-        int currentUser = mUserTracker.getUserId();
        String overlayPackageJson = mSecureSettings.getStringForUser(
                Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES,
                currentUser);
@@ -278,10 +313,9 @@ public class ThemeOverlayController extends SystemUI implements Dumpable {
        @Override
        public void onReceive(Context context, Intent intent) {
            boolean newWorkProfile = Intent.ACTION_MANAGED_PROFILE_ADDED.equals(intent.getAction());
-            boolean userStarted = Intent.ACTION_USER_SWITCHED.equals(intent.getAction());
            boolean isManagedProfile = mUserManager.isManagedProfile(
                    intent.getIntExtra(Intent.EXTRA_USER_HANDLE, 0));
-            if (userStarted || newWorkProfile) {
+            if (newWorkProfile) {
                if (!mDeviceProvisionedController.isCurrentUserSetup() && isManagedProfile) {
                    Log.i(TAG, "User setup not finished when " + intent.getAction()
                            + " was received. Deferring... Managed profile? " + isManagedProfile);
@@ -330,7 +364,6 @@ public class ThemeOverlayController extends SystemUI implements Dumpable {
    public void start() {
        if (DEBUG) Log.d(TAG, "Start");
        final IntentFilter filter = new IntentFilter();
-        filter.addAction(Intent.ACTION_USER_SWITCHED);
        filter.addAction(Intent.ACTION_MANAGED_PROFILE_ADDED);
        filter.addAction(Intent.ACTION_WALLPAPER_CHANGED);
        mBroadcastDispatcher.registerReceiver(mBroadcastReceiver, filter, mMainExecutor,
@@ -365,15 +398,17 @@ public class ThemeOverlayController extends SystemUI implements Dumpable {
            return;
        }

+        mUserTracker.addCallback(mUserTrackerCallback, mMainExecutor);
+
        mDeviceProvisionedController.addCallback(mDeviceProvisionedListener);

        // Upon boot, make sure we have the most up to date colors
        Runnable updateColors = () -> {
            WallpaperColors systemColor = mWallpaperManager.getWallpaperColors(
-                    getLatestWallpaperType());
+                    getLatestWallpaperType(mUserTracker.getUserId()));
            Runnable applyColors = () -> {
                if (DEBUG) Log.d(TAG, "Boot colors: " + systemColor);
-                mCurrentColors = systemColor;
+                mCurrentColors.put(mUserTracker.getUserId(), systemColor);
                reevaluateSystemTheme(false /* forceReload */);
            };
            if (mDeviceProvisionedController.isCurrentUserSetup()) {
@@ -395,21 +430,22 @@ public class ThemeOverlayController extends SystemUI implements Dumpable {
        mWakefulnessLifecycle.addObserver(new WakefulnessLifecycle.Observer() {
            @Override
            public void onFinishedGoingToSleep() {
-                if (mDeferredWallpaperColors != null) {
-                    WallpaperColors colors = mDeferredWallpaperColors;
-                    int flags = mDeferredWallpaperColorsFlags;
+                final int userId = mUserTracker.getUserId();
+                final WallpaperColors colors = mDeferredWallpaperColors.get(userId);
+                if (colors != null) {
+                    int flags = mDeferredWallpaperColorsFlags.get(userId);

-                    mDeferredWallpaperColors = null;
-                    mDeferredWallpaperColorsFlags = 0;
+                    mDeferredWallpaperColors.put(userId, null);
+                    mDeferredWallpaperColorsFlags.put(userId, 0);

-                    handleWallpaperColors(colors, flags);
+                    handleWallpaperColors(colors, flags, userId);
                }
            }
        });
    }

    private void reevaluateSystemTheme(boolean forceReload) {
-        final WallpaperColors currentColors = mCurrentColors;
+        final WallpaperColors currentColors = mCurrentColors.get(mUserTracker.getUserId());
        final int mainColor;
        final int accentCandidate;
        if (currentColors == null) {
--- a/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java
@@ -76,6 +76,9 @@ import java.util.concurrent.Executor;
@RunWith(AndroidTestingRunner.class)
 public class ThemeOverlayControllerTest extends SysuiTestCase {

+    private static final int USER_SYSTEM = UserHandle.USER_SYSTEM;
+    private static final int USER_SECONDARY = 10;
+
    private ThemeOverlayController mThemeOverlayController;
    @Mock
    private Executor mBgExecutor;
@@ -111,6 +114,9 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
    private ArgumentCaptor<DeviceProvisionedListener> mDeviceProvisionedListener;
    @Captor
    private ArgumentCaptor<WakefulnessLifecycle.Observer> mWakefulnessLifecycleObserver;
+    @Captor
+    private ArgumentCaptor<UserTracker.Callback> mUserTrackerCallback;
+
    @Before
    public void setup() {
        MockitoAnnotations.initMocks(this);
@@ -133,6 +139,7 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {

        mWakefulnessLifecycle.dispatchFinishedWakingUp();
        mThemeOverlayController.start();
+        verify(mUserTracker).addCallback(mUserTrackerCallback.capture(), eq(mMainExecutor));
        verify(mWallpaperManager).addOnColorsChangedListener(mColorsListener.capture(), eq(null),
                eq(UserHandle.USER_ALL));
        verify(mBroadcastDispatcher).registerReceiver(mBroadcastReceiver.capture(), any(),
@@ -156,7 +163,8 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        // Should ask for a new theme when wallpaper colors change
        WallpaperColors mainColors = new WallpaperColors(Color.valueOf(Color.RED),
                Color.valueOf(Color.BLUE), null);
-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);
        ArgumentCaptor<Map<String, OverlayIdentifier>> themeOverlays =
                ArgumentCaptor.forClass(Map.class);

@@ -170,12 +178,13 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
                .isEqualTo(new OverlayIdentifier("ffff0000"));

        // Should not ask again if changed to same value
-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);
        verifyNoMoreInteractions(mThemeOverlayApplier);

        // Should not ask again even for new colors until we change wallpapers
        mColorsListener.getValue().onColorsChanged(new WallpaperColors(Color.valueOf(Color.BLACK),
-                null, null), WallpaperManager.FLAG_SYSTEM);
+                null, null), WallpaperManager.FLAG_SYSTEM, USER_SYSTEM);
        verifyNoMoreInteractions(mThemeOverlayApplier);

        // But should change theme after changing wallpapers
@@ -184,7 +193,7 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        intent.putExtra(WallpaperManager.EXTRA_FROM_FOREGROUND_APP, true);
        mBroadcastReceiver.getValue().onReceive(null, intent);
        mColorsListener.getValue().onColorsChanged(new WallpaperColors(Color.valueOf(Color.BLACK),
-                null, null), WallpaperManager.FLAG_SYSTEM);
+                null, null), WallpaperManager.FLAG_SYSTEM, USER_SYSTEM);
        verify(mThemeOverlayApplier).applyCurrentUserOverlays(any(), any(), anyInt(), any());
    }

@@ -193,7 +202,8 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        // Should ask for a new theme when wallpaper colors change
        WallpaperColors mainColors = new WallpaperColors(Color.valueOf(Color.RED),
                Color.valueOf(Color.BLUE), null);
-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);
        ArgumentCaptor<Map<String, OverlayIdentifier>> themeOverlays =
                ArgumentCaptor.forClass(Map.class);

@@ -211,7 +221,7 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        clearInvocations(mThemeOverlayApplier);
        mBroadcastReceiver.getValue().onReceive(null, new Intent(Intent.ACTION_WALLPAPER_CHANGED));
        mColorsListener.getValue().onColorsChanged(new WallpaperColors(Color.valueOf(Color.BLACK),
-                null, null), WallpaperManager.FLAG_SYSTEM);
+                null, null), WallpaperManager.FLAG_SYSTEM, USER_SYSTEM);
        verify(mThemeOverlayApplier, never())
                .applyCurrentUserOverlays(any(), any(), anyInt(), any());
    }
@@ -229,7 +239,8 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), anyInt()))
                .thenReturn(jsonString);

-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);
        ArgumentCaptor<Map<String, OverlayIdentifier>> themeOverlays =
                ArgumentCaptor.forClass(Map.class);

@@ -257,7 +268,8 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), anyInt()))
                .thenReturn(jsonString);

-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);

        ArgumentCaptor<String> updatedSetting = ArgumentCaptor.forClass(String.class);
        verify(mSecureSettings).putString(
@@ -289,10 +301,13 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        when(mSecureSettings.getStringForUser(
                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), anyInt()))
                .thenReturn(jsonString);
-        when(mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_LOCK)).thenReturn(20);
-        when(mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_SYSTEM)).thenReturn(21);
+        when(mWallpaperManager.getWallpaperIdForUser(WallpaperManager.FLAG_LOCK, USER_SYSTEM))
+                .thenReturn(20);
+        when(mWallpaperManager.getWallpaperIdForUser(WallpaperManager.FLAG_SYSTEM, USER_SYSTEM))
+                .thenReturn(21);

-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);

        ArgumentCaptor<String> updatedSetting = ArgumentCaptor.forClass(String.class);
        verify(mSecureSettings).putString(
@@ -320,10 +335,11 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        when(mSecureSettings.getStringForUser(
                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), anyInt()))
                .thenReturn(jsonString);
-        when(mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_LOCK)).thenReturn(-1);
+        when(mWallpaperManager.getWallpaperIdForUser(WallpaperManager.FLAG_LOCK, USER_SYSTEM))
+                .thenReturn(-1);

        mColorsListener.getValue().onColorsChanged(mainColors,
-                WallpaperManager.FLAG_SYSTEM | WallpaperManager.FLAG_LOCK);
+                WallpaperManager.FLAG_SYSTEM | WallpaperManager.FLAG_LOCK, USER_SYSTEM);

        ArgumentCaptor<String> updatedSetting = ArgumentCaptor.forClass(String.class);
        verify(mSecureSettings).putString(
@@ -349,9 +365,11 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        when(mSecureSettings.getStringForUser(
                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), anyInt()))
                .thenReturn(jsonString);
-        when(mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_LOCK)).thenReturn(1);
+        when(mWallpaperManager.getWallpaperIdForUser(WallpaperManager.FLAG_LOCK, USER_SYSTEM))
+                .thenReturn(1);

-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_LOCK);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_LOCK,
+                USER_SYSTEM);

        ArgumentCaptor<String> updatedSetting = ArgumentCaptor.forClass(String.class);
        verify(mSecureSettings).putString(
@@ -377,9 +395,11 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        when(mSecureSettings.getStringForUser(
                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), anyInt()))
                .thenReturn(jsonString);
-        when(mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_LOCK)).thenReturn(-1);
+        when(mWallpaperManager.getWallpaperIdForUser(WallpaperManager.FLAG_LOCK, USER_SYSTEM))
+                .thenReturn(-1);

-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);

        ArgumentCaptor<String> updatedSetting = ArgumentCaptor.forClass(String.class);
        verify(mSecureSettings).putString(
@@ -407,11 +427,14 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        when(mSecureSettings.getStringForUser(
                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), anyInt()))
                .thenReturn(jsonString);
-        when(mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_LOCK)).thenReturn(1);
+        when(mWallpaperManager.getWallpaperIdForUser(WallpaperManager.FLAG_LOCK, USER_SYSTEM))
+                .thenReturn(1);
        // SYSTEM wallpaper is the last applied one
-        when(mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_SYSTEM)).thenReturn(2);
+        when(mWallpaperManager.getWallpaperIdForUser(WallpaperManager.FLAG_SYSTEM, USER_SYSTEM))
+                .thenReturn(2);

-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);

        ArgumentCaptor<String> updatedSetting = ArgumentCaptor.forClass(String.class);
        verify(mSecureSettings).putString(
@@ -437,11 +460,14 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        when(mSecureSettings.getStringForUser(
                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), anyInt()))
                .thenReturn(jsonString);
-        when(mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_LOCK)).thenReturn(1);
+        when(mWallpaperManager.getWallpaperIdForUser(WallpaperManager.FLAG_LOCK, USER_SYSTEM))
+                .thenReturn(1);
        // SYSTEM wallpaper is the last applied one
-        when(mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_SYSTEM)).thenReturn(2);
+        when(mWallpaperManager.getWallpaperIdForUser(WallpaperManager.FLAG_SYSTEM, USER_SYSTEM))
+                .thenReturn(2);

-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);

        ArgumentCaptor<String> updatedSetting = ArgumentCaptor.forClass(String.class);
        verify(mSecureSettings, never()).putString(
@@ -467,11 +493,14 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        when(mSecureSettings.getStringForUser(
                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), anyInt()))
                .thenReturn(jsonString);
-        when(mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_LOCK)).thenReturn(1);
+        when(mWallpaperManager.getWallpaperIdForUser(WallpaperManager.FLAG_LOCK, USER_SYSTEM))
+                .thenReturn(1);
        // SYSTEM wallpaper is the last applied one
-        when(mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_SYSTEM)).thenReturn(2);
+        when(mWallpaperManager.getWallpaperIdForUser(WallpaperManager.FLAG_SYSTEM, USER_SYSTEM))
+                .thenReturn(2);

-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_LOCK);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_LOCK,
+                USER_SYSTEM);

        verify(mSecureSettings, never()).putString(
                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), any());
@@ -481,6 +510,34 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
                .applyCurrentUserOverlays(any(), any(), anyInt(), any());
    }

+    @Test
+    public void onUserSwitching_setsTheme() {
+        // Setup users with different colors
+        WallpaperColors mainColors = new WallpaperColors(Color.valueOf(Color.RED), null, null);
+        WallpaperColors secondaryColors =
+                new WallpaperColors(Color.valueOf(Color.BLUE), null, null);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(secondaryColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SECONDARY);
+
+        // When changing users
+        clearInvocations(mThemeOverlayApplier);
+        when(mUserTracker.getUserId()).thenReturn(USER_SECONDARY);
+        mUserTrackerCallback.getValue().onUserChanged(USER_SECONDARY, mContext);
+
+        ArgumentCaptor<Map<String, OverlayIdentifier>> themeOverlays =
+                ArgumentCaptor.forClass(Map.class);
+        verify(mThemeOverlayApplier)
+                .applyCurrentUserOverlays(themeOverlays.capture(), any(), anyInt(), any());
+
+        // Assert that we received secondary user colors
+        assertThat(themeOverlays.getValue().get(OVERLAY_CATEGORY_SYSTEM_PALETTE))
+                .isEqualTo(new OverlayIdentifier("ff0000ff"));
+        assertThat(themeOverlays.getValue().get(OVERLAY_CATEGORY_ACCENT_COLOR))
+                .isEqualTo(new OverlayIdentifier("ff0000ff"));
+    }
+
    @Test
    public void onProfileAdded_setsTheme() {
        mBroadcastReceiver.getValue().onReceive(null,
@@ -515,7 +572,8 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        reset(mDeviceProvisionedController);
        WallpaperColors mainColors = new WallpaperColors(Color.valueOf(Color.RED),
                Color.valueOf(Color.BLUE), null);
-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);

        verify(mThemeOverlayApplier).applyCurrentUserOverlays(any(), any(), anyInt(), any());

@@ -525,11 +583,11 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        Intent intent = new Intent(Intent.ACTION_WALLPAPER_CHANGED);
        intent.putExtra(WallpaperManager.EXTRA_FROM_FOREGROUND_APP, true);
        mBroadcastReceiver.getValue().onReceive(null, intent);
-        mColorsListener.getValue().onColorsChanged(null, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(null, WallpaperManager.FLAG_SYSTEM, USER_SYSTEM);
        verify(mThemeOverlayApplier, never()).applyCurrentUserOverlays(any(), any(), anyInt(),
                any());
        mColorsListener.getValue().onColorsChanged(new WallpaperColors(Color.valueOf(Color.GREEN),
-                null, null), WallpaperManager.FLAG_SYSTEM);
+                null, null), WallpaperManager.FLAG_SYSTEM, USER_SYSTEM);
        verify(mThemeOverlayApplier, never()).applyCurrentUserOverlays(any(), any(), anyInt(),
                any());
    }
@@ -607,7 +665,8 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {

        WallpaperColors mainColors = new WallpaperColors(Color.valueOf(Color.RED),
                Color.valueOf(Color.BLUE), null);
-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);

        // Defers event because we already have initial colors.
        verify(mThemeOverlayApplier, never())
@@ -628,14 +687,16 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        // Second color application is not applied.
        WallpaperColors mainColors = new WallpaperColors(Color.valueOf(Color.RED),
                Color.valueOf(Color.BLUE), null);
-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);

        clearInvocations(mThemeOverlayApplier);

        // Device went to sleep and second set of colors was applied.
        mainColors =  new WallpaperColors(Color.valueOf(Color.BLUE),
                Color.valueOf(Color.RED), null);
-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);
        verify(mThemeOverlayApplier, never())
                .applyCurrentUserOverlays(any(), any(), anyInt(), any());

@@ -652,14 +713,16 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
        // Second color application is not applied.
        WallpaperColors mainColors = new WallpaperColors(Color.valueOf(Color.RED),
                Color.valueOf(Color.BLUE), null);
-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);

        clearInvocations(mThemeOverlayApplier);

        // Device went to sleep and second set of colors was applied.
        mainColors =  new WallpaperColors(Color.valueOf(Color.BLUE),
                Color.valueOf(Color.RED), null);
-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);
        verify(mThemeOverlayApplier, never())
                .applyCurrentUserOverlays(any(), any(), anyInt(), any());

@@ -678,7 +741,8 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), anyInt()))
                .thenReturn(jsonString);

-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM,
+                USER_SYSTEM);
        ArgumentCaptor<Map<String, OverlayIdentifier>> themeOverlays =
                ArgumentCaptor.forClass(Map.class);