Allow first color event during SUW

SetupWizard should be able to use system colors, but should block any
other color events until the user is fully done with the flow.

Test: manual
Test: atest ThemeOverlayControllerTest
Fixes: 184704620
Change-Id: Iea17b653f3089351ae21274e1748bd3de21213d0

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());
     }