Keyguard: Factor out Screen and Wakefulness lifecycles

Test: runtest systemui
Bug: 62446740
Change-Id: Ife34c1f4299b152a6352445adc8c9fc3c757e87c
diff --git a/core/java/com/android/internal/policy/IKeyguardService.aidl b/core/java/com/android/internal/policy/IKeyguardService.aidl
index 3f35c91..69184c3 100644
--- a/core/java/com/android/internal/policy/IKeyguardService.aidl
+++ b/core/java/com/android/internal/policy/IKeyguardService.aidl
@@ -64,6 +64,11 @@
     void onStartedWakingUp();
 
     /**
+     * Called when the device has finished waking up.
+     */
+    void onFinishedWakingUp();
+
+    /**
      * Called when the device screen is turning on.
      */
     void onScreenTurningOn(IKeyguardDrawnCallback callback);
@@ -74,6 +79,11 @@
     void onScreenTurnedOn();
 
     /**
+     * Called when the screen starts turning off.
+     */
+    void onScreenTurningOff();
+
+    /**
      * Called when the screen has turned off.
      */
     void onScreenTurnedOff();
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitorCallback.java b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitorCallback.java
index 14d6b59..5a02178 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitorCallback.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitorCallback.java
@@ -154,13 +154,19 @@
 
     /**
      * Called when the device has started waking up.
+     *
+     * @deprecated use {@link com.android.systemui.keyguard.WakefulnessLifecycle}.
      */
+    @Deprecated
     public void onStartedWakingUp() { }
 
     /**
      * Called when the device has started going to sleep.
      * @param why see {@link #onFinishedGoingToSleep(int)}
+     *
+     * @deprecated use {@link com.android.systemui.keyguard.WakefulnessLifecycle}.
      */
+    @Deprecated
     public void onStartedGoingToSleep(int why) { }
 
     /**
@@ -168,17 +174,26 @@
      * @param why either {@link WindowManagerPolicy#OFF_BECAUSE_OF_ADMIN},
      * {@link WindowManagerPolicy#OFF_BECAUSE_OF_USER}, or
      * {@link WindowManagerPolicy#OFF_BECAUSE_OF_TIMEOUT}.
+     *
+     * @deprecated use {@link com.android.systemui.keyguard.WakefulnessLifecycle}.
      */
+    @Deprecated
     public void onFinishedGoingToSleep(int why) { }
 
     /**
      * Called when the screen has been turned on.
+     *
+     * @deprecated use {@link com.android.systemui.keyguard.ScreenLifecycle}.
      */
+    @Deprecated
     public void onScreenTurnedOn() { }
 
     /**
      * Called when the screen has been turned off.
+     *
+     * @deprecated use {@link com.android.systemui.keyguard.ScreenLifecycle}.
      */
+    @Deprecated
     public void onScreenTurnedOff() { }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/Dependency.java b/packages/SystemUI/src/com/android/systemui/Dependency.java
index 776d076..e5c729f 100644
--- a/packages/SystemUI/src/com/android/systemui/Dependency.java
+++ b/packages/SystemUI/src/com/android/systemui/Dependency.java
@@ -30,6 +30,8 @@
 import com.android.systemui.assist.AssistManager;
 import com.android.systemui.colorextraction.SysuiColorExtractor;
 import com.android.systemui.fragments.FragmentService;
+import com.android.systemui.keyguard.ScreenLifecycle;
+import com.android.systemui.keyguard.WakefulnessLifecycle;
 import com.android.systemui.plugins.ActivityStarter;
 import com.android.systemui.plugins.PluginActivityManager;
 import com.android.systemui.plugins.PluginDependencyProvider;
@@ -248,6 +250,12 @@
         mProviders.put(StatusBarIconController.class, () ->
                 new StatusBarIconControllerImpl(mContext));
 
+        mProviders.put(ScreenLifecycle.class, () ->
+                new ScreenLifecycle());
+
+        mProviders.put(WakefulnessLifecycle.class, () ->
+                new WakefulnessLifecycle());
+
         mProviders.put(FragmentService.class, () ->
                 new FragmentService(mContext));
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardLifecyclesDispatcher.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardLifecyclesDispatcher.java
new file mode 100644
index 0000000..4c98c08
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardLifecyclesDispatcher.java
@@ -0,0 +1,86 @@
+/*
+ * 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.keyguard;
+
+import android.os.Handler;
+import android.os.Message;
+
+import com.android.internal.policy.IKeyguardDrawnCallback;
+
+/**
+ * Dispatches the lifecycles keyguard gets from WindowManager on the main thread.
+ */
+public class KeyguardLifecyclesDispatcher {
+
+    static final int SCREEN_TURNING_ON = 0;
+    static final int SCREEN_TURNED_ON = 1;
+    static final int SCREEN_TURNING_OFF = 2;
+    static final int SCREEN_TURNED_OFF = 3;
+
+    static final int STARTED_WAKING_UP = 4;
+    static final int FINISHED_WAKING_UP = 5;
+    static final int STARTED_GOING_TO_SLEEP = 6;
+    static final int FINISHED_GOING_TO_SLEEP = 7;
+
+    private final ScreenLifecycle mScreenLifecycle;
+    private final WakefulnessLifecycle mWakefulnessLifecycle;
+
+    public KeyguardLifecyclesDispatcher(ScreenLifecycle screenLifecycle,
+            WakefulnessLifecycle wakefulnessLifecycle) {
+        mScreenLifecycle = screenLifecycle;
+        mWakefulnessLifecycle = wakefulnessLifecycle;
+    }
+
+    void dispatch(int what) {
+        mHandler.obtainMessage(what).sendToTarget();
+    }
+
+    private Handler mHandler = new Handler() {
+        @Override
+        public void handleMessage(Message msg) {
+            switch (msg.what) {
+                case SCREEN_TURNING_ON:
+                    mScreenLifecycle.dispatchScreenTurningOn();
+                    break;
+                case SCREEN_TURNED_ON:
+                    mScreenLifecycle.dispatchScreenTurnedOn();
+                    break;
+                case SCREEN_TURNING_OFF:
+                    mScreenLifecycle.dispatchScreenTurningOff();
+                    break;
+                case SCREEN_TURNED_OFF:
+                    mScreenLifecycle.dispatchScreenTurnedOff();
+                    break;
+                case STARTED_WAKING_UP:
+                    mWakefulnessLifecycle.dispatchStartedWakingUp();
+                    break;
+                case FINISHED_WAKING_UP:
+                    mWakefulnessLifecycle.dispatchFinishedWakingUp();
+                    break;
+                case STARTED_GOING_TO_SLEEP:
+                    mWakefulnessLifecycle.dispatchStartedGoingToSleep();
+                    break;
+                case FINISHED_GOING_TO_SLEEP:
+                    mWakefulnessLifecycle.dispatchFinishedGoingToSleep();
+                    break;
+                default:
+                    throw new IllegalArgumentException("Unknown message: " + msg);
+            }
+        }
+    };
+
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
index 2d5d198..2a5ae0d 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
@@ -31,6 +31,7 @@
 import com.android.internal.policy.IKeyguardExitCallback;
 import com.android.internal.policy.IKeyguardService;
 import com.android.internal.policy.IKeyguardStateCallback;
+import com.android.systemui.Dependency;
 import com.android.systemui.SystemUIApplication;
 
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
@@ -40,12 +41,17 @@
     static final String PERMISSION = android.Manifest.permission.CONTROL_KEYGUARD;
 
     private KeyguardViewMediator mKeyguardViewMediator;
+    private KeyguardLifecyclesDispatcher mKeyguardLifecyclesDispatcher;
 
     @Override
     public void onCreate() {
         ((SystemUIApplication) getApplication()).startServicesIfNeeded();
         mKeyguardViewMediator =
                 ((SystemUIApplication) getApplication()).getComponent(KeyguardViewMediator.class);
+        mKeyguardLifecyclesDispatcher = new KeyguardLifecyclesDispatcher(
+                Dependency.get(ScreenLifecycle.class),
+                Dependency.get(WakefulnessLifecycle.class));
+
     }
 
     @Override
@@ -111,12 +117,16 @@
         public void onStartedGoingToSleep(int reason) {
             checkPermission();
             mKeyguardViewMediator.onStartedGoingToSleep(reason);
+            mKeyguardLifecyclesDispatcher.dispatch(
+                    KeyguardLifecyclesDispatcher.STARTED_GOING_TO_SLEEP);
         }
 
         @Override // Binder interface
         public void onFinishedGoingToSleep(int reason, boolean cameraGestureTriggered) {
             checkPermission();
             mKeyguardViewMediator.onFinishedGoingToSleep(reason, cameraGestureTriggered);
+            mKeyguardLifecyclesDispatcher.dispatch(
+                    KeyguardLifecyclesDispatcher.FINISHED_GOING_TO_SLEEP);
         }
 
         @Override // Binder interface
@@ -124,6 +134,15 @@
             Trace.beginSection("KeyguardService.mBinder#onStartedWakingUp");
             checkPermission();
             mKeyguardViewMediator.onStartedWakingUp();
+            mKeyguardLifecyclesDispatcher.dispatch(KeyguardLifecyclesDispatcher.STARTED_WAKING_UP);
+            Trace.endSection();
+        }
+
+        @Override // Binder interface
+        public void onFinishedWakingUp() {
+            Trace.beginSection("KeyguardService.mBinder#onFinishedWakingUp");
+            checkPermission();
+            mKeyguardLifecyclesDispatcher.dispatch(KeyguardLifecyclesDispatcher.FINISHED_WAKING_UP);
             Trace.endSection();
         }
 
@@ -132,6 +151,7 @@
             Trace.beginSection("KeyguardService.mBinder#onScreenTurningOn");
             checkPermission();
             mKeyguardViewMediator.onScreenTurningOn(callback);
+            mKeyguardLifecyclesDispatcher.dispatch(KeyguardLifecyclesDispatcher.SCREEN_TURNING_ON);
             Trace.endSection();
         }
 
@@ -140,13 +160,21 @@
             Trace.beginSection("KeyguardService.mBinder#onScreenTurnedOn");
             checkPermission();
             mKeyguardViewMediator.onScreenTurnedOn();
+            mKeyguardLifecyclesDispatcher.dispatch(KeyguardLifecyclesDispatcher.SCREEN_TURNED_ON);
             Trace.endSection();
         }
 
         @Override // Binder interface
+        public void onScreenTurningOff() {
+            checkPermission();
+            mKeyguardLifecyclesDispatcher.dispatch(KeyguardLifecyclesDispatcher.SCREEN_TURNING_OFF);
+        }
+
+        @Override // Binder interface
         public void onScreenTurnedOff() {
             checkPermission();
             mKeyguardViewMediator.onScreenTurnedOff();
+            mKeyguardLifecyclesDispatcher.dispatch(KeyguardLifecyclesDispatcher.SCREEN_TURNED_OFF);
         }
 
         @Override // Binder interface
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/Lifecycle.java b/packages/SystemUI/src/com/android/systemui/keyguard/Lifecycle.java
new file mode 100644
index 0000000..1b20cfb
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/Lifecycle.java
@@ -0,0 +1,42 @@
+/*
+ * 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.keyguard;
+
+import java.util.ArrayList;
+import java.util.function.Consumer;
+
+/**
+ * Base class for lifecycles with observers.
+ */
+public class Lifecycle<T> {
+
+    private ArrayList<T> mObservers = new ArrayList<>();
+
+    public void addObserver(T observer) {
+        mObservers.add(observer);
+    }
+
+    public void removeObserver(T observer) {
+        mObservers.remove(observer);
+    }
+
+    public void dispatch(Consumer<T> consumer) {
+        for (int i = 0; i < mObservers.size(); i++) {
+            consumer.accept(mObservers.get(i));
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ScreenLifecycle.java b/packages/SystemUI/src/com/android/systemui/keyguard/ScreenLifecycle.java
new file mode 100644
index 0000000..3f39dfe
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ScreenLifecycle.java
@@ -0,0 +1,73 @@
+/*
+ * 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.keyguard;
+
+import com.android.systemui.Dumpable;
+
+import java.io.FileDescriptor;
+import java.io.PrintWriter;
+
+/**
+ * Tracks the screen lifecycle.
+ */
+public class ScreenLifecycle extends Lifecycle<ScreenLifecycle.Observer> implements Dumpable {
+
+    public static final int SCREEN_OFF = 0;
+    public static final int SCREEN_TURNING_ON = 1;
+    public static final int SCREEN_ON = 2;
+    public static final int SCREEN_TURNING_OFF = 3;
+
+    private int mScreenState = SCREEN_OFF;
+
+    public int getScreenState() {
+        return mScreenState;
+    }
+
+    public void dispatchScreenTurningOn() {
+        mScreenState = SCREEN_TURNING_ON;
+        dispatch(Observer::onScreenTurningOn);
+    }
+
+    public void dispatchScreenTurnedOn() {
+        mScreenState = SCREEN_ON;
+        dispatch(Observer::onScreenTurnedOn);
+    }
+
+    public void dispatchScreenTurningOff() {
+        mScreenState = SCREEN_TURNING_OFF;
+        dispatch(Observer::onScreenTurningOff);
+    }
+
+    public void dispatchScreenTurnedOff() {
+        mScreenState = SCREEN_OFF;
+        dispatch(Observer::onScreenTurnedOff);
+    }
+
+    @Override
+    public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
+        pw.println("ScreenLifecycle:");
+        pw.println("  mScreenState=" + mScreenState);
+    }
+
+    public interface Observer {
+        default void onScreenTurningOn() {}
+        default void onScreenTurnedOn() {}
+        default void onScreenTurningOff() {}
+        default void onScreenTurnedOff() {}
+    }
+
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/WakefulnessLifecycle.java b/packages/SystemUI/src/com/android/systemui/keyguard/WakefulnessLifecycle.java
new file mode 100644
index 0000000..578e6fb
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/WakefulnessLifecycle.java
@@ -0,0 +1,73 @@
+/*
+ * 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.keyguard;
+
+import com.android.systemui.Dumpable;
+
+import java.io.FileDescriptor;
+import java.io.PrintWriter;
+
+/**
+ * Tracks the wakefulness lifecycle.
+ */
+public class WakefulnessLifecycle extends Lifecycle<WakefulnessLifecycle.Observer> implements
+        Dumpable {
+
+    public static final int WAKEFULNESS_ASLEEP = 0;
+    public static final int WAKEFULNESS_WAKING = 1;
+    public static final int WAKEFULNESS_AWAKE = 2;
+    public static final int WAKEFULNESS_GOING_TO_SLEEP = 3;
+
+    private int mWakefulness = WAKEFULNESS_ASLEEP;
+
+    public int getWakefulness() {
+        return mWakefulness;
+    }
+
+    public void dispatchStartedWakingUp() {
+        mWakefulness = WAKEFULNESS_WAKING;
+        dispatch(Observer::onStartedWakingUp);
+    }
+
+    public void dispatchFinishedWakingUp() {
+        mWakefulness = WAKEFULNESS_AWAKE;
+        dispatch(Observer::onFinishedWakingUp);
+    }
+
+    public void dispatchStartedGoingToSleep() {
+        mWakefulness = WAKEFULNESS_GOING_TO_SLEEP;
+        dispatch(Observer::onStartedGoingToSleep);
+    }
+
+    public void dispatchFinishedGoingToSleep() {
+        mWakefulness = WAKEFULNESS_ASLEEP;
+        dispatch(Observer::onFinishedGoingToSleep);
+    }
+
+    @Override
+    public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
+        pw.println("WakefulnessLifecycle:");
+        pw.println("  mWakefulness=" + mWakefulness);
+    }
+
+    public interface Observer {
+        default void onStartedWakingUp() {}
+        default void onFinishedWakingUp() {}
+        default void onStartedGoingToSleep() {}
+        default void onFinishedGoingToSleep() {}
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/LifecycleTest.java b/packages/SystemUI/tests/src/com/android/systemui/keyguard/LifecycleTest.java
new file mode 100644
index 0000000..e0807d6
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/LifecycleTest.java
@@ -0,0 +1,79 @@
+/*
+ * 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.keyguard;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import android.support.test.filters.SmallTest;
+import android.testing.AndroidTestingRunner;
+
+import com.android.systemui.SysuiTestCase;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.ArrayList;
+
+@RunWith(AndroidTestingRunner.class)
+@SmallTest
+public class LifecycleTest extends SysuiTestCase {
+
+    private final Object mObj1 = new Object();
+    private final Object mObj2 = new Object();
+
+    private Lifecycle<Object> mLifecycle;
+    private ArrayList<Object> mDispatchedObjects;
+
+    @Before
+    public void setUp() throws Exception {
+        mLifecycle = new Lifecycle<>();
+        mDispatchedObjects = new ArrayList<>();
+    }
+
+    @Test
+    public void addObserver_addsObserver() throws Exception {
+        mLifecycle.addObserver(mObj1);
+
+        mLifecycle.dispatch(mDispatchedObjects::add);
+
+        assertTrue(mDispatchedObjects.contains(mObj1));
+    }
+
+    @Test
+    public void removeObserver() throws Exception {
+        mLifecycle.addObserver(mObj1);
+        mLifecycle.removeObserver(mObj1);
+
+        mLifecycle.dispatch(mDispatchedObjects::add);
+
+        assertFalse(mDispatchedObjects.contains(mObj1));
+    }
+
+    @Test
+    public void dispatch() throws Exception {
+        mLifecycle.addObserver(mObj1);
+        mLifecycle.addObserver(mObj2);
+
+        mLifecycle.dispatch(mDispatchedObjects::add);
+
+        assertTrue(mDispatchedObjects.contains(mObj1));
+        assertTrue(mDispatchedObjects.contains(mObj2));
+    }
+
+}
\ No newline at end of file
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ScreenLifecycleTest.java b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ScreenLifecycleTest.java
new file mode 100644
index 0000000..8c918f6
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ScreenLifecycleTest.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.keyguard;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import android.support.test.filters.SmallTest;
+import android.testing.AndroidTestingRunner;
+
+import com.android.systemui.SysuiTestCase;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.ByteArrayOutputStream;
+import java.io.PrintWriter;
+
+@RunWith(AndroidTestingRunner.class)
+@SmallTest
+public class ScreenLifecycleTest extends SysuiTestCase {
+
+    private ScreenLifecycle mScreen;
+    private ScreenLifecycle.Observer mScreenObserverMock;
+
+    @Before
+    public void setUp() throws Exception {
+        mScreen = new ScreenLifecycle();
+        mScreenObserverMock = mock(ScreenLifecycle.Observer.class);
+        mScreen.addObserver(mScreenObserverMock);
+    }
+
+    @Test
+    public void baseState() throws Exception {
+        assertEquals(ScreenLifecycle.SCREEN_OFF, mScreen.getScreenState());
+        verifyNoMoreInteractions(mScreenObserverMock);
+    }
+
+    @Test
+    public void screenTurningOn() throws Exception {
+        mScreen.dispatchScreenTurningOn();
+
+        assertEquals(ScreenLifecycle.SCREEN_TURNING_ON, mScreen.getScreenState());
+        verify(mScreenObserverMock).onScreenTurningOn();
+    }
+
+    @Test
+    public void screenTurnedOn() throws Exception {
+        mScreen.dispatchScreenTurningOn();
+        mScreen.dispatchScreenTurnedOn();
+
+        assertEquals(ScreenLifecycle.SCREEN_ON, mScreen.getScreenState());
+        verify(mScreenObserverMock).onScreenTurnedOn();
+    }
+
+    @Test
+    public void screenTurningOff() throws Exception {
+        mScreen.dispatchScreenTurningOn();
+        mScreen.dispatchScreenTurnedOn();
+        mScreen.dispatchScreenTurningOff();
+
+        assertEquals(ScreenLifecycle.SCREEN_TURNING_OFF, mScreen.getScreenState());
+        verify(mScreenObserverMock).onScreenTurningOff();
+    }
+
+    @Test
+    public void screenTurnedOff() throws Exception {
+        mScreen.dispatchScreenTurningOn();
+        mScreen.dispatchScreenTurnedOn();
+        mScreen.dispatchScreenTurningOff();
+        mScreen.dispatchScreenTurnedOff();
+
+        assertEquals(ScreenLifecycle.SCREEN_OFF, mScreen.getScreenState());
+        verify(mScreenObserverMock).onScreenTurnedOff();
+    }
+
+    @Test
+    public void dump() throws Exception {
+        mScreen.dump(null, new PrintWriter(new ByteArrayOutputStream()), new String[0]);
+    }
+}
\ No newline at end of file
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/WakefulnessLifecycleTest.java b/packages/SystemUI/tests/src/com/android/systemui/keyguard/WakefulnessLifecycleTest.java
new file mode 100644
index 0000000..e15e0b4
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/WakefulnessLifecycleTest.java
@@ -0,0 +1,106 @@
+/*
+ * 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.keyguard;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import android.support.test.filters.SmallTest;
+import android.testing.AndroidTestingRunner;
+
+import com.android.systemui.SysuiTestCase;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.ByteArrayOutputStream;
+import java.io.PrintWriter;
+
+@RunWith(AndroidTestingRunner.class)
+@SmallTest
+public class WakefulnessLifecycleTest extends SysuiTestCase {
+
+    private WakefulnessLifecycle mWakefulness;
+    private WakefulnessLifecycle.Observer mWakefulnessObserver;
+
+    @Before
+    public void setUp() throws Exception {
+        mWakefulness = new WakefulnessLifecycle();
+        mWakefulnessObserver = mock(WakefulnessLifecycle.Observer.class);
+        mWakefulness.addObserver(mWakefulnessObserver);
+    }
+
+    @Test
+    public void baseState() throws Exception {
+        assertEquals(WakefulnessLifecycle.WAKEFULNESS_ASLEEP, mWakefulness.getWakefulness());
+
+        verifyNoMoreInteractions(mWakefulnessObserver);
+    }
+
+    @Test
+    public void dispatchStartedWakingUp() throws Exception {
+        mWakefulness.dispatchStartedWakingUp();
+
+        assertEquals(WakefulnessLifecycle.WAKEFULNESS_WAKING, mWakefulness.getWakefulness());
+
+        verify(mWakefulnessObserver).onStartedWakingUp();
+    }
+
+    @Test
+    public void dispatchFinishedWakingUp() throws Exception {
+        mWakefulness.dispatchStartedWakingUp();
+        mWakefulness.dispatchFinishedWakingUp();
+
+        assertEquals(WakefulnessLifecycle.WAKEFULNESS_AWAKE, mWakefulness.getWakefulness());
+
+        verify(mWakefulnessObserver).onFinishedWakingUp();
+    }
+
+    @Test
+    public void dispatchStartedGoingToSleep() throws Exception {
+        mWakefulness.dispatchStartedWakingUp();
+        mWakefulness.dispatchFinishedWakingUp();
+        mWakefulness.dispatchStartedGoingToSleep();
+
+        assertEquals(WakefulnessLifecycle.WAKEFULNESS_GOING_TO_SLEEP,
+                mWakefulness.getWakefulness());
+
+        verify(mWakefulnessObserver).onStartedGoingToSleep();
+    }
+
+    @Test
+    public void dispatchFinishedGoingToSleep() throws Exception {
+        mWakefulness.dispatchStartedWakingUp();
+        mWakefulness.dispatchFinishedWakingUp();
+        mWakefulness.dispatchStartedGoingToSleep();
+        mWakefulness.dispatchFinishedGoingToSleep();
+
+        assertEquals(WakefulnessLifecycle.WAKEFULNESS_ASLEEP,
+                mWakefulness.getWakefulness());
+
+        verify(mWakefulnessObserver).onFinishedGoingToSleep();
+    }
+
+    @Test
+    public void dump() throws Exception {
+        mWakefulness.dump(null, new PrintWriter(new ByteArrayOutputStream()), new String[0]);
+    }
+
+}
\ No newline at end of file
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index 75fc25a..bc28f97 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -6606,6 +6606,10 @@
     @Override
     public void finishedWakingUp() {
         if (DEBUG_WAKEUP) Slog.i(TAG, "Finished waking up...");
+
+        if (mKeyguardDelegate != null) {
+            mKeyguardDelegate.onFinishedWakingUp();
+        }
     }
 
     private void wakeUpFromPowerKey(long eventTime) {
@@ -6714,6 +6718,11 @@
     @Override
     public void screenTurningOff(ScreenOffListener screenOffListener) {
         mWindowManagerFuncs.screenTurningOff(screenOffListener);
+        synchronized (mLock) {
+            if (mKeyguardDelegate != null) {
+                mKeyguardDelegate.onScreenTurningOff();
+            }
+        }
     }
 
     private void reportScreenStateToVrManager(boolean isScreenOn) {
diff --git a/services/core/java/com/android/server/policy/keyguard/KeyguardServiceDelegate.java b/services/core/java/com/android/server/policy/keyguard/KeyguardServiceDelegate.java
index 0121ee1..50e5e7b 100644
--- a/services/core/java/com/android/server/policy/keyguard/KeyguardServiceDelegate.java
+++ b/services/core/java/com/android/server/policy/keyguard/KeyguardServiceDelegate.java
@@ -35,10 +35,12 @@
     private static final int SCREEN_STATE_OFF = 0;
     private static final int SCREEN_STATE_TURNING_ON = 1;
     private static final int SCREEN_STATE_ON = 2;
+    private static final int SCREEN_STATE_TURNING_OFF = 3;
 
     private static final int INTERACTIVE_STATE_SLEEP = 0;
-    private static final int INTERACTIVE_STATE_AWAKE = 1;
-    private static final int INTERACTIVE_STATE_GOING_TO_SLEEP = 2;
+    private static final int INTERACTIVE_STATE_WAKING = 1;
+    private static final int INTERACTIVE_STATE_AWAKE = 2;
+    private static final int INTERACTIVE_STATE_GOING_TO_SLEEP = 3;
 
     protected KeyguardServiceWrapper mKeyguardService;
     private final Context mContext;
@@ -164,9 +166,13 @@
                     mKeyguardService.setCurrentUser(mKeyguardState.currentUser);
                 }
                 // This is used to hide the scrim once keyguard displays.
-                if (mKeyguardState.interactiveState == INTERACTIVE_STATE_AWAKE) {
+                if (mKeyguardState.interactiveState == INTERACTIVE_STATE_AWAKE
+                        || mKeyguardState.interactiveState == INTERACTIVE_STATE_WAKING) {
                     mKeyguardService.onStartedWakingUp();
                 }
+                if (mKeyguardState.interactiveState == INTERACTIVE_STATE_AWAKE) {
+                    mKeyguardService.onFinishedWakingUp();
+                }
                 if (mKeyguardState.screenState == SCREEN_STATE_ON
                         || mKeyguardState.screenState == SCREEN_STATE_TURNING_ON) {
                     mKeyguardService.onScreenTurningOn(
@@ -277,9 +283,25 @@
             if (DEBUG) Log.v(TAG, "onStartedWakingUp()");
             mKeyguardService.onStartedWakingUp();
         }
+        mKeyguardState.interactiveState = INTERACTIVE_STATE_WAKING;
+    }
+
+    public void onFinishedWakingUp() {
+        if (mKeyguardService != null) {
+            if (DEBUG) Log.v(TAG, "onFinishedWakingUp()");
+            mKeyguardService.onFinishedWakingUp();
+        }
         mKeyguardState.interactiveState = INTERACTIVE_STATE_AWAKE;
     }
 
+    public void onScreenTurningOff() {
+        if (mKeyguardService != null) {
+            if (DEBUG) Log.v(TAG, "onScreenTurningOff()");
+            mKeyguardService.onScreenTurningOff();
+        }
+        mKeyguardState.screenState = SCREEN_STATE_TURNING_OFF;
+    }
+
     public void onScreenTurnedOff() {
         if (mKeyguardService != null) {
             if (DEBUG) Log.v(TAG, "onScreenTurnedOff()");
diff --git a/services/core/java/com/android/server/policy/keyguard/KeyguardServiceWrapper.java b/services/core/java/com/android/server/policy/keyguard/KeyguardServiceWrapper.java
index 425be54..952e0b0 100644
--- a/services/core/java/com/android/server/policy/keyguard/KeyguardServiceWrapper.java
+++ b/services/core/java/com/android/server/policy/keyguard/KeyguardServiceWrapper.java
@@ -128,6 +128,15 @@
     }
 
     @Override
+    public void onFinishedWakingUp() {
+        try {
+            mService.onFinishedWakingUp();
+        } catch (RemoteException e) {
+            Slog.w(TAG , "Remote Exception", e);
+        }
+    }
+
+    @Override
     public void onScreenTurningOn(IKeyguardDrawnCallback callback) {
         try {
             mService.onScreenTurningOn(callback);
@@ -146,6 +155,15 @@
     }
 
     @Override
+    public void onScreenTurningOff() {
+        try {
+            mService.onScreenTurningOff();
+        } catch (RemoteException e) {
+            Slog.w(TAG , "Remote Exception", e);
+        }
+    }
+
+    @Override
     public void onScreenTurnedOff() {
         try {
             mService.onScreenTurnedOff();