Merge "Correctly implement multi-user theming support" into sc-v2-dev am: 8f2f9a2ff1

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/16516205

Change-Id: Icc53036eb6b3a26408dcd6ed2b6e2e509a520e8c
diff --git a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
index 592fa15..70cb9d3 100644
--- 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.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 @@
     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 @@
     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 @@
                 }
             };
 
-    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 @@
         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 @@
             } 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 @@
         @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 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 @@
             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 @@
         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) {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java
index 3ff5666..b357c78 100644
--- 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 @@
 @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 @@
     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 @@
 
         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 @@
         // 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 @@
                 .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 @@
         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 @@
         // 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 @@
         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 @@
                 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 @@
                 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 @@
         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 @@
         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 @@
         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 @@
         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 @@
         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 @@
         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 @@
         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());
@@ -482,6 +511,34 @@
     }
 
     @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,
                 new Intent(Intent.ACTION_MANAGED_PROFILE_ADDED));
@@ -515,7 +572,8 @@
         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 @@
         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 @@
 
         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 @@
         // 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 @@
         // 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 @@
                 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);