2 files changed, 17 insertions, 1 deletions
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java
index e50f1fb2a3ee..915359374bfe 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java
@@ -92,7 +92,7 @@ public class DozeScreenState implements DozeMachine.Part {
         final boolean pulseEnding = oldState == DOZE_PULSE_DONE && newState.isAlwaysOn();
         final boolean turningOn = (oldState == DOZE_AOD_PAUSED || oldState == DOZE)
                 && newState.isAlwaysOn();
-        final boolean turningOff = (newState.isAlwaysOn() && newState == DOZE)
+        final boolean turningOff = (oldState.isAlwaysOn() && newState == DOZE)
                 || (oldState == DOZE_AOD_PAUSING && newState == DOZE_AOD_PAUSED);
         final boolean justInitialized = oldState == DozeMachine.State.INITIALIZED;
         if (messagePending || justInitialized || pulseEnding || turningOn) {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java
index e8a3c0ec2e98..41d7fd64fe7a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java
@@ -217,4 +217,20 @@ public class DozeScreenStateTest extends SysuiTestCase {
         assertEquals(Display.STATE_OFF, mServiceFake.screenState);
     }
 
+    @Test
+    public void test_animatesOff() {
+        ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
+        doAnswer(invocation -> null).when(mDozeHost).prepareForGentleSleep(captor.capture());
+        mHandlerFake.setMode(QUEUEING);
+
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE_AOD);
+        mScreen.transitionTo(DOZE_AOD, DOZE);
+
+        mHandlerFake.dispatchQueuedMessages();
+        verify(mDozeHost).prepareForGentleSleep(eq(captor.getValue()));
+        captor.getValue().run();
+        assertEquals(Display.STATE_OFF, mServiceFake.screenState);
+    }
+
 }
 \ No newline at end of file