Merge "Add an API to BrightnessConfiguration to enable colour sampling."
diff --git a/api/system-current.txt b/api/system-current.txt
index aed13b7..d91ecd4 100644
--- a/api/system-current.txt
+++ b/api/system-current.txt
@@ -2001,6 +2001,7 @@
     method @Nullable public android.hardware.display.BrightnessCorrection getCorrectionByCategory(int);
     method @Nullable public android.hardware.display.BrightnessCorrection getCorrectionByPackageName(@NonNull String);
     method public android.util.Pair<float[],float[]> getCurve();
+    method public boolean shouldCollectColorSamples();
     method public void writeToParcel(android.os.Parcel, int);
     field @NonNull public static final android.os.Parcelable.Creator<android.hardware.display.BrightnessConfiguration> CREATOR;
   }
@@ -2013,6 +2014,7 @@
     method public int getMaxCorrectionsByCategory();
     method public int getMaxCorrectionsByPackageName();
     method @NonNull public android.hardware.display.BrightnessConfiguration.Builder setDescription(@Nullable String);
+    method @NonNull public android.hardware.display.BrightnessConfiguration.Builder setShouldCollectColorSamples(boolean);
   }
 
   public final class BrightnessCorrection implements android.os.Parcelable {
diff --git a/api/test-current.txt b/api/test-current.txt
index 44f736c..c6b02b9 100644
--- a/api/test-current.txt
+++ b/api/test-current.txt
@@ -1018,6 +1018,7 @@
     method @Nullable public android.hardware.display.BrightnessCorrection getCorrectionByCategory(int);
     method @Nullable public android.hardware.display.BrightnessCorrection getCorrectionByPackageName(@NonNull String);
     method public android.util.Pair<float[],float[]> getCurve();
+    method public boolean shouldCollectColorSamples();
     method public void writeToParcel(android.os.Parcel, int);
     field @NonNull public static final android.os.Parcelable.Creator<android.hardware.display.BrightnessConfiguration> CREATOR;
   }
@@ -1030,6 +1031,7 @@
     method public int getMaxCorrectionsByCategory();
     method public int getMaxCorrectionsByPackageName();
     method @NonNull public android.hardware.display.BrightnessConfiguration.Builder setDescription(@Nullable String);
+    method @NonNull public android.hardware.display.BrightnessConfiguration.Builder setShouldCollectColorSamples(boolean);
   }
 
   public final class BrightnessCorrection implements android.os.Parcelable {
diff --git a/core/java/android/hardware/display/BrightnessChangeEvent.java b/core/java/android/hardware/display/BrightnessChangeEvent.java
index 21fcc63..a6a44be 100644
--- a/core/java/android/hardware/display/BrightnessChangeEvent.java
+++ b/core/java/android/hardware/display/BrightnessChangeEvent.java
@@ -79,7 +79,8 @@
     /**
      * Histogram counting how many times a pixel of a given value was displayed onscreen for the
      * Value component of HSV if the device supports color sampling, if the device does not support
-     * color sampling the value will be null.
+     * color sampling or {@link BrightnessConfiguration#shouldCollectColorSamples()} is false the
+     * value will be null.
      *
      * The buckets of the histogram are evenly weighted, the number of buckets is device specific.
      * The units are in pixels * milliseconds, with 1 pixel millisecond being 1 pixel displayed
@@ -94,7 +95,8 @@
 
     /**
      * How many milliseconds of data are contained in the colorValueBuckets, if the device does
-     * not support color sampling the value will be 0L.
+     * not support color sampling or {@link BrightnessConfiguration#shouldCollectColorSamples()} is
+     * false the value will be 0L.
      *
      * {@see #colorValueBuckets}
      */
diff --git a/core/java/android/hardware/display/BrightnessConfiguration.java b/core/java/android/hardware/display/BrightnessConfiguration.java
index 4c2e297..139be8e 100644
--- a/core/java/android/hardware/display/BrightnessConfiguration.java
+++ b/core/java/android/hardware/display/BrightnessConfiguration.java
@@ -49,26 +49,31 @@
     private static final String TAG_BRIGHTNESS_POINT = "brightness-point";
     private static final String TAG_BRIGHTNESS_CORRECTIONS = "brightness-corrections";
     private static final String TAG_BRIGHTNESS_CORRECTION = "brightness-correction";
+    private static final String TAG_BRIGHTNESS_PARAMS = "brightness-params";
     private static final String ATTR_LUX = "lux";
     private static final String ATTR_NITS = "nits";
     private static final String ATTR_DESCRIPTION = "description";
     private static final String ATTR_PACKAGE_NAME = "package-name";
     private static final String ATTR_CATEGORY = "category";
+    private static final String ATTR_COLLECT_COLOR = "collect-color";
 
     private final float[] mLux;
     private final float[] mNits;
     private final Map<String, BrightnessCorrection> mCorrectionsByPackageName;
     private final Map<Integer, BrightnessCorrection> mCorrectionsByCategory;
     private final String mDescription;
+    private final boolean mShouldCollectColorSamples;
 
     private BrightnessConfiguration(float[] lux, float[] nits,
             Map<String, BrightnessCorrection> correctionsByPackageName,
-            Map<Integer, BrightnessCorrection> correctionsByCategory, String description) {
+            Map<Integer, BrightnessCorrection> correctionsByCategory, String description,
+            boolean shouldCollectColorSamples) {
         mLux = lux;
         mNits = nits;
         mCorrectionsByPackageName = correctionsByPackageName;
         mCorrectionsByCategory = correctionsByCategory;
         mDescription = description;
+        mShouldCollectColorSamples = shouldCollectColorSamples;
     }
 
     /**
@@ -119,6 +124,14 @@
         return mDescription;
     }
 
+    /**
+     * Returns whether color samples should be collected in
+     * {@link BrightnessChangeEvent#colorValueBuckets}.
+     */
+    public boolean shouldCollectColorSamples() {
+        return mShouldCollectColorSamples;
+    }
+
     @Override
     public void writeToParcel(Parcel dest, int flags) {
         dest.writeFloatArray(mLux);
@@ -138,6 +151,7 @@
             correction.writeToParcel(dest, flags);
         }
         dest.writeString(mDescription);
+        dest.writeBoolean(mShouldCollectColorSamples);
     }
 
     @Override
@@ -167,6 +181,7 @@
         if (mDescription != null) {
             sb.append(mDescription);
         }
+        sb.append(", shouldCollectColorSamples = " + mShouldCollectColorSamples);
         sb.append("'}");
         return sb.toString();
     }
@@ -181,6 +196,7 @@
         if (mDescription != null) {
             result = result * 31 + mDescription.hashCode();
         }
+        result = result * 31 + Boolean.hashCode(mShouldCollectColorSamples);
         return result;
     }
 
@@ -196,7 +212,8 @@
         return Arrays.equals(mLux, other.mLux) && Arrays.equals(mNits, other.mNits)
                 && mCorrectionsByPackageName.equals(other.mCorrectionsByPackageName)
                 && mCorrectionsByCategory.equals(other.mCorrectionsByCategory)
-                && Objects.equals(mDescription, other.mDescription);
+                && Objects.equals(mDescription, other.mDescription)
+                && mShouldCollectColorSamples == other.mShouldCollectColorSamples;
     }
 
     public static final @android.annotation.NonNull Creator<BrightnessConfiguration> CREATOR =
@@ -224,6 +241,8 @@
 
             final String description = in.readString();
             builder.setDescription(description);
+            final boolean shouldCollectColorSamples = in.readBoolean();
+            builder.setShouldCollectColorSamples(shouldCollectColorSamples);
             return builder.build();
         }
 
@@ -252,6 +271,7 @@
             serializer.endTag(null, TAG_BRIGHTNESS_POINT);
         }
         serializer.endTag(null, TAG_BRIGHTNESS_CURVE);
+
         serializer.startTag(null, TAG_BRIGHTNESS_CORRECTIONS);
         for (Map.Entry<String, BrightnessCorrection> entry :
                 mCorrectionsByPackageName.entrySet()) {
@@ -271,6 +291,12 @@
             serializer.endTag(null, TAG_BRIGHTNESS_CORRECTION);
         }
         serializer.endTag(null, TAG_BRIGHTNESS_CORRECTIONS);
+
+        serializer.startTag(null, TAG_BRIGHTNESS_PARAMS);
+        if (mShouldCollectColorSamples) {
+            serializer.attribute(null, ATTR_COLLECT_COLOR, Boolean.toString(true));
+        }
+        serializer.endTag(null, TAG_BRIGHTNESS_PARAMS);
     }
 
     /**
@@ -293,6 +319,7 @@
         List<Float> nitsList = new ArrayList<>();
         Map<String, BrightnessCorrection> correctionsByPackageName = new HashMap<>();
         Map<Integer, BrightnessCorrection> correctionsByCategory = new HashMap<>();
+        boolean shouldCollectColorSamples = false;
         final int configDepth = parser.getDepth();
         while (XmlUtils.nextElementWithin(parser, configDepth)) {
             if (TAG_BRIGHTNESS_CURVE.equals(parser.getName())) {
@@ -307,8 +334,7 @@
                     luxList.add(lux);
                     nitsList.add(nits);
                 }
-            }
-            if (TAG_BRIGHTNESS_CORRECTIONS.equals(parser.getName())) {
+            } else if (TAG_BRIGHTNESS_CORRECTIONS.equals(parser.getName())) {
                 final int correctionsDepth = parser.getDepth();
                 while (XmlUtils.nextElementWithin(parser, correctionsDepth)) {
                     if (!TAG_BRIGHTNESS_CORRECTION.equals(parser.getName())) {
@@ -328,6 +354,9 @@
                         }
                     }
                 }
+            } else if (TAG_BRIGHTNESS_PARAMS.equals(parser.getName())) {
+                shouldCollectColorSamples =
+                        Boolean.parseBoolean(parser.getAttributeValue(null, ATTR_COLLECT_COLOR));
             }
         }
         final int n = luxList.size();
@@ -350,6 +379,7 @@
             final BrightnessCorrection correction = entry.getValue();
             builder.addCorrectionByCategory(category, correction);
         }
+        builder.setShouldCollectColorSamples(shouldCollectColorSamples);
         return builder.build();
     }
 
@@ -374,6 +404,7 @@
         private Map<String, BrightnessCorrection> mCorrectionsByPackageName;
         private Map<Integer, BrightnessCorrection> mCorrectionsByCategory;
         private String mDescription;
+        private boolean mShouldCollectColorSamples;
 
         /**
          * Constructs the builder with the control points for the brightness curve.
@@ -498,6 +529,19 @@
         }
 
         /**
+         * Control whether screen color samples should be returned in
+         * {@link BrightnessChangeEvent#colorValueBuckets} if supported by the device.
+         *
+         * @param shouldCollectColorSamples true if color samples should be collected.
+         * @return
+         */
+        @NonNull
+        public Builder setShouldCollectColorSamples(boolean shouldCollectColorSamples) {
+            mShouldCollectColorSamples = shouldCollectColorSamples;
+            return this;
+        }
+
+        /**
          * Builds the {@link BrightnessConfiguration}.
          */
         @NonNull
@@ -506,7 +550,7 @@
                 throw new IllegalStateException("A curve must be set!");
             }
             return new BrightnessConfiguration(mCurveLux, mCurveNits, mCorrectionsByPackageName,
-                    mCorrectionsByCategory, mDescription);
+                    mCorrectionsByCategory, mDescription, mShouldCollectColorSamples);
         }
 
         private static void checkMonotonic(float[] vals, boolean strictlyIncreasing, String name) {
diff --git a/core/tests/coretests/src/android/hardware/display/BrightnessConfigurationTest.java b/core/tests/coretests/src/android/hardware/display/BrightnessConfigurationTest.java
index 823fca5..85aa118 100644
--- a/core/tests/coretests/src/android/hardware/display/BrightnessConfigurationTest.java
+++ b/core/tests/coretests/src/android/hardware/display/BrightnessConfigurationTest.java
@@ -23,13 +23,23 @@
 
 import android.os.Parcel;
 import android.util.Pair;
+import android.util.Xml;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.internal.util.FastXmlSerializer;
+
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+import org.xmlpull.v1.XmlSerializer;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 
 @SmallTest
@@ -104,11 +114,15 @@
         });
     }
 
-
     @Test
     public void testParceledConfigIsEquivalent() {
         BrightnessConfiguration.Builder builder =
                 new BrightnessConfiguration.Builder(LUX_LEVELS, NITS_LEVELS);
+        builder.setShouldCollectColorSamples(true);
+        builder.addCorrectionByCategory(3,
+                BrightnessCorrection.createScaleAndTranslateLog(1.0f, 2.0f));
+        builder.addCorrectionByPackageName("a.package.name",
+                BrightnessCorrection.createScaleAndTranslateLog(1.0f, 2.0f));
         BrightnessConfiguration config = builder.build();
         Parcel p = Parcel.obtain();
         p.writeParcelable(config, 0 /*flags*/);
@@ -119,12 +133,49 @@
     }
 
     @Test
+    public void testWriteReadXml() throws IOException, XmlPullParserException {
+        BrightnessConfiguration.Builder builder =
+                new BrightnessConfiguration.Builder(LUX_LEVELS, NITS_LEVELS);
+        builder.setShouldCollectColorSamples(true);
+        builder.addCorrectionByCategory(3,
+                BrightnessCorrection.createScaleAndTranslateLog(1.0f, 2.0f));
+        builder.addCorrectionByPackageName("a.package.name",
+                BrightnessCorrection.createScaleAndTranslateLog(1.0f, 2.0f));
+        BrightnessConfiguration config = builder.build();
+
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        XmlSerializer out = new FastXmlSerializer();
+        out.setOutput(baos, StandardCharsets.UTF_8.name());
+        out.startDocument(null, true);
+        config.saveToXml(out);
+        out.endDocument();
+        baos.flush();
+
+        ByteArrayInputStream input = new ByteArrayInputStream(baos.toByteArray());
+        XmlPullParser parser = Xml.newPullParser();
+        parser.setInput(input, StandardCharsets.UTF_8.name());
+        BrightnessConfiguration loadedConfig = BrightnessConfiguration.loadFromXml(parser);
+
+        assertEquals(config, loadedConfig);
+    }
+
+    @Test
     public void testEquals() {
         BrightnessConfiguration.Builder builder =
                 new BrightnessConfiguration.Builder(LUX_LEVELS, NITS_LEVELS);
+        builder.setShouldCollectColorSamples(true);
+        builder.addCorrectionByCategory(3,
+                BrightnessCorrection.createScaleAndTranslateLog(1.0f, 2.0f));
+        builder.addCorrectionByPackageName("a.package.name",
+                BrightnessCorrection.createScaleAndTranslateLog(1.0f, 2.0f));
         BrightnessConfiguration baseConfig = builder.build();
 
         builder = new BrightnessConfiguration.Builder(LUX_LEVELS, NITS_LEVELS);
+        builder.setShouldCollectColorSamples(true);
+        builder.addCorrectionByCategory(3,
+                BrightnessCorrection.createScaleAndTranslateLog(1.0f, 2.0f));
+        builder.addCorrectionByPackageName("a.package.name",
+                BrightnessCorrection.createScaleAndTranslateLog(1.0f, 2.0f));
         BrightnessConfiguration identicalConfig = builder.build();
         assertEquals(baseConfig, identicalConfig);
         assertEquals("hashCodes must be equal for identical configs",
@@ -133,14 +184,37 @@
         float[] lux = Arrays.copyOf(LUX_LEVELS, LUX_LEVELS.length);
         lux[lux.length - 1] = lux[lux.length - 1] * 2;
         builder = new BrightnessConfiguration.Builder(lux, NITS_LEVELS);
+        builder.setShouldCollectColorSamples(true);
+        builder.addCorrectionByCategory(3,
+                BrightnessCorrection.createScaleAndTranslateLog(1.0f, 2.0f));
+        builder.addCorrectionByPackageName("a.package.name",
+                BrightnessCorrection.createScaleAndTranslateLog(1.0f, 2.0f));
         BrightnessConfiguration luxDifferConfig = builder.build();
         assertNotEquals(baseConfig, luxDifferConfig);
 
         float[] nits = Arrays.copyOf(NITS_LEVELS, NITS_LEVELS.length);
         nits[nits.length - 1] = nits[nits.length - 1] * 2;
         builder = new BrightnessConfiguration.Builder(LUX_LEVELS, nits);
+        builder.setShouldCollectColorSamples(true);
+        builder.addCorrectionByCategory(3,
+                BrightnessCorrection.createScaleAndTranslateLog(1.0f, 2.0f));
+        builder.addCorrectionByPackageName("a.package.name",
+                BrightnessCorrection.createScaleAndTranslateLog(1.0f, 2.0f));
         BrightnessConfiguration nitsDifferConfig = builder.build();
         assertNotEquals(baseConfig, nitsDifferConfig);
+
+        builder = new BrightnessConfiguration.Builder(LUX_LEVELS, NITS_LEVELS);
+        builder.addCorrectionByCategory(3,
+                BrightnessCorrection.createScaleAndTranslateLog(1.0f, 2.0f));
+        builder.addCorrectionByPackageName("a.package.name",
+                BrightnessCorrection.createScaleAndTranslateLog(1.0f, 2.0f));
+        BrightnessConfiguration colorCollectionDiffers = builder.build();
+        assertNotEquals(baseConfig, colorCollectionDiffers);
+
+        builder = new BrightnessConfiguration.Builder(LUX_LEVELS, NITS_LEVELS);
+        builder.setShouldCollectColorSamples(true);
+        BrightnessConfiguration correctionsDiffer = builder.build();
+        assertNotEquals(baseConfig, correctionsDiffer);
     }
 
     private static void assertArrayEquals(float[] expected, float[] actual, String name) {
diff --git a/services/core/java/com/android/server/display/BrightnessTracker.java b/services/core/java/com/android/server/display/BrightnessTracker.java
index 126beef..eea1980 100644
--- a/services/core/java/com/android/server/display/BrightnessTracker.java
+++ b/services/core/java/com/android/server/display/BrightnessTracker.java
@@ -34,6 +34,7 @@
 import android.hardware.SensorManager;
 import android.hardware.display.AmbientBrightnessDayStats;
 import android.hardware.display.BrightnessChangeEvent;
+import android.hardware.display.BrightnessConfiguration;
 import android.hardware.display.ColorDisplayManager;
 import android.hardware.display.DisplayManager;
 import android.hardware.display.DisplayManagerInternal;
@@ -125,6 +126,7 @@
     private static final int MSG_BRIGHTNESS_CHANGED = 1;
     private static final int MSG_STOP_SENSOR_LISTENER = 2;
     private static final int MSG_START_SENSOR_LISTENER = 3;
+    private static final int MSG_BRIGHTNESS_CONFIG_CHANGED = 4;
 
     private static final SimpleDateFormat FORMAT = new SimpleDateFormat("MM-dd HH:mm:ss.SSS");
 
@@ -158,6 +160,7 @@
     private boolean mColorSamplingEnabled;
     private int mNoFramesToSample;
     private float mFrameRate;
+    private BrightnessConfiguration mBrightnessConfiguration;
     // End of block of members that should only be accessed on the mBgHandler thread.
 
     private @UserIdInt int mCurrentUserId = UserHandle.USER_NULL;
@@ -202,6 +205,14 @@
         mBgHandler.obtainMessage(MSG_BACKGROUND_START, (Float) initialBrightness).sendToTarget();
     }
 
+    /**
+     * Update tracker with new brightness configuration.
+     */
+    public void setBrightnessConfiguration(BrightnessConfiguration brightnessConfiguration) {
+        mBgHandler.obtainMessage(MSG_BRIGHTNESS_CONFIG_CHANGED,
+                brightnessConfiguration).sendToTarget();
+    }
+
     private void backgroundStart(float initialBrightness) {
         readEvents();
         readAmbientBrightnessStats();
@@ -759,7 +770,9 @@
     private void enableColorSampling() {
         if (!mInjector.isBrightnessModeAutomatic(mContentResolver)
                 || !mInjector.isInteractive(mContext)
-                || mColorSamplingEnabled) {
+                || mColorSamplingEnabled
+                || mBrightnessConfiguration == null
+                || !mBrightnessConfiguration.shouldCollectColorSamples()) {
             return;
         }
 
@@ -977,6 +990,18 @@
                     stopSensorListener();
                     disableColorSampling();
                     break;
+                case MSG_BRIGHTNESS_CONFIG_CHANGED:
+                    mBrightnessConfiguration = (BrightnessConfiguration) msg.obj;
+                    boolean shouldCollectColorSamples =
+                            mBrightnessConfiguration != null
+                                    && mBrightnessConfiguration.shouldCollectColorSamples();
+                    if (shouldCollectColorSamples && !mColorSamplingEnabled) {
+                        enableColorSampling();
+                    } else if (!shouldCollectColorSamples && mColorSamplingEnabled) {
+                        disableColorSampling();
+                    }
+                    break;
+
             }
         }
     }
diff --git a/services/core/java/com/android/server/display/DisplayPowerController.java b/services/core/java/com/android/server/display/DisplayPowerController.java
index 5804fc8..e42545e 100644
--- a/services/core/java/com/android/server/display/DisplayPowerController.java
+++ b/services/core/java/com/android/server/display/DisplayPowerController.java
@@ -931,6 +931,10 @@
                     autoBrightnessAdjustmentChanged, mPowerRequest.policy);
         }
 
+        if (mBrightnessTracker != null) {
+            mBrightnessTracker.setBrightnessConfiguration(mBrightnessConfiguration);
+        }
+
         // Apply auto-brightness.
         boolean slowChange = false;
         if (brightness < 0) {
diff --git a/services/tests/servicestests/src/com/android/server/display/BrightnessTrackerTest.java b/services/tests/servicestests/src/com/android/server/display/BrightnessTrackerTest.java
index 7081d2e..1d04c83 100644
--- a/services/tests/servicestests/src/com/android/server/display/BrightnessTrackerTest.java
+++ b/services/tests/servicestests/src/com/android/server/display/BrightnessTrackerTest.java
@@ -37,6 +37,7 @@
 import android.hardware.SensorEventListener;
 import android.hardware.display.AmbientBrightnessDayStats;
 import android.hardware.display.BrightnessChangeEvent;
+import android.hardware.display.BrightnessConfiguration;
 import android.hardware.display.DisplayManager;
 import android.hardware.display.DisplayedContentSample;
 import android.hardware.display.DisplayedContentSamplingAttributes;
@@ -78,6 +79,7 @@
 @RunWith(AndroidJUnit4.class)
 public class BrightnessTrackerTest {
     private static final float DEFAULT_INITIAL_BRIGHTNESS = 2.5f;
+    private static final boolean DEFAULT_COLOR_SAMPLING_ENABLED = true;
     private static final float FLOAT_DELTA = 0.01f;
 
     private BrightnessTracker mTracker;
@@ -151,6 +153,40 @@
     }
 
     @Test
+    public void testModifyBrightnessConfiguration() {
+        mInjector.mInteractive = true;
+        // Start with tracker not listening for color samples.
+        startTracker(mTracker, DEFAULT_INITIAL_BRIGHTNESS, /* collectColorSamples= */ false);
+        assertFalse(mInjector.mColorSamplingEnabled);
+
+        // Update brightness config to enabled color sampling.
+        mTracker.setBrightnessConfiguration(buildBrightnessConfiguration(
+                /* collectColorSamples= */ true));
+        mInjector.waitForHandler();
+        assertTrue(mInjector.mColorSamplingEnabled);
+
+        // Update brightness config to disable color sampling.
+        mTracker.setBrightnessConfiguration(buildBrightnessConfiguration(
+                /* collectColorSamples= */ false));
+        mInjector.waitForHandler();
+        assertFalse(mInjector.mColorSamplingEnabled);
+
+        // Pretend screen is off, update config to turn on color sampling.
+        mInjector.sendScreenChange(/*screen on */ false);
+        mTracker.setBrightnessConfiguration(buildBrightnessConfiguration(
+                /* collectColorSamples= */ true));
+        mInjector.waitForHandler();
+        assertFalse(mInjector.mColorSamplingEnabled);
+
+        // Pretend screen is on.
+        mInjector.sendScreenChange(/*screen on */ true);
+        assertTrue(mInjector.mColorSamplingEnabled);
+
+        mTracker.stop();
+        assertFalse(mInjector.mColorSamplingEnabled);
+    }
+
+    @Test
     public void testNoColorSampling_WrongPixelFormat() {
         mInjector.mDefaultSamplingAttributes =
                 new DisplayedContentSamplingAttributes(
@@ -278,7 +314,7 @@
         mInjector.mSecureIntSettings.put(Settings.Secure.NIGHT_DISPLAY_ACTIVATED, 1);
         mInjector.mSecureIntSettings.put(Settings.Secure.NIGHT_DISPLAY_COLOR_TEMPERATURE, 3333);
 
-        startTracker(mTracker, initialBrightness);
+        startTracker(mTracker, initialBrightness, DEFAULT_COLOR_SAMPLING_ENABLED);
         mInjector.mBroadcastReceiver.onReceive(InstrumentationRegistry.getContext(),
                 batteryChangeEvent(30, 60));
         mInjector.mSensorListener.onSensorChanged(createSensorEvent(1000.0f));
@@ -311,7 +347,7 @@
     @Test
     public void testIgnoreAutomaticBrightnessChange() {
         final int initialBrightness = 30;
-        startTracker(mTracker, initialBrightness);
+        startTracker(mTracker, initialBrightness, DEFAULT_COLOR_SAMPLING_ENABLED);
         mInjector.mSensorListener.onSensorChanged(createSensorEvent(1.0f));
         mInjector.incrementTime(TimeUnit.SECONDS.toMillis(1));
 
@@ -750,11 +786,13 @@
     }
 
     private void startTracker(BrightnessTracker tracker) {
-        startTracker(tracker, DEFAULT_INITIAL_BRIGHTNESS);
+        startTracker(tracker, DEFAULT_INITIAL_BRIGHTNESS,  DEFAULT_COLOR_SAMPLING_ENABLED);
     }
 
-    private void startTracker(BrightnessTracker tracker, float initialBrightness) {
+    private void startTracker(BrightnessTracker tracker, float initialBrightness,
+            boolean collectColorSamples) {
         tracker.start(initialBrightness);
+        tracker.setBrightnessConfiguration(buildBrightnessConfiguration(collectColorSamples));
         mInjector.waitForHandler();
     }
 
@@ -772,6 +810,14 @@
         mInjector.waitForHandler();
     }
 
+    private BrightnessConfiguration buildBrightnessConfiguration(boolean collectColorSamples) {
+        BrightnessConfiguration.Builder builder = new BrightnessConfiguration.Builder(
+                /* lux = */ new float[] {0f, 10f, 100f},
+                /* nits = */ new float[] {1f, 90f, 100f});
+        builder.setShouldCollectColorSamples(collectColorSamples);
+        return builder.build();
+    }
+
     private static final class Idle implements MessageQueue.IdleHandler {
         private boolean mIdle;