2 files changed, 76 insertions, 11 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
index c9011f439352..fdd929cae17d 100644
--- a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
+++ b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
@@ -138,19 +138,25 @@ public class ThemeOverlayController extends SystemUI implements Dumpable {
             mAcceptColorEvents = false;
         }
 
+        final boolean hadWallpaperColors = mSystemColors != null;
         if ((which & WallpaperManager.FLAG_SYSTEM) != 0) {
             mSystemColors = wallpaperColors;
-            if (DEBUG) {
-                Log.d(TAG, "got new lock colors: " + wallpaperColors + " where: " + which);
-            }
+            if (DEBUG) Log.d(TAG, "got new colors: " + wallpaperColors + " where: " + which);
         }
 
         if (mDeviceProvisionedController != null
                 && !mDeviceProvisionedController.isCurrentUserSetup()) {
-            Log.i(TAG, "Wallpaper color event deferred until setup is finished: "
-                    + wallpaperColors);
-            mDeferredThemeEvaluation = true;
-            return;
+            if (hadWallpaperColors) {
+                Log.i(TAG, "Wallpaper color event deferred until setup is finished: "
+                        + wallpaperColors);
+                mDeferredThemeEvaluation = true;
+                return;
+            } else {
+                if (DEBUG) {
+                    Log.i(TAG, "During user setup, but allowing first color event: had? "
+                            + hadWallpaperColors + " has? " + (mSystemColors != null));
+                }
+            }
         }
         reevaluateSystemTheme(false /* forceReload */);
     };
@@ -221,17 +227,31 @@ public class ThemeOverlayController extends SystemUI implements Dumpable {
                     }
                 },
                 UserHandle.USER_ALL);
+
+        if (!mIsMonetEnabled) {
+            return;
+        }
+
         mDeviceProvisionedController.addCallback(mDeviceProvisionedListener);
 
         // Upon boot, make sure we have the most up to date colors
-        mBgExecutor.execute(() -> {
+        Runnable updateColors = () -> {
             WallpaperColors systemColor = mWallpaperManager.getWallpaperColors(
                     WallpaperManager.FLAG_SYSTEM);
             mMainExecutor.execute(() -> {
+                if (DEBUG) Log.d(TAG, "Boot colors: " + systemColor);
                 mSystemColors = systemColor;
                 reevaluateSystemTheme(false /* forceReload */);
             });
-        });
+        };
+
+        // Whenever we're going directly to setup wizard, we need to process colors synchronously,
+        // otherwise we'll see some jank when the activity is recreated.
+        if (!mDeviceProvisionedController.isCurrentUserSetup()) {
+            mMainExecutor.execute(updateColors);
+        } else {
+            mBgExecutor.execute(updateColors);
+        }
         mWallpaperManager.addOnColorsChangedListener(mOnColorsChangedListener, null,
                 UserHandle.USER_ALL);
     }
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 ddf05372ebfa..ddf39d1e7c6e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java
@@ -57,6 +57,8 @@ import com.android.systemui.statusbar.policy.DeviceProvisionedController;
 import com.android.systemui.statusbar.policy.DeviceProvisionedController.DeviceProvisionedListener;
 import com.android.systemui.util.settings.SecureSettings;
 
+import com.google.common.util.concurrent.MoreExecutors;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -201,18 +203,61 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
     }
 
     @Test
-    public void onWallpaperColorsChanged_defersUntilSetupIsCompleted() {
+    public void onWallpaperColorsChanged_firstEventBeforeUserSetup_shouldBeAccepted() {
+        // By default, on setup() we make this controller return that the user finished setup
+        // wizard. This test on the other hand, is testing the setup flow.
         reset(mDeviceProvisionedController);
         WallpaperColors mainColors = new WallpaperColors(Color.valueOf(Color.RED),
                 Color.valueOf(Color.BLUE), null);
         mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
 
+        verify(mThemeOverlayApplier).applyCurrentUserOverlays(any(), any(), anyInt(), any());
+    }
+
+    @Test
+    public void onWallpaperColorsChanged_defersUntilSetupIsCompleted_ifHasColors() {
+        mDeviceProvisionedController = mock(DeviceProvisionedController.class);
+        mThemeOverlayApplier = mock(ThemeOverlayApplier.class);
+        mWallpaperManager = mock(WallpaperManager.class);
+
+        // Assume we have some wallpaper colors at boot.
+        when(mWallpaperManager.getWallpaperColors(anyInt()))
+                .thenReturn(new WallpaperColors(Color.valueOf(Color.GRAY), null, null));
+
+        Executor executor = MoreExecutors.directExecutor();
+        mThemeOverlayController = new ThemeOverlayController(null /* context */,
+                mBroadcastDispatcher, mBgHandler, executor, executor, mThemeOverlayApplier,
+                mSecureSettings, mWallpaperManager, mUserManager, mDeviceProvisionedController,
+                mUserTracker, mDumpManager, mFeatureFlags) {
+            @Nullable
+            @Override
+            protected FabricatedOverlay getOverlay(int color, int type) {
+                FabricatedOverlay overlay = mock(FabricatedOverlay.class);
+                when(overlay.getIdentifier())
+                        .thenReturn(new OverlayIdentifier(Integer.toHexString(color | 0xff000000)));
+                return overlay;
+            }
+        };
+        mThemeOverlayController.start();
+        verify(mWallpaperManager).addOnColorsChangedListener(mColorsListener.capture(), eq(null),
+                eq(UserHandle.USER_ALL));
+        verify(mDeviceProvisionedController).addCallback(mDeviceProvisionedListener.capture());
+
+        // Colors were applied during controller initialization.
+        verify(mThemeOverlayApplier).applyCurrentUserOverlays(any(), any(), anyInt(), any());
+        clearInvocations(mThemeOverlayApplier);
+
+        WallpaperColors mainColors = new WallpaperColors(Color.valueOf(Color.RED),
+                Color.valueOf(Color.BLUE), null);
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+
+        // Defers event because we already have initial colors.
         verify(mThemeOverlayApplier, never())
                 .applyCurrentUserOverlays(any(), any(), anyInt(), any());
 
+        // Then event happens after setup phase is over.
         when(mDeviceProvisionedController.isCurrentUserSetup()).thenReturn(true);
         mDeviceProvisionedListener.getValue().onUserSetupChanged();
-
         verify(mThemeOverlayApplier).applyCurrentUserOverlays(any(), any(), anyInt(), any());
     }