Doze: Fix WakeLock bug

Change-Id: I3d687bbf82173723c420058867531a970d1f46cb
Fixes: 32857549
Test: runtest -x $(gettop)/frameworks/base/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java

2 files changed, 15 insertions, 0 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java b/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java
index c633aa15c8e3..951b27f75aeb 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java
@@ -225,8 +225,10 @@ public class DozeMachine {
         boolean newPolicy = wakeLockPolicy(newState);
         if (mWakeLockHeldForCurrentState && !newPolicy) {
             mWakeLock.release();
+            mWakeLockHeldForCurrentState = false;
         } else if (!mWakeLockHeldForCurrentState && newPolicy) {
             mWakeLock.acquire();
+            mWakeLockHeldForCurrentState = true;
         }
     }
 
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 ba7c923efe71..863f0e5a9897 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeMachineTest.java
@@ -189,6 +189,19 @@ public class DozeMachineTest {
 
     @Test
     @UiThreadTest
+    public void testWakeLock_releasedAfterPulse() {
+        mMachine.requestState(INITIALIZED);
+
+        mMachine.requestState(DOZE);
+        mMachine.requestState(DOZE_REQUEST_PULSE);
+        mMachine.requestState(DOZE_PULSING);
+        mMachine.requestState(DOZE_PULSE_DONE);
+
+        assertFalse(mWakeLockFake.isHeld());
+    }
+
+    @Test
+    @UiThreadTest
     public void testScreen_offInDoze() {
         mMachine.requestState(INITIALIZED);