Merge "Fix unreadable clock" into qt-dev
am: c442c8e5f5

Change-Id: Ifc3d184fbd856df515bfa527028c8e736572e14e
diff --git a/core/java/com/android/internal/colorextraction/ColorExtractor.java b/core/java/com/android/internal/colorextraction/ColorExtractor.java
index d9fd3b5..b27c11b 100644
--- a/core/java/com/android/internal/colorextraction/ColorExtractor.java
+++ b/core/java/com/android/internal/colorextraction/ColorExtractor.java
@@ -53,11 +53,13 @@
     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 @@
 
         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 @@
         }
     }
 
-    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 79c691c..a7ccc3a 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 @@
         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 @@
     @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 9a04222..67c4458 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 0bb9e744..7ec1bda 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardClockSwitch.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardClockSwitch.java
@@ -374,7 +374,7 @@
 
     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 f468eca..558ac4b 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 @@
         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 61a6952..bdf9dc4 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 @@
         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 ce1f09c..98679ade 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 @@
         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 835ffc9..6f56a53 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 @@
         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 @@
         }
     }
 
-    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 @@
             // 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 @@
      * - 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 @@
         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 75ef185..6c36ab9 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 @@
         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 @@
         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 2bd09c89..95eaf4f 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 @@
     @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 5fff054..80fbda0 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 @@
 
         // 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 891bb35..c73ed60 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 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 @@
             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 9c2c822..41747f4 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 org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 /**
@@ -57,6 +58,8 @@
             ColorExtractor.TYPE_DARK,
             ColorExtractor.TYPE_EXTRA_DARK};
 
+    @Mock
+    private WallpaperManager mWallpaperManager;
     private ColorExtractor.GradientColors mColors;
     private SysuiColorExtractor mColorExtractor;
 
@@ -72,32 +75,15 @@
                     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 @@
     @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 @@
         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 17fa931..45ddc3e 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 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 @@
 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 @@
                     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 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),