AOD: Refactor doze brigthness controllers, add AOD auto brightness support

- Factor out doze brightness and screen state controllers
- Add support for setting the screen brightness from a
  sensor in AOD / when pulsing
- Refactor FakeSensorManager for generic sensors, add a light sensor

Bug: 38354633
Test: runtest -x packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenTest.java
Change-Id: I55f2bee253cee23a76ba857ba1f19052ff88a753
diff --git a/packages/SystemUI/res/values/config.xml b/packages/SystemUI/res/values/config.xml
index 2ad6f2d..91d6b24 100644
--- a/packages/SystemUI/res/values/config.xml
+++ b/packages/SystemUI/res/values/config.xml
@@ -247,6 +247,11 @@
     -->
     <string name="doze_pickup_subtype_performs_proximity_check"></string>
 
+    <!-- Type of a sensor that provides a low-power estimate of the desired display
+         brightness, suitable to listen to while the device is asleep (e.g. during
+         always-on display) -->
+    <string name="doze_brightness_sensor_type" translatable="false"></string>
+
     <!-- Doze: pulse parameter - how long does it take to fade in? -->
     <integer name="doze_pulse_duration_in">900</integer>
 
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeFactory.java b/packages/SystemUI/src/com/android/systemui/doze/DozeFactory.java
index 1cc10c2..53a8be5 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeFactory.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeFactory.java
@@ -19,10 +19,12 @@
 import android.app.AlarmManager;
 import android.app.Application;
 import android.content.Context;
+import android.hardware.Sensor;
 import android.hardware.SensorManager;
 import android.os.Handler;
 
 import com.android.internal.hardware.AmbientDisplayConfiguration;
+import com.android.systemui.R;
 import com.android.systemui.SystemUIApplication;
 import com.android.systemui.statusbar.phone.DozeParameters;
 import com.android.systemui.util.wakelock.DelayedWakeLock;
@@ -46,20 +48,31 @@
         WakeLock wakeLock = new DelayedWakeLock(handler,
                 WakeLock.createPartial(context, "Doze"));
 
-        DozeMachine machine = new DozeMachine(
-                DozeSuspendScreenStatePreventingAdapter.wrapIfNeeded(
-                        DozeScreenStatePreventingAdapter.wrapIfNeeded(dozeService, params), params),
-                config,
-                wakeLock);
+        DozeMachine.Service wrappedService = DozeSuspendScreenStatePreventingAdapter.wrapIfNeeded(
+                DozeScreenStatePreventingAdapter.wrapIfNeeded(dozeService, params), params);
+        DozeMachine machine = new DozeMachine(wrappedService, config, wakeLock);
         machine.setParts(new DozeMachine.Part[]{
                 createDozeTriggers(context, sensorManager, host, config, params, handler, wakeLock,
                         machine),
                 createDozeUi(context, host, wakeLock, machine, handler, alarmManager),
+                createDozeScreenState(wrappedService),
+                createDozeScreenBrightness(context, wrappedService, sensorManager, handler),
         });
 
         return machine;
     }
 
+    private DozeMachine.Part createDozeScreenState(DozeMachine.Service service) {
+        return new DozeScreenState(service);
+    }
+
+    private DozeMachine.Part createDozeScreenBrightness(Context context,
+            DozeMachine.Service service, SensorManager sensorManager, Handler handler) {
+        Sensor sensor = DozeSensors.findSensorWithType(sensorManager,
+                context.getString(R.string.doze_brightness_sensor_type));
+        return new DozeScreenBrightness(context, service, sensorManager, sensor, handler);
+    }
+
     private DozeTriggers createDozeTriggers(Context context, SensorManager sensorManager,
             DozeHost host, AmbientDisplayConfiguration config, DozeParameters params,
             Handler handler, WakeLock wakeLock, DozeMachine machine) {
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java b/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java
index 5526e6b..e72abf3 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java
@@ -222,7 +222,6 @@
 
         updatePulseReason(newState, oldState, pulseReason);
         performTransitionOnComponents(oldState, newState);
-        updateScreenState(newState);
         updateWakeLockState(newState);
 
         resolveIntermediateState(newState);
@@ -307,13 +306,6 @@
         }
     }
 
-    private void updateScreenState(State newState) {
-        int state = newState.screenState();
-        if (state != Display.STATE_UNKNOWN) {
-            mDozeService.setDozeScreenState(state);
-        }
-    }
-
     private void resolveIntermediateState(State state) {
         switch (state) {
             case INITIALIZED:
@@ -360,5 +352,36 @@
 
         /** Request waking up. */
         void requestWakeUp();
+
+        /** Set screen brightness */
+        void setDozeScreenBrightness(int brightness);
+
+        class Delegate implements Service {
+            private final Service mDelegate;
+
+            public Delegate(Service delegate) {
+                mDelegate = delegate;
+            }
+
+            @Override
+            public void finish() {
+                mDelegate.finish();
+            }
+
+            @Override
+            public void setDozeScreenState(int state) {
+                mDelegate.setDozeScreenState(state);
+            }
+
+            @Override
+            public void requestWakeUp() {
+                mDelegate.requestWakeUp();
+            }
+
+            @Override
+            public void setDozeScreenBrightness(int brightness) {
+                mDelegate.setDozeScreenBrightness(brightness);
+            }
+        }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeScreenBrightness.java b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenBrightness.java
new file mode 100644
index 0000000..e461986
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenBrightness.java
@@ -0,0 +1,92 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.doze;
+
+import android.content.Context;
+import android.hardware.Sensor;
+import android.hardware.SensorEvent;
+import android.hardware.SensorEventListener;
+import android.hardware.SensorManager;
+import android.os.Handler;
+
+/**
+ * Controls the screen brightness when dozing.
+ */
+public class DozeScreenBrightness implements DozeMachine.Part, SensorEventListener {
+    private final Context mContext;
+    private final DozeMachine.Service mDozeService;
+    private final Handler mHandler;
+    private final SensorManager mSensorManager;
+    private final Sensor mLightSensor;
+    private boolean mRegistered;
+
+    public DozeScreenBrightness(Context context, DozeMachine.Service service,
+            SensorManager sensorManager, Sensor lightSensor, Handler handler) {
+        mContext = context;
+        mDozeService = service;
+        mSensorManager = sensorManager;
+        mLightSensor = lightSensor;
+        mHandler = handler;
+    }
+
+    @Override
+    public void transitionTo(DozeMachine.State oldState, DozeMachine.State newState) {
+        switch (newState) {
+            case INITIALIZED:
+                resetBrightnessToDefault();
+                break;
+            case DOZE_AOD:
+            case DOZE_REQUEST_PULSE:
+                setLightSensorEnabled(true);
+                break;
+            case DOZE:
+            case DOZE_AOD_PAUSED:
+                setLightSensorEnabled(false);
+                resetBrightnessToDefault();
+                break;
+            case FINISH:
+                setLightSensorEnabled(false);
+                break;
+        }
+    }
+
+    @Override
+    public void onSensorChanged(SensorEvent event) {
+        if (mRegistered) {
+            mDozeService.setDozeScreenBrightness(Math.max(1, (int) event.values[0]));
+        }
+    }
+
+    @Override
+    public void onAccuracyChanged(Sensor sensor, int accuracy) {
+    }
+
+    private void resetBrightnessToDefault() {
+        mDozeService.setDozeScreenBrightness(mContext.getResources().getInteger(
+                com.android.internal.R.integer.config_screenBrightnessDoze));
+    }
+
+    private void setLightSensorEnabled(boolean enabled) {
+        if (enabled && !mRegistered && mLightSensor != null) {
+            mRegistered = mSensorManager.registerListener(this, mLightSensor,
+                    SensorManager.SENSOR_DELAY_NORMAL, mHandler);
+        } else if (!enabled && mRegistered) {
+            mSensorManager.unregisterListener(this);
+            mRegistered = false;
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java
new file mode 100644
index 0000000..846ec27
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License
+ */
+
+package com.android.systemui.doze;
+
+import android.content.Context;
+import android.view.Display;
+
+/**
+ * Controls the screen when dozing.
+ */
+public class DozeScreenState implements DozeMachine.Part {
+    private final DozeMachine.Service mDozeService;
+
+    public DozeScreenState(DozeMachine.Service service) {
+        mDozeService = service;
+    }
+
+    @Override
+    public void transitionTo(DozeMachine.State oldState, DozeMachine.State newState) {
+        int screenState = newState.screenState();
+        if (screenState != Display.STATE_UNKNOWN) {
+            mDozeService.setDozeScreenState(screenState);
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeScreenStatePreventingAdapter.java b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenStatePreventingAdapter.java
index ad5897a..5d0a9d7 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeScreenStatePreventingAdapter.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenStatePreventingAdapter.java
@@ -24,18 +24,11 @@
 /**
  * Prevents usage of doze screen states on devices that don't support them.
  */
-public class DozeScreenStatePreventingAdapter implements DozeMachine.Service {
-
-    private final DozeMachine.Service mInner;
+public class DozeScreenStatePreventingAdapter extends DozeMachine.Service.Delegate {
 
     @VisibleForTesting
     DozeScreenStatePreventingAdapter(DozeMachine.Service inner) {
-        mInner = inner;
-    }
-
-    @Override
-    public void finish() {
-        mInner.finish();
+        super(inner);
     }
 
     @Override
@@ -43,12 +36,7 @@
         if (state == Display.STATE_DOZE || state == Display.STATE_DOZE_SUSPEND) {
             state = Display.STATE_ON;
         }
-        mInner.setDozeScreenState(state);
-    }
-
-    @Override
-    public void requestWakeUp() {
-        mInner.requestWakeUp();
+        super.setDozeScreenState(state);
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeSensors.java b/packages/SystemUI/src/com/android/systemui/doze/DozeSensors.java
index 23da716..b0b77c6 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeSensors.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeSensors.java
@@ -100,10 +100,14 @@
     }
 
     private Sensor findSensorWithType(String type) {
+        return findSensorWithType(mSensorManager, type);
+    }
+
+    static Sensor findSensorWithType(SensorManager sensorManager, String type) {
         if (TextUtils.isEmpty(type)) {
             return null;
         }
-        List<Sensor> sensorList = mSensorManager.getSensorList(Sensor.TYPE_ALL);
+        List<Sensor> sensorList = sensorManager.getSensorList(Sensor.TYPE_ALL);
         for (Sensor s : sensorList) {
             if (type.equals(s.getStringType())) {
                 return s;
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeService.java b/packages/SystemUI/src/com/android/systemui/doze/DozeService.java
index d9fb087..98b1106 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeService.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeService.java
@@ -76,8 +76,6 @@
         super.onDreamingStarted();
         mDozeMachine.requestState(DozeMachine.State.INITIALIZED);
         startDozing();
-        setDozeScreenBrightness(getResources().getInteger(
-                com.android.internal.R.integer.config_screenBrightnessDoze));
         if (mDozePlugin != null) {
             mDozePlugin.onDreamingStarted();
         }
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeSuspendScreenStatePreventingAdapter.java b/packages/SystemUI/src/com/android/systemui/doze/DozeSuspendScreenStatePreventingAdapter.java
index 1e06797..1c6521f 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeSuspendScreenStatePreventingAdapter.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeSuspendScreenStatePreventingAdapter.java
@@ -24,18 +24,11 @@
 /**
  * Prevents usage of doze screen states on devices that don't support them.
  */
-public class DozeSuspendScreenStatePreventingAdapter implements DozeMachine.Service {
-
-    private final DozeMachine.Service mInner;
+public class DozeSuspendScreenStatePreventingAdapter extends DozeMachine.Service.Delegate {
 
     @VisibleForTesting
     DozeSuspendScreenStatePreventingAdapter(DozeMachine.Service inner) {
-        mInner = inner;
-    }
-
-    @Override
-    public void finish() {
-        mInner.finish();
+        super(inner);
     }
 
     @Override
@@ -43,12 +36,7 @@
         if (state == Display.STATE_DOZE_SUSPEND) {
             state = Display.STATE_DOZE;
         }
-        mInner.setDozeScreenState(state);
-    }
-
-    @Override
-    public void requestWakeUp() {
-        mInner.requestWakeUp();
+        super.setDozeScreenState(state);
     }
 
     /**
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java
index 3a09ce0..368c814f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java
@@ -215,54 +215,6 @@
     }
 
     @Test
-    public void testScreen_offInDoze() {
-        mMachine.requestState(INITIALIZED);
-
-        mMachine.requestState(DOZE);
-
-        assertEquals(Display.STATE_OFF, mServiceFake.screenState);
-    }
-
-    @Test
-    public void testScreen_onInAod() {
-        mMachine.requestState(INITIALIZED);
-
-        mMachine.requestState(DOZE_AOD);
-
-        assertEquals(Display.STATE_DOZE_SUSPEND, mServiceFake.screenState);
-    }
-
-    @Test
-    public void testScreen_onInPulse() {
-        mMachine.requestState(INITIALIZED);
-
-        mMachine.requestPulse(DozeLog.PULSE_REASON_NOTIFICATION);
-        mMachine.requestState(DOZE_PULSING);
-
-        assertEquals(Display.STATE_ON, mServiceFake.screenState);
-    }
-
-    @Test
-    public void testScreen_offInRequestPulseWithoutAoD() {
-        mMachine.requestState(INITIALIZED);
-
-        mMachine.requestState(DOZE);
-        mMachine.requestPulse(DozeLog.PULSE_REASON_NOTIFICATION);
-
-        assertEquals(Display.STATE_OFF, mServiceFake.screenState);
-    }
-
-    @Test
-    public void testScreen_onInRequestPulseWithoutAoD() {
-        mMachine.requestState(INITIALIZED);
-
-        mMachine.requestState(DOZE_AOD);
-        mMachine.requestPulse(DozeLog.PULSE_REASON_NOTIFICATION);
-
-        assertEquals(Display.STATE_DOZE_SUSPEND, mServiceFake.screenState);
-    }
-
-    @Test
     public void testTransitions_canRequestTransitions() {
         mMachine.requestState(INITIALIZED);
         mMachine.requestState(DOZE);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenBrightnessTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenBrightnessTest.java
new file mode 100644
index 0000000..514dc8b
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenBrightnessTest.java
@@ -0,0 +1,162 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.doze;
+
+import static com.android.systemui.doze.DozeMachine.State.DOZE;
+import static com.android.systemui.doze.DozeMachine.State.DOZE_AOD;
+import static com.android.systemui.doze.DozeMachine.State.DOZE_AOD_PAUSED;
+import static com.android.systemui.doze.DozeMachine.State.DOZE_PULSE_DONE;
+import static com.android.systemui.doze.DozeMachine.State.DOZE_PULSING;
+import static com.android.systemui.doze.DozeMachine.State.DOZE_REQUEST_PULSE;
+import static com.android.systemui.doze.DozeMachine.State.FINISH;
+import static com.android.systemui.doze.DozeMachine.State.INITIALIZED;
+import static com.android.systemui.doze.DozeMachine.State.UNINITIALIZED;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertTrue;
+
+import android.os.PowerManager;
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
+
+import com.android.systemui.SysuiTestCase;
+import com.android.systemui.utils.hardware.FakeSensorManager;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class DozeScreenBrightnessTest extends SysuiTestCase {
+
+    DozeServiceFake mServiceFake;
+    DozeScreenBrightness mScreen;
+    FakeSensorManager.FakeGenericSensor mSensor;
+    FakeSensorManager mSensorManager;
+
+    @Before
+    public void setUp() throws Exception {
+        mServiceFake = new DozeServiceFake();
+        mSensorManager = new FakeSensorManager(mContext);
+        mSensor = mSensorManager.getFakeLightSensor();
+        mScreen = new DozeScreenBrightness(mContext, mServiceFake, mSensorManager,
+                mSensor.getSensor(), null /* handler */);
+    }
+
+    @Test
+    public void testInitialize_setsScreenBrightnessToValidValue() throws Exception {
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+
+        assertNotEquals(PowerManager.BRIGHTNESS_DEFAULT, mServiceFake.screenBrightness);
+        assertTrue(mServiceFake.screenBrightness <= PowerManager.BRIGHTNESS_ON);
+    }
+
+    @Test
+    public void testAod_usesLightSensor() throws Exception {
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE_AOD);
+
+        mSensor.sendSensorEvent(1000);
+
+        assertEquals(1000, mServiceFake.screenBrightness);
+    }
+
+    @Test
+    public void testPausingAod_pausesLightSensor() throws Exception {
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE_AOD);
+
+        mSensor.sendSensorEvent(1000);
+
+        mScreen.transitionTo(DOZE_AOD, DOZE_AOD_PAUSED);
+
+        mSensor.sendSensorEvent(1001);
+
+        assertNotEquals(1001, mServiceFake.screenBrightness);
+    }
+
+    @Test
+    public void testPausingAod_resetsBrightness() throws Exception {
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE_AOD);
+
+        mSensor.sendSensorEvent(1000);
+
+        mScreen.transitionTo(DOZE_AOD, DOZE_AOD_PAUSED);
+
+        assertNotEquals(1000, mServiceFake.screenBrightness);
+    }
+
+    @Test
+    public void testPulsing_usesLightSensor() throws Exception {
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE);
+        mScreen.transitionTo(DOZE, DOZE_REQUEST_PULSE);
+
+        mSensor.sendSensorEvent(1000);
+
+        assertEquals(1000, mServiceFake.screenBrightness);
+    }
+
+    @Test
+    public void testDozingAfterPulsing_pausesLightSensor() throws Exception {
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE);
+        mScreen.transitionTo(DOZE, DOZE_REQUEST_PULSE);
+        mScreen.transitionTo(DOZE_REQUEST_PULSE, DOZE_PULSING);
+        mScreen.transitionTo(DOZE_PULSING, DOZE_PULSE_DONE);
+        mScreen.transitionTo(DOZE_PULSE_DONE, DOZE);
+
+        mSensor.sendSensorEvent(1000);
+
+        assertNotEquals(1000, mServiceFake.screenBrightness);
+    }
+
+    @Test
+    public void testNullSensor() throws Exception {
+        mScreen = new DozeScreenBrightness(mContext, mServiceFake, mSensorManager,
+                null /* sensor */, null /* handler */);
+
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE_AOD);
+        mScreen.transitionTo(DOZE_AOD, DOZE_AOD_PAUSED);
+    }
+
+    @Test
+    public void testNoBrightnessDeliveredAfterFinish() throws Exception {
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE_AOD);
+        mScreen.transitionTo(DOZE_AOD, FINISH);
+
+        mSensor.sendSensorEvent(1000);
+
+        assertNotEquals(1000, mServiceFake.screenBrightness);
+    }
+
+    @Test
+    public void testBrightness_atLeastOne() throws Exception {
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE_AOD);
+
+        mSensor.sendSensorEvent(0);
+
+        assertTrue("Brightness must be at least 1, but was " + mServiceFake.screenBrightness,
+                mServiceFake.screenBrightness >= 1);
+    }
+}
\ No newline at end of file
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java
new file mode 100644
index 0000000..203876b
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java
@@ -0,0 +1,98 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.doze;
+
+import static com.android.systemui.doze.DozeMachine.State.DOZE;
+import static com.android.systemui.doze.DozeMachine.State.DOZE_AOD;
+import static com.android.systemui.doze.DozeMachine.State.DOZE_PULSING;
+import static com.android.systemui.doze.DozeMachine.State.DOZE_REQUEST_PULSE;
+import static com.android.systemui.doze.DozeMachine.State.INITIALIZED;
+import static com.android.systemui.doze.DozeMachine.State.UNINITIALIZED;
+
+import static org.junit.Assert.assertEquals;
+
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
+import android.view.Display;
+
+import com.android.systemui.SysuiTestCase;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class DozeScreenStateTest extends SysuiTestCase {
+
+    DozeServiceFake mServiceFake;
+    DozeScreenState mScreen;
+
+    @Before
+    public void setUp() throws Exception {
+        mServiceFake = new DozeServiceFake();
+        mScreen = new DozeScreenState(mServiceFake);
+    }
+
+    @Test
+    public void testScreen_offInDoze() {
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE);
+
+        assertEquals(Display.STATE_OFF, mServiceFake.screenState);
+    }
+
+    @Test
+    public void testScreen_onInAod() {
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE_AOD);
+
+        assertEquals(Display.STATE_DOZE_SUSPEND, mServiceFake.screenState);
+    }
+
+    @Test
+    public void testScreen_onInPulse() {
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE);
+
+        mScreen.transitionTo(DOZE, DOZE_REQUEST_PULSE);
+        mScreen.transitionTo(DOZE_REQUEST_PULSE, DOZE_PULSING);
+
+        assertEquals(Display.STATE_ON, mServiceFake.screenState);
+    }
+
+    @Test
+    public void testScreen_offInRequestPulseWithoutAoD() {
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE);
+
+        mScreen.transitionTo(DOZE, DOZE_REQUEST_PULSE);
+
+        assertEquals(Display.STATE_OFF, mServiceFake.screenState);
+    }
+
+    @Test
+    public void testScreen_onInRequestPulseWithAoD() {
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE_AOD);
+
+        mScreen.transitionTo(DOZE, DOZE_REQUEST_PULSE);
+
+        assertEquals(Display.STATE_DOZE_SUSPEND, mServiceFake.screenState);
+    }
+
+}
\ No newline at end of file
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeServiceFake.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeServiceFake.java
index c1e7fe4..34026b0 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeServiceFake.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeServiceFake.java
@@ -16,6 +16,7 @@
 
 package com.android.systemui.doze;
 
+import android.os.PowerManager;
 import android.view.Display;
 
 public class DozeServiceFake implements DozeMachine.Service {
@@ -23,6 +24,7 @@
     public boolean finished;
     public int screenState;
     public boolean requestedWakeup;
+    public int screenBrightness;
 
     public DozeServiceFake() {
         reset();
@@ -38,13 +40,19 @@
         screenState = state;
     }
 
-    public void reset() {
-        finished = false;
-        screenState = Display.STATE_UNKNOWN;
-    }
-
     @Override
     public void requestWakeUp() {
         requestedWakeup = true;
     }
+
+    @Override
+    public void setDozeScreenBrightness(int brightness) {
+        screenBrightness = brightness;
+    }
+
+    public void reset() {
+        finished = false;
+        screenState = Display.STATE_UNKNOWN;
+        screenBrightness = PowerManager.BRIGHTNESS_DEFAULT;
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/utils/hardware/FakeSensorManager.java b/packages/SystemUI/tests/src/com/android/systemui/utils/hardware/FakeSensorManager.java
index 30be665..a4ae166 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/utils/hardware/FakeSensorManager.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/utils/hardware/FakeSensorManager.java
@@ -30,13 +30,15 @@
 import android.os.SystemClock;
 import android.util.ArraySet;
 
-import com.google.android.collect.Lists;
+import com.android.internal.util.Preconditions;
 
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
+import java.util.stream.Collectors;
 
 /**
  * Rudimentary fake for SensorManager
@@ -49,6 +51,8 @@
 public class FakeSensorManager extends SensorManager {
 
     private final MockProximitySensor mMockProximitySensor;
+    private final FakeGenericSensor mFakeLightSensor;
+    private final FakeGenericSensor[] mSensors;
 
     public FakeSensorManager(Context context) throws Exception {
         Sensor proxSensor = context.getSystemService(SensorManager.class)
@@ -57,13 +61,21 @@
             // No prox? Let's create a fake one!
             proxSensor = createSensor(Sensor.TYPE_PROXIMITY);
         }
-        mMockProximitySensor = new MockProximitySensor(proxSensor);
+
+        mSensors = new FakeGenericSensor[]{
+                mMockProximitySensor = new MockProximitySensor(proxSensor),
+                mFakeLightSensor = new FakeGenericSensor(createSensor(Sensor.TYPE_LIGHT)),
+        };
     }
 
     public MockProximitySensor getMockProximitySensor() {
         return mMockProximitySensor;
     }
 
+    public FakeGenericSensor getFakeLightSensor() {
+        return mFakeLightSensor;
+    }
+
     @Override
     public Sensor getDefaultSensor(int type) {
         Sensor s = super.getDefaultSensor(type);
@@ -77,7 +89,10 @@
 
     @Override
     protected List<Sensor> getFullSensorList() {
-        return Lists.newArrayList(mMockProximitySensor.sensor);
+        return Arrays
+                .stream(mSensors)
+                .map(i -> i.mSensor)
+                .collect(Collectors.toList());
     }
 
     @Override
@@ -87,8 +102,11 @@
 
     @Override
     protected void unregisterListenerImpl(SensorEventListener listener, Sensor sensor) {
-        if (sensor == mMockProximitySensor.sensor || sensor == null) {
-            mMockProximitySensor.listeners.remove(listener);
+        Preconditions.checkNotNull(listener);
+        for (FakeGenericSensor s : mSensors) {
+            if (sensor == null || s.mSensor == sensor) {
+                s.mListeners.remove(listener);
+            }
         }
     }
 
@@ -96,9 +114,13 @@
     protected boolean registerListenerImpl(SensorEventListener listener, Sensor sensor,
             int delayUs,
             Handler handler, int maxReportLatencyUs, int reservedFlags) {
-        if (sensor == mMockProximitySensor.sensor) {
-            mMockProximitySensor.listeners.add(listener);
-            return true;
+        Preconditions.checkNotNull(sensor);
+        Preconditions.checkNotNull(listener);
+        for (FakeGenericSensor s : mSensors) {
+            if (s.mSensor == sensor) {
+                s.mListeners.add(listener);
+                return true;
+            }
         }
         return false;
     }
@@ -196,18 +218,35 @@
         setter.invoke(sensor, type);
     }
 
-    public class MockProximitySensor {
-        final Sensor sensor;
-        final ArraySet<SensorEventListener> listeners = new ArraySet<>();
+    public class MockProximitySensor extends FakeGenericSensor {
 
         private MockProximitySensor(Sensor sensor) {
-            this.sensor = sensor;
+            super(sensor);
         }
 
         public void sendProximityResult(boolean far) {
-            SensorEvent event = createSensorEvent(1);
-            event.values[0] = far ? sensor.getMaximumRange() : 0;
-            for (SensorEventListener listener : listeners) {
+            sendSensorEvent(far ? getSensor().getMaximumRange() : 0);
+        }
+    }
+
+    public class FakeGenericSensor {
+
+        private final Sensor mSensor;
+        private final ArraySet<SensorEventListener> mListeners = new ArraySet<>();
+
+        public FakeGenericSensor(
+                Sensor sensor) {
+            this.mSensor = sensor;
+        }
+
+        public Sensor getSensor() {
+            return mSensor;
+        }
+
+        public void sendSensorEvent(float... values) {
+            SensorEvent event = createSensorEvent(values.length);
+            System.arraycopy(values, 0, event.values, 0, values.length);
+            for (SensorEventListener listener : mListeners) {
                 listener.onSensorChanged(event);
             }
         }
@@ -222,7 +261,7 @@
             } catch (Exception e) {
                 throw new RuntimeException(e);
             }
-            event.sensor = sensor;
+            event.sensor = mSensor;
             event.timestamp = SystemClock.elapsedRealtimeNanos();
 
             return event;