From a476c79c3a7c22a839b109dc166112f88cc7952c Mon Sep 17 00:00:00 2001
From: Lucas Dupin <dupin@google.com>
Date: Mon, 10 Jun 2019 18:00:22 -0700
Subject: [PATCH] Fix unreadable clock

- Always use white text with shadows on media
- Use lock screen colors for shade theme
- Remove legacy wallpaper visibility listener
- Remove shadow from black title

Test: set white wallpaper
Test: set dark wallpaper
Test: set dark wallpaper on home, light on ls
Test: set dark wallpaper on ls, light on home
Test: atest SysuiColorExtractorTests
Fixes: 134119878
Change-Id: I64ee2dee14322b2c02891c3ffe5fc0b25d1f100b
Merged-In: I64ee2dee14322b2c02891c3ffe5fc0b25d1f100b
---
 .../colorextraction/ColorExtractor.java       |   9 +-
 .../systemui/recents/RecentsActivity.java     |   4 +-
 .../SystemUI/res-keyguard/values/styles.xml   |   2 +-
 .../android/keyguard/KeyguardClockSwitch.java |   2 +-
 .../keyguard/clock/AnalogClockController.java |   2 +-
 .../keyguard/clock/BubbleClockController.java |   2 +-
 .../clock/DefaultClockController.java         |   2 +-
 .../colorextraction/SysuiColorExtractor.java  | 166 +++++-------------
 .../statusbar/NotificationMediaManager.java   |   3 +-
 .../stack/NotificationStackScrollLayout.java  |  10 +-
 .../systemui/statusbar/phone/StatusBar.java   |   3 +-
 .../phone/StatusBarWindowController.java      |  13 +-
 .../SysuiColorExtractorTests.java             |  29 +--
 .../colorextraction/ColorExtractorTest.java   |  12 +-
 14 files changed, 72 insertions(+), 187 deletions(-)

diff --git a/core/java/com/android/internal/colorextraction/ColorExtractor.java b/core/java/com/android/internal/colorextraction/ColorExtractor.java
index d9fd3b5bd6d82..b27c11b524e50 100644
--- a/core/java/com/android/internal/colorextraction/ColorExtractor.java
+++ b/core/java/com/android/internal/colorextraction/ColorExtractor.java
@@ -53,11 +53,13 @@ public class ColorExtractor implements WallpaperManager.OnColorsChangedListener
     protected WallpaperColors mLockColors;
 
     public ColorExtractor(Context context) {
-        this(context, new Tonal(context), true /* immediately */);
+        this(context, new Tonal(context), true /* immediately */,
+                context.getSystemService(WallpaperManager.class));
     }
 
     @VisibleForTesting
-    public ColorExtractor(Context context, ExtractionType extractionType, boolean immediately) {
+    public ColorExtractor(Context context, ExtractionType extractionType, boolean immediately,
+            WallpaperManager wallpaperManager) {
         mContext = context;
         mExtractionType = extractionType;
 
@@ -72,7 +74,6 @@ public class ColorExtractor implements WallpaperManager.OnColorsChangedListener
 
         mOnColorsChangedListeners = new ArrayList<>();
 
-        WallpaperManager wallpaperManager = mContext.getSystemService(WallpaperManager.class);
         if (wallpaperManager == null) {
             Log.w(TAG, "Can't listen to color changes!");
         } else {
@@ -110,7 +111,7 @@ public class ColorExtractor implements WallpaperManager.OnColorsChangedListener
         }
     }
 
-    private void extractWallpaperColors() {
+    protected void extractWallpaperColors() {
         GradientColors[] systemColors = mGradientColors.get(WallpaperManager.FLAG_SYSTEM);
         GradientColors[] lockColors = mGradientColors.get(WallpaperManager.FLAG_LOCK);
         extractInto(mSystemColors,
diff --git a/packages/SystemUI/legacy/recents/src/com/android/systemui/recents/RecentsActivity.java b/packages/SystemUI/legacy/recents/src/com/android/systemui/recents/RecentsActivity.java
index 79c691cf45e13..a7ccc3a490734 100644
--- a/packages/SystemUI/legacy/recents/src/com/android/systemui/recents/RecentsActivity.java
+++ b/packages/SystemUI/legacy/recents/src/com/android/systemui/recents/RecentsActivity.java
@@ -323,7 +323,7 @@ public class RecentsActivity extends Activity implements ViewTreeObserver.OnPreD
         mColorExtractor = Dependency.get(SysuiColorExtractor.class);
         mColorExtractor.addOnColorsChangedListener(this);
         mUsingDarkText = mColorExtractor.getColors(ColorExtractor.TYPE_DARK,
-                WallpaperManager.FLAG_SYSTEM, true).supportsDarkText();
+                WallpaperManager.FLAG_SYSTEM).supportsDarkText();
         setTheme(mUsingDarkText ? R.style.RecentsTheme_Wallpaper_Light
                 : R.style.RecentsTheme_Wallpaper);
 
@@ -394,8 +394,6 @@ public class RecentsActivity extends Activity implements ViewTreeObserver.OnPreD
     @Override
     public void onColorsChanged(ColorExtractor colorExtractor, int which) {
         if ((which & WallpaperManager.FLAG_SYSTEM) != 0) {
-            // Recents doesn't care about the wallpaper being visible or not, it always
-            // wants to scrim with wallpaper colors
             ColorExtractor.GradientColors colors = mColorExtractor.getNeutralColors();
             boolean darkText = colors.supportsDarkText();
             if (darkText != mUsingDarkText) {
diff --git a/packages/SystemUI/res-keyguard/values/styles.xml b/packages/SystemUI/res-keyguard/values/styles.xml
index 9a042228435b2..67c4458de2f25 100644
--- a/packages/SystemUI/res-keyguard/values/styles.xml
+++ b/packages/SystemUI/res-keyguard/values/styles.xml
@@ -111,7 +111,7 @@
         <item name="android:colorBackground">@*android:color/background_material_dark</item>
     </style>
 
-    <style name="TextAppearance.Keyguard" parent="Theme.SystemUI">
+    <style name="TextAppearance.Keyguard">
         <item name="android:textSize">@dimen/widget_title_font_size</item>
         <item name="android:gravity">center</item>
         <item name="android:ellipsize">end</item>
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardClockSwitch.java b/packages/SystemUI/src/com/android/keyguard/KeyguardClockSwitch.java
index 0bb9e744f2b62..7ec1bda26cf83 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardClockSwitch.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardClockSwitch.java
@@ -374,7 +374,7 @@ public class KeyguardClockSwitch extends RelativeLayout {
 
     private void updateColors() {
         ColorExtractor.GradientColors colors = mSysuiColorExtractor.getColors(
-                WallpaperManager.FLAG_LOCK, true);
+                WallpaperManager.FLAG_LOCK);
         mSupportsDarkText = colors.supportsDarkText();
         mColorPalette = colors.getColorPalette();
         if (mClockPlugin != null) {
diff --git a/packages/SystemUI/src/com/android/keyguard/clock/AnalogClockController.java b/packages/SystemUI/src/com/android/keyguard/clock/AnalogClockController.java
index f468ecaae4c1b..558ac4b564d15 100644
--- a/packages/SystemUI/src/com/android/keyguard/clock/AnalogClockController.java
+++ b/packages/SystemUI/src/com/android/keyguard/clock/AnalogClockController.java
@@ -130,7 +130,7 @@ public class AnalogClockController implements ClockPlugin {
         setDarkAmount(1f);
         setTextColor(Color.WHITE);
         ColorExtractor.GradientColors colors = mColorExtractor.getColors(
-                WallpaperManager.FLAG_LOCK, true);
+                WallpaperManager.FLAG_LOCK);
         setColorPalette(colors.supportsDarkText(), colors.getColorPalette());
         onTimeTick();
 
diff --git a/packages/SystemUI/src/com/android/keyguard/clock/BubbleClockController.java b/packages/SystemUI/src/com/android/keyguard/clock/BubbleClockController.java
index 61a6952cacac1..bdf9dc4865b2d 100644
--- a/packages/SystemUI/src/com/android/keyguard/clock/BubbleClockController.java
+++ b/packages/SystemUI/src/com/android/keyguard/clock/BubbleClockController.java
@@ -130,7 +130,7 @@ public class BubbleClockController implements ClockPlugin {
         setDarkAmount(1f);
         setTextColor(Color.WHITE);
         ColorExtractor.GradientColors colors = mColorExtractor.getColors(
-                WallpaperManager.FLAG_LOCK, true);
+                WallpaperManager.FLAG_LOCK);
         setColorPalette(colors.supportsDarkText(), colors.getColorPalette());
         onTimeTick();
 
diff --git a/packages/SystemUI/src/com/android/keyguard/clock/DefaultClockController.java b/packages/SystemUI/src/com/android/keyguard/clock/DefaultClockController.java
index ce1f09c9355c4..98679ade10225 100644
--- a/packages/SystemUI/src/com/android/keyguard/clock/DefaultClockController.java
+++ b/packages/SystemUI/src/com/android/keyguard/clock/DefaultClockController.java
@@ -124,7 +124,7 @@ public class DefaultClockController implements ClockPlugin {
         setDarkAmount(1f);
         setTextColor(Color.WHITE);
         ColorExtractor.GradientColors colors = mColorExtractor.getColors(
-                WallpaperManager.FLAG_LOCK, true);
+                WallpaperManager.FLAG_LOCK);
         setColorPalette(colors.supportsDarkText(), colors.getColorPalette());
         onTimeTick();
 
diff --git a/packages/SystemUI/src/com/android/systemui/colorextraction/SysuiColorExtractor.java b/packages/SystemUI/src/com/android/systemui/colorextraction/SysuiColorExtractor.java
index 835ffc976e9f7..6f56a53c1c494 100644
--- a/packages/SystemUI/src/com/android/systemui/colorextraction/SysuiColorExtractor.java
+++ b/packages/SystemUI/src/com/android/systemui/colorextraction/SysuiColorExtractor.java
@@ -16,18 +16,11 @@
 
 package com.android.systemui.colorextraction;
 
-import android.annotation.ColorInt;
 import android.app.WallpaperColors;
 import android.app.WallpaperManager;
 import android.content.Context;
-import android.os.Handler;
-import android.os.RemoteException;
+import android.graphics.Color;
 import android.os.UserHandle;
-import android.util.Log;
-import android.view.Display;
-import android.view.IWallpaperVisibilityListener;
-import android.view.IWindowManager;
-import android.view.WindowManagerGlobal;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.colorextraction.ColorExtractor;
@@ -52,46 +45,28 @@ public class SysuiColorExtractor extends ColorExtractor implements Dumpable,
         ConfigurationController.ConfigurationListener {
     private static final String TAG = "SysuiColorExtractor";
     private final Tonal mTonal;
-    private boolean mWallpaperVisible;
-    private boolean mHasBackdrop;
-    // Colors to return when the wallpaper isn't visible
-    private final GradientColors mWpHiddenColors;
+    private boolean mHasMediaArtwork;
+    private final GradientColors mNeutralColorsLock;
+    private final GradientColors mBackdropColors;
 
     @Inject
     public SysuiColorExtractor(Context context, ConfigurationController configurationController) {
-        this(context, new Tonal(context), configurationController, true);
+        this(context, new Tonal(context), configurationController,
+                context.getSystemService(WallpaperManager.class), false /* immediately */);
     }
 
     @VisibleForTesting
     public SysuiColorExtractor(Context context, ExtractionType type,
-            ConfigurationController configurationController, boolean registerVisibility) {
-        super(context, type, false /* immediately */);
+            ConfigurationController configurationController,
+            WallpaperManager wallpaperManager, boolean immediately) {
+        super(context, type, immediately, wallpaperManager);
         mTonal = type instanceof Tonal ? (Tonal) type : new Tonal(context);
-        mWpHiddenColors = new GradientColors();
+        mNeutralColorsLock = new GradientColors();
         configurationController.addCallback(this);
 
-        WallpaperColors systemColors = getWallpaperColors(WallpaperManager.FLAG_SYSTEM);
-        updateDefaultGradients(systemColors);
+        mBackdropColors = new GradientColors();
+        mBackdropColors.setMainColor(Color.BLACK);
 
-        if (registerVisibility) {
-            try {
-                IWindowManager windowManagerService = WindowManagerGlobal.getWindowManagerService();
-                Handler handler = Handler.getMain();
-                boolean visible = windowManagerService.registerWallpaperVisibilityListener(
-                        new IWallpaperVisibilityListener.Stub() {
-                            @Override
-                            public void onWallpaperVisibilityChanged(boolean newVisibility,
-                                    int displayId) throws RemoteException {
-                                handler.post(() -> setWallpaperVisible(newVisibility));
-                            }
-                        }, Display.DEFAULT_DISPLAY);
-                setWallpaperVisible(visible);
-            } catch (RemoteException e) {
-                Log.w(TAG, "Can't listen to wallpaper visibility changes", e);
-            }
-        }
-
-        WallpaperManager wallpaperManager = context.getSystemService(WallpaperManager.class);
         if (wallpaperManager != null) {
             // Listen to all users instead of only the current one.
             wallpaperManager.removeOnColorsChangedListener(this);
@@ -100,8 +75,14 @@ public class SysuiColorExtractor extends ColorExtractor implements Dumpable,
         }
     }
 
-    private void updateDefaultGradients(WallpaperColors colors) {
-        mTonal.applyFallback(colors, mWpHiddenColors);
+    @Override
+    protected void extractWallpaperColors() {
+        super.extractWallpaperColors();
+        // mTonal is final but this method will be invoked by the base class during its ctor.
+        if (mTonal == null) {
+            return;
+        }
+        mTonal.applyFallback(mLockColors == null ? mSystemColors : mLockColors, mNeutralColorsLock);
     }
 
     @Override
@@ -110,27 +91,28 @@ public class SysuiColorExtractor extends ColorExtractor implements Dumpable,
             // Colors do not belong to current user, ignoring.
             return;
         }
-
-        super.onColorsChanged(colors, which);
-
-        if ((which & WallpaperManager.FLAG_SYSTEM) != 0) {
-            @ColorInt int oldColor = mWpHiddenColors.getMainColor();
-            updateDefaultGradients(colors);
-            if (oldColor != mWpHiddenColors.getMainColor()) {
-                triggerColorsChanged(WallpaperManager.FLAG_SYSTEM);
-            }
+        if ((which & WallpaperManager.FLAG_LOCK) != 0) {
+            mTonal.applyFallback(colors, mNeutralColorsLock);
         }
+        super.onColorsChanged(colors, which);
     }
 
     @Override
     public void onUiModeChanged() {
-        WallpaperColors systemColors = getWallpaperColors(WallpaperManager.FLAG_SYSTEM);
-        updateDefaultGradients(systemColors);
-        triggerColorsChanged(WallpaperManager.FLAG_SYSTEM);
+        extractWallpaperColors();
+        triggerColorsChanged(WallpaperManager.FLAG_SYSTEM | WallpaperManager.FLAG_LOCK);
+    }
+
+    @Override
+    public GradientColors getColors(int which, int type) {
+        if (mHasMediaArtwork && (which & WallpaperManager.FLAG_LOCK) != 0) {
+            return mBackdropColors;
+        }
+        return super.getColors(which, type);
     }
 
     /**
-     * Colors the should be using for scrims.
+     * Colors that should be using for scrims.
      *
      * They will be:
      * - A light gray if the wallpaper is light
@@ -138,81 +120,12 @@ public class SysuiColorExtractor extends ColorExtractor implements Dumpable,
      * - Black otherwise
      */
     public GradientColors getNeutralColors() {
-        return mWpHiddenColors;
+        return mHasMediaArtwork ? mBackdropColors : mNeutralColorsLock;
     }
 
-    /**
-     * Get TYPE_NORMAL colors when wallpaper is visible, or fallback otherwise.
-     *
-     * @param which FLAG_LOCK or FLAG_SYSTEM
-     * @return colors
-     */
-    @Override
-    public GradientColors getColors(int which) {
-        return getColors(which, TYPE_DARK);
-    }
-
-    /**
-     * Wallpaper colors when the wallpaper is visible, fallback otherwise.
-     *
-     * @param which FLAG_LOCK or FLAG_SYSTEM
-     * @param type TYPE_NORMAL, TYPE_DARK or TYPE_EXTRA_DARK
-     * @return colors
-     */
-    @Override
-    public GradientColors getColors(int which, int type) {
-        return getColors(which, type, false /* ignoreVisibility */);
-    }
-
-    /**
-     * Get TYPE_NORMAL colors, possibly ignoring wallpaper visibility.
-     *
-     * @param which FLAG_LOCK or FLAG_SYSTEM
-     * @param ignoreWallpaperVisibility whether you want fallback colors or not if the wallpaper
-     *                                  isn't visible
-     * @return
-     */
-    public GradientColors getColors(int which, boolean ignoreWallpaperVisibility) {
-        return getColors(which, TYPE_NORMAL, ignoreWallpaperVisibility);
-    }
-
-    /**
-     *
-     * @param which FLAG_LOCK or FLAG_SYSTEM
-     * @param type TYPE_NORMAL, TYPE_DARK or TYPE_EXTRA_DARK
-     * @param ignoreWallpaperVisibility true if true wallpaper colors should be returning
-     *                                  if it's visible or not
-     * @return colors
-     */
-    public GradientColors getColors(int which, int type, boolean ignoreWallpaperVisibility) {
-        // mWallpaperVisible only handles the "system wallpaper" and will be always set to false
-        // if we have different lock and system wallpapers.
-        if (which == WallpaperManager.FLAG_SYSTEM) {
-            if (mWallpaperVisible || ignoreWallpaperVisibility) {
-                return super.getColors(which, type);
-            } else {
-                return mWpHiddenColors;
-            }
-        } else {
-            if (mHasBackdrop) {
-                return mWpHiddenColors;
-            } else {
-                return super.getColors(which, type);
-            }
-        }
-    }
-
-    @VisibleForTesting
-    void setWallpaperVisible(boolean visible) {
-        if (mWallpaperVisible != visible) {
-            mWallpaperVisible = visible;
-            triggerColorsChanged(WallpaperManager.FLAG_SYSTEM);
-        }
-    }
-
-    public void setHasBackdrop(boolean hasBackdrop) {
-        if (mHasBackdrop != hasBackdrop) {
-            mHasBackdrop = hasBackdrop;
+    public void setHasMediaArtwork(boolean hasBackdrop) {
+        if (mHasMediaArtwork != hasBackdrop) {
+            mHasMediaArtwork = hasBackdrop;
             triggerColorsChanged(WallpaperManager.FLAG_LOCK);
         }
     }
@@ -230,7 +143,8 @@ public class SysuiColorExtractor extends ColorExtractor implements Dumpable,
         pw.println("  Gradients:");
         pw.println("    system: " + Arrays.toString(system));
         pw.println("    lock: " + Arrays.toString(lock));
-        pw.println("  Default scrim: " + mWpHiddenColors);
+        pw.println("  Neutral colors: " + mNeutralColorsLock);
+        pw.println("  Has media backdrop: " + mHasMediaArtwork);
 
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationMediaManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationMediaManager.java
index 75ef18545fdf4..6c36ab95f9231 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationMediaManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationMediaManager.java
@@ -488,6 +488,7 @@ public class NotificationMediaManager implements Dumpable {
         if (bmp != null) {
             artworkDrawable = new BitmapDrawable(mBackdropBack.getResources(), bmp);
         }
+        boolean hasMediaArtwork = artworkDrawable != null;
         boolean allowWhenShade = false;
         if (ENABLE_LOCKSCREEN_WALLPAPER && artworkDrawable == null) {
             Bitmap lockWallpaper =
@@ -506,7 +507,7 @@ public class NotificationMediaManager implements Dumpable {
         boolean hideBecauseOccluded = shadeController != null && shadeController.isOccluded();
 
         final boolean hasArtwork = artworkDrawable != null;
-        mColorExtractor.setHasBackdrop(hasArtwork);
+        mColorExtractor.setHasMediaArtwork(hasMediaArtwork);
         if (mScrimController != null) {
             mScrimController.setHasBackdrop(hasArtwork);
         }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
index d219ac920eec7..0fbc9ffaa17e5 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
@@ -654,15 +654,7 @@ public class NotificationStackScrollLayout extends ViewGroup implements ScrollAd
     @Override
     @ShadeViewRefactor(RefactorComponent.SHADE_VIEW)
     public void onThemeChanged() {
-        int which;
-        if (mStatusBarState == StatusBarState.KEYGUARD
-                || mStatusBarState == StatusBarState.SHADE_LOCKED) {
-            which = WallpaperManager.FLAG_LOCK;
-        } else {
-            which = WallpaperManager.FLAG_SYSTEM;
-        }
-        final boolean useDarkText = mColorExtractor.getColors(which,
-                true /* ignoreVisibility */).supportsDarkText();
+        final boolean useDarkText = mColorExtractor.getNeutralColors().supportsDarkText();
         updateDecorViews(useDarkText);
 
         updateFooter();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
index 5fff054b3bc2e..80fbda0809ba3 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
@@ -3222,8 +3222,7 @@ public class StatusBar extends SystemUI implements DemoMode,
 
         // Lock wallpaper defines the color of the majority of the views, hence we'll use it
         // to set our default theme.
-        final boolean lockDarkText = mColorExtractor.getColors(WallpaperManager.FLAG_LOCK, true
-                /* ignoreVisibility */).supportsDarkText();
+        final boolean lockDarkText = mColorExtractor.getNeutralColors().supportsDarkText();
         final int themeResId = lockDarkText ? R.style.Theme_SystemUI_Light : R.style.Theme_SystemUI;
         if (mContext.getThemeResId() != themeResId) {
             mContext.setTheme(themeResId);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarWindowController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarWindowController.java
index 891bb35127ee1..c73ed60f161dd 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarWindowController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarWindowController.java
@@ -25,7 +25,6 @@ import static com.android.systemui.statusbar.NotificationRemoteInputManager.ENAB
 
 import android.app.ActivityManager;
 import android.app.IActivityManager;
-import android.app.WallpaperManager;
 import android.content.Context;
 import android.content.pm.ActivityInfo;
 import android.content.res.Resources;
@@ -556,17 +555,7 @@ public class StatusBarWindowController implements Callback, Dumpable, Configurat
             return;
         }
 
-        StatusBarStateController state = Dependency.get(StatusBarStateController.class);
-        int which;
-        if (state.getState() == StatusBarState.KEYGUARD
-                || state.getState() == StatusBarState.SHADE_LOCKED) {
-            which = WallpaperManager.FLAG_LOCK;
-        } else {
-            which = WallpaperManager.FLAG_SYSTEM;
-        }
-        final boolean useDarkText = mColorExtractor.getColors(which,
-                true /* ignoreVisibility */).supportsDarkText();
-
+        final boolean useDarkText = mColorExtractor.getNeutralColors().supportsDarkText();
         // Make sure we have the correct navbar/statusbar colors.
         setKeyguardDark(useDarkText);
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/colorextraction/SysuiColorExtractorTests.java b/packages/SystemUI/tests/src/com/android/systemui/colorextraction/SysuiColorExtractorTests.java
index 9c2c822571739..41747f4075469 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/colorextraction/SysuiColorExtractorTests.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/colorextraction/SysuiColorExtractorTests.java
@@ -40,6 +40,7 @@ import com.android.systemui.statusbar.policy.ConfigurationController;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 /**
@@ -57,6 +58,8 @@ public class SysuiColorExtractorTests extends SysuiTestCase {
             ColorExtractor.TYPE_DARK,
             ColorExtractor.TYPE_EXTRA_DARK};
 
+    @Mock
+    private WallpaperManager mWallpaperManager;
     private ColorExtractor.GradientColors mColors;
     private SysuiColorExtractor mColorExtractor;
 
@@ -72,32 +75,15 @@ public class SysuiColorExtractorTests extends SysuiTestCase {
                     outGradientColorsNormal.set(mColors);
                     outGradientColorsDark.set(mColors);
                     outGradientColorsExtraDark.set(mColors);
-                }, mock(ConfigurationController.class), false);
+                }, mock(ConfigurationController.class), mWallpaperManager, true /* immediately */);
     }
 
     @Test
-    public void getColors_usesGreyIfWallpaperNotVisible() {
-        simulateEvent(mColorExtractor);
-        mColorExtractor.setWallpaperVisible(false);
-
-        ColorExtractor.GradientColors fallbackColors = mColorExtractor.getNeutralColors();
-
-        for (int type : sTypes) {
-            assertEquals("Not using fallback!",
-                    mColorExtractor.getColors(WallpaperManager.FLAG_SYSTEM, type), fallbackColors);
-            assertNotEquals("Wallpaper visibility event should not affect lock wallpaper.",
-                    mColorExtractor.getColors(WallpaperManager.FLAG_LOCK, type), fallbackColors);
-        }
-    }
-
-    @Test
-    public void getColors_doesntUseFallbackIfVisible() {
+    public void getColors() {
         mColors.setMainColor(Color.RED);
         mColors.setSecondaryColor(Color.RED);
 
         simulateEvent(mColorExtractor);
-        mColorExtractor.setWallpaperVisible(true);
-
         for (int which : sWhich) {
             for (int type : sTypes) {
                 assertEquals("Not using extracted colors!",
@@ -109,8 +95,7 @@ public class SysuiColorExtractorTests extends SysuiTestCase {
     @Test
     public void getColors_fallbackWhenMediaIsVisible() {
         simulateEvent(mColorExtractor);
-        mColorExtractor.setWallpaperVisible(true);
-        mColorExtractor.setHasBackdrop(true);
+        mColorExtractor.setHasMediaArtwork(true);
 
         ColorExtractor.GradientColors fallbackColors = mColorExtractor.getNeutralColors();
 
@@ -127,7 +112,7 @@ public class SysuiColorExtractorTests extends SysuiTestCase {
         Tonal tonal = mock(Tonal.class);
         ConfigurationController configurationController = mock(ConfigurationController.class);
         SysuiColorExtractor sysuiColorExtractor = new SysuiColorExtractor(getContext(),
-                tonal, configurationController, false /* registerVisibility */);
+                tonal, configurationController, mWallpaperManager, true /* immediately */);
         verify(configurationController).addCallback(eq(sysuiColorExtractor));
 
         reset(tonal);
diff --git a/tests/Internal/src/com/android/internal/colorextraction/ColorExtractorTest.java b/tests/Internal/src/com/android/internal/colorextraction/ColorExtractorTest.java
index 17fa93135c7d8..45ddc3eed39c8 100644
--- a/tests/Internal/src/com/android/internal/colorextraction/ColorExtractorTest.java
+++ b/tests/Internal/src/com/android/internal/colorextraction/ColorExtractorTest.java
@@ -39,6 +39,8 @@ import com.android.internal.colorextraction.types.Tonal;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 
 /**
  * Tests color extraction generation.
@@ -48,16 +50,19 @@ import org.junit.runner.RunWith;
 public class ColorExtractorTest {
 
     Context mContext;
+    @Mock
+    WallpaperManager mWallpaperManager;
 
     @Before
     public void setup() {
+        MockitoAnnotations.initMocks(this);
         mContext = InstrumentationRegistry.getContext();
     }
 
     @Test
     public void ColorExtractor_extractWhenInitialized() {
         ExtractionType type = mock(Tonal.class);
-        new ColorExtractor(mContext, type, true);
+        new ColorExtractor(mContext, type, true, mWallpaperManager);
         // 1 for lock and 1 for system
         verify(type, times(2))
                 .extractInto(any(), any(), any(), any());
@@ -84,7 +89,7 @@ public class ColorExtractorTest {
                     outGradientColorsDark.set(colorsExpectedDark);
                     outGradientColorsExtraDark.set(colorsExpectedExtraDark);
                 };
-        ColorExtractor extractor = new ColorExtractor(mContext, type, true);
+        ColorExtractor extractor = new ColorExtractor(mContext, type, true, mWallpaperManager);
 
         GradientColors colors = extractor.getColors(WallpaperManager.FLAG_SYSTEM,
                 ColorExtractor.TYPE_NORMAL);
@@ -99,7 +104,8 @@ public class ColorExtractorTest {
     public void addOnColorsChangedListener_invokesListener() {
         ColorExtractor.OnColorsChangedListener mockedListeners =
                 mock(ColorExtractor.OnColorsChangedListener.class);
-        ColorExtractor extractor = new ColorExtractor(mContext, new Tonal(mContext), true);
+        ColorExtractor extractor = new ColorExtractor(mContext, new Tonal(mContext), true,
+                mWallpaperManager);
         extractor.addOnColorsChangedListener(mockedListeners);
 
         extractor.onColorsChanged(new WallpaperColors(Color.valueOf(Color.RED), null, null),