Fix flaky test

Test was registering a listener, and events could come from wallpaper
manager making it unstable.

Fixes: 134485306
Test: atest SysuiColorExtractorTests
Change-Id: Ie838cbee0887f3503ce328d5c5b5603c413d5458

4 files changed, 27 insertions, 26 deletions

diff --git a/core/java/com/android/internal/colorextraction/ColorExtractor.java b/core/java/com/android/internal/colorextraction/ColorExtractor.java
index d9fd3b5bd6d8..a6286c0f03f7 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;
 
@@ -71,14 +73,8 @@ 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 {
-            wallpaperManager.addOnColorsChangedListener(this, null /* handler */);
-            initExtractColors(wallpaperManager, immediately);
-        }
+        wallpaperManager.addOnColorsChangedListener(this, null /* handler */);
+        initExtractColors(wallpaperManager, immediately);
     }
 
     private void initExtractColors(WallpaperManager wallpaperManager, boolean immediately) {
diff --git a/packages/SystemUI/src/com/android/systemui/colorextraction/SysuiColorExtractor.java b/packages/SystemUI/src/com/android/systemui/colorextraction/SysuiColorExtractor.java
index 835ffc976e9f..6579c0b6de27 100644
--- a/packages/SystemUI/src/com/android/systemui/colorextraction/SysuiColorExtractor.java
+++ b/packages/SystemUI/src/com/android/systemui/colorextraction/SysuiColorExtractor.java
@@ -59,13 +59,15 @@ public class SysuiColorExtractor extends ColorExtractor implements Dumpable,
 
     @Inject
     public SysuiColorExtractor(Context context, ConfigurationController configurationController) {
-        this(context, new Tonal(context), configurationController, true);
+        this(context, new Tonal(context), configurationController, true,
+                context.getSystemService(WallpaperManager.class));
     }
 
     @VisibleForTesting
     public SysuiColorExtractor(Context context, ExtractionType type,
-            ConfigurationController configurationController, boolean registerVisibility) {
-        super(context, type, false /* immediately */);
+            ConfigurationController configurationController, boolean registerVisibility,
+            WallpaperManager wallpaperManager) {
+        super(context, type, false /* immediately */, wallpaperManager);
         mTonal = type instanceof Tonal ? (Tonal) type : new Tonal(context);
         mWpHiddenColors = new GradientColors();
         configurationController.addCallback(this);
@@ -91,13 +93,10 @@ public class SysuiColorExtractor extends ColorExtractor implements Dumpable,
             }
         }
 
-        WallpaperManager wallpaperManager = context.getSystemService(WallpaperManager.class);
-        if (wallpaperManager != null) {
-            // Listen to all users instead of only the current one.
-            wallpaperManager.removeOnColorsChangedListener(this);
-            wallpaperManager.addOnColorsChangedListener(this, null /* handler */,
-                    UserHandle.USER_ALL);
-        }
+        // Listen to all users instead of only the current one.
+        wallpaperManager.removeOnColorsChangedListener(this);
+        wallpaperManager.addOnColorsChangedListener(this, null /* handler */,
+                UserHandle.USER_ALL);
     }
 
     private void updateDefaultGradients(WallpaperColors colors) {
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 9c2c82257173..e3a162c7b10d 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,7 +75,7 @@ public class SysuiColorExtractorTests extends SysuiTestCase {
                     outGradientColorsNormal.set(mColors);
                     outGradientColorsDark.set(mColors);
                     outGradientColorsExtraDark.set(mColors);
-                }, mock(ConfigurationController.class), false);
+                }, mock(ConfigurationController.class), false, mWallpaperManager);
     }
 
     @Test
@@ -127,7 +130,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, false /* registerVisibility */, mWallpaperManager);
         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 17fa93135c7d..f88a7c41c02b 100644
--- a/tests/Internal/src/com/android/internal/colorextraction/ColorExtractorTest.java
+++ b/tests/Internal/src/com/android/internal/colorextraction/ColorExtractorTest.java
@@ -47,17 +47,19 @@ import org.junit.runner.RunWith;
 @RunWith(AndroidJUnit4.class)
 public class ColorExtractorTest {
 
-    Context mContext;
+    private Context mContext;
+    private WallpaperManager mWallpaperManager;
 
     @Before
     public void setup() {
         mContext = InstrumentationRegistry.getContext();
+        mWallpaperManager = mock(WallpaperManager.class);
     }
 
     @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 +86,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 +101,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),