Don't trust DozeState onScreenState changes

In DozeScreenBrightness, we should only rely on the passed
displayScreenState because our DozeState (updated by transitionTo)
may not be updated yet. Therefore, only register the brightness sensor
when the display state is Display.STATE_DOZE/DOZE_SUSPENDED.

Test: atest DozeScreenBrightnessTest
Bug: 168689990
Change-Id: Ie32f91c5a46bd987649a8a17e6543071847ad97c

packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java

@@ -432,8 +432,12 @@ public class DozeMachine {
         /** Give the Part a chance to clean itself up. */
         default void destroy() {}
 
-        /** Alerts that the screenstate is being changed. */
+        /**
-        default void onScreenState(int state) {}
+         *  Alerts that the screenstate is being changed.
+         *  Note: This may be called from within a call to transitionTo, so local DozeState may not
+         *  be accurate nor match with the new displayState.
+         */
+        default void onScreenState(int displayState) {}
 
         /** Sets the {@link DozeMachine} when this Part is associated with one. */
         default void setDozeMachine(DozeMachine dozeMachine) {}

packages/SystemUI/src/com/android/systemui/doze/DozeScreenBrightness.java

@@ -73,7 +73,6 @@ public class DozeScreenBrightness extends BroadcastReceiver implements DozeMachi
      * --ei brightness_bucket 1}
      */
     private int mDebugBrightnessBucket = -1;
-    private DozeMachine.State mState;
 
     @Inject
     public DozeScreenBrightness(Context context, @WrappedService DozeMachine.Service service,
@@ -94,7 +93,6 @@ public class DozeScreenBrightness extends BroadcastReceiver implements DozeMachi
 
     @Override
     public void transitionTo(DozeMachine.State oldState, DozeMachine.State newState) {
-        mState = newState;
         switch (newState) {
             case INITIALIZED:
             case DOZE:
@@ -112,10 +110,7 @@ public class DozeScreenBrightness extends BroadcastReceiver implements DozeMachi
 
     @Override
     public void onScreenState(int state) {
-        if (!mScreenOff
+        if (state == Display.STATE_DOZE || state == Display.STATE_DOZE_SUSPEND) {
-                && (mState == DozeMachine.State.DOZE_AOD
-                     || mState == DozeMachine.State.DOZE_AOD_DOCKED)
-                && (state == Display.STATE_DOZE || state == Display.STATE_DOZE_SUSPEND)) {
             setLightSensorEnabled(true);
         } else {
             setLightSensorEnabled(false);

packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenBrightnessTest.java

@@ -115,8 +115,6 @@ public class DozeScreenBrightnessTest extends SysuiTestCase {
 
     @Test
     public void testAod_usesLightSensor() {
-        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
-        mScreen.transitionTo(INITIALIZED, DOZE_AOD);
         mScreen.onScreenState(Display.STATE_DOZE);
         waitForSensorManager();
 
@@ -127,8 +125,6 @@ public class DozeScreenBrightnessTest extends SysuiTestCase {
 
     @Test
     public void testAod_usesDebugValue() throws Exception {
-        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
-        mScreen.transitionTo(INITIALIZED, DOZE_AOD);
         mScreen.onScreenState(Display.STATE_DOZE);
         waitForSensorManager();
 
@@ -181,14 +177,13 @@ public class DozeScreenBrightnessTest extends SysuiTestCase {
     }
 
     @Test
-    public void testDozingAfterPulsing_pausesLightSensor() throws Exception {
+    public void testScreenOffAfterPulsing_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);
-        mScreen.onScreenState(Display.STATE_DOZE);
         waitForSensorManager();
 
         mSensor.sendSensorEvent(1);
@@ -196,6 +191,18 @@ public class DozeScreenBrightnessTest extends SysuiTestCase {
         assertEquals(DEFAULT_BRIGHTNESS, mServiceFake.screenBrightness);
     }
 
+    @Test
+    public void testOnScreenStateSetBeforeTransition_stillRegistersSensor() {
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.onScreenState(Display.STATE_DOZE);
+        mScreen.transitionTo(INITIALIZED, DOZE_AOD);
+        waitForSensorManager();
+
+        mSensor.sendSensorEvent(1);
+
+        assertEquals(1, mServiceFake.screenBrightness);
+    }
+
     @Test
     public void testNullSensor() throws Exception {
         mScreen = new DozeScreenBrightness(mContext, mServiceFake, mSensorManager,
@@ -206,12 +213,15 @@ public class DozeScreenBrightnessTest extends SysuiTestCase {
         mScreen.transitionTo(INITIALIZED, DOZE_AOD);
         mScreen.transitionTo(DOZE_AOD, DOZE_AOD_PAUSING);
         mScreen.transitionTo(DOZE_AOD_PAUSING, DOZE_AOD_PAUSED);
+        mScreen.onScreenState(Display.STATE_DOZE);
+        mScreen.onScreenState(Display.STATE_OFF);
     }
 
     @Test
     public void testNoBrightnessDeliveredAfterFinish() throws Exception {
         mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
         mScreen.transitionTo(INITIALIZED, DOZE_AOD);
+        mScreen.onScreenState(Display.STATE_DOZE);
         mScreen.transitionTo(DOZE_AOD, FINISH);
         waitForSensorManager();