Merge "KeyguardViewMediator refactor"
diff --git a/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt b/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt
new file mode 100644
index 0000000..705cf6d
--- /dev/null
+++ b/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt
@@ -0,0 +1,98 @@
+/*
+ * Copyright (C) 2022 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.keyguard.mediator
+
+import android.os.Trace
+
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.keyguard.ScreenLifecycle
+import com.android.systemui.util.concurrency.Execution
+import com.android.systemui.util.concurrency.PendingTasksContainer
+import com.android.systemui.unfold.SysUIUnfoldComponent
+import com.android.systemui.util.kotlin.getOrNull
+
+import java.util.Optional
+
+import javax.inject.Inject
+
+/**
+ * Coordinates screen on/turning on animations for the KeyguardViewMediator. Specifically for
+ * screen on events, this will invoke the onDrawn Runnable after all tasks have completed. This
+ * should route back to the KeyguardService, which informs the system_server that keyguard has
+ * drawn.
+ */
+@SysUISingleton
+class ScreenOnCoordinator @Inject constructor(
+    screenLifecycle: ScreenLifecycle,
+    unfoldComponent: Optional<SysUIUnfoldComponent>,
+    private val execution: Execution
+) : ScreenLifecycle.Observer {
+
+    private val unfoldLightRevealAnimation = unfoldComponent.map(
+        SysUIUnfoldComponent::getUnfoldLightRevealOverlayAnimation).getOrNull()
+    private val foldAodAnimationController = unfoldComponent.map(
+        SysUIUnfoldComponent::getFoldAodAnimationController).getOrNull()
+    private val pendingTasks = PendingTasksContainer()
+
+    private var wakeAndUnlockingTask: Runnable? = null
+    var wakeAndUnlocking = false
+        set(value) {
+            if (!value && field) {
+                // When updating the value back to false, mark the task complete in order to
+                // callback onDrawn
+                wakeAndUnlockingTask?.run()
+                wakeAndUnlockingTask = null
+            }
+            field = value
+        }
+
+    init {
+        screenLifecycle.addObserver(this)
+    }
+
+    /**
+     * When turning on, registers tasks that may need to run before invoking [onDrawn].
+     */
+    override fun onScreenTurningOn(onDrawn: Runnable) {
+        execution.assertIsMainThread()
+        Trace.beginSection("ScreenOnCoordinator#onScreenTurningOn")
+
+        pendingTasks.reset()
+
+        unfoldLightRevealAnimation?.onScreenTurningOn(pendingTasks.registerTask("unfold-reveal"))
+        foldAodAnimationController?.onScreenTurningOn(pendingTasks.registerTask("fold-to-aod"))
+
+        if (wakeAndUnlocking) {
+            wakeAndUnlockingTask = pendingTasks.registerTask("wake-and-unlocking")
+        }
+
+        pendingTasks.onTasksComplete { onDrawn.run() }
+        Trace.endSection()
+    }
+
+    override fun onScreenTurnedOn() {
+        execution.assertIsMainThread()
+
+        foldAodAnimationController?.onScreenTurnedOn()
+
+        pendingTasks.reset()
+    }
+
+    override fun onScreenTurnedOff() {
+        wakeAndUnlockingTask = null
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardLifecyclesDispatcher.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardLifecyclesDispatcher.java
index 5019e65..32b58c2 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardLifecyclesDispatcher.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardLifecyclesDispatcher.java
@@ -19,7 +19,10 @@
 import android.os.Handler;
 import android.os.Message;
 import android.os.PowerManager;
+import android.os.RemoteException;
+import android.util.Log;
 
+import com.android.internal.policy.IKeyguardDrawnCallback;
 import com.android.systemui.dagger.SysUISingleton;
 
 import javax.inject.Inject;
@@ -39,6 +42,7 @@
     static final int FINISHED_WAKING_UP = 5;
     static final int STARTED_GOING_TO_SLEEP = 6;
     static final int FINISHED_GOING_TO_SLEEP = 7;
+    private static final String TAG = "KeyguardLifecyclesDispatcher";
 
     private final ScreenLifecycle mScreenLifecycle;
     private final WakefulnessLifecycle mWakefulnessLifecycle;
@@ -65,12 +69,38 @@
         message.sendToTarget();
     }
 
+    /**
+     * @param what Message to send.
+     * @param object Object to send with the message
+     */
+    void dispatch(int what, Object object) {
+        mHandler.obtainMessage(what, object).sendToTarget();
+    }
+
     private Handler mHandler = new Handler() {
         @Override
         public void handleMessage(Message msg) {
+            final Object obj = msg.obj;
             switch (msg.what) {
                 case SCREEN_TURNING_ON:
-                    mScreenLifecycle.dispatchScreenTurningOn();
+                    // Ensure the drawn callback is only ever called once
+                    mScreenLifecycle.dispatchScreenTurningOn(new Runnable() {
+                            boolean mInvoked;
+                            @Override
+                            public void run() {
+                                if (obj == null) return;
+                                if (!mInvoked) {
+                                    mInvoked = true;
+                                    try {
+                                        ((IKeyguardDrawnCallback) obj).onDrawn();
+                                    } catch (RemoteException e) {
+                                        Log.w(TAG, "Exception calling onDrawn():", e);
+                                    }
+                                } else {
+                                    Log.w(TAG, "KeyguardDrawnCallback#onDrawn() invoked > 1 times");
+                                }
+                            }
+                        });
                     break;
                 case SCREEN_TURNED_ON:
                     mScreenLifecycle.dispatchScreenTurnedOn();
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
index fb31d88..e88011e 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
@@ -500,8 +500,8 @@
         public void onScreenTurningOn(IKeyguardDrawnCallback callback) {
             Trace.beginSection("KeyguardService.mBinder#onScreenTurningOn");
             checkPermission();
-            mKeyguardViewMediator.onScreenTurningOn(callback);
-            mKeyguardLifecyclesDispatcher.dispatch(KeyguardLifecyclesDispatcher.SCREEN_TURNING_ON);
+            mKeyguardLifecyclesDispatcher.dispatch(KeyguardLifecyclesDispatcher.SCREEN_TURNING_ON,
+                    callback);
             Trace.endSection();
         }
 
@@ -509,7 +509,6 @@
         public void onScreenTurnedOn() {
             Trace.beginSection("KeyguardService.mBinder#onScreenTurnedOn");
             checkPermission();
-            mKeyguardViewMediator.onScreenTurnedOn();
             mKeyguardLifecyclesDispatcher.dispatch(KeyguardLifecyclesDispatcher.SCREEN_TURNED_ON);
             Trace.endSection();
         }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
index 7bbd8ee..4658a74 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
@@ -87,7 +87,6 @@
 import com.android.internal.jank.InteractionJankMonitor;
 import com.android.internal.jank.InteractionJankMonitor.Configuration;
 import com.android.internal.policy.IKeyguardDismissCallback;
-import com.android.internal.policy.IKeyguardDrawnCallback;
 import com.android.internal.policy.IKeyguardExitCallback;
 import com.android.internal.policy.IKeyguardStateCallback;
 import com.android.internal.util.LatencyTracker;
@@ -99,6 +98,7 @@
 import com.android.keyguard.KeyguardUpdateMonitorCallback;
 import com.android.keyguard.KeyguardViewController;
 import com.android.keyguard.ViewMediatorCallback;
+import com.android.keyguard.mediator.ScreenOnCoordinator;
 import com.android.systemui.CoreStartable;
 import com.android.systemui.Dumpable;
 import com.android.systemui.animation.Interpolators;
@@ -122,15 +122,11 @@
 import com.android.systemui.statusbar.phone.panelstate.PanelExpansionStateManager;
 import com.android.systemui.statusbar.policy.KeyguardStateController;
 import com.android.systemui.statusbar.policy.UserSwitcherController;
-import com.android.systemui.unfold.FoldAodAnimationController;
-import com.android.systemui.unfold.SysUIUnfoldComponent;
-import com.android.systemui.unfold.UnfoldLightRevealOverlayAnimation;
 import com.android.systemui.util.DeviceConfigProxy;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.ArrayList;
-import java.util.Optional;
 import java.util.concurrent.Executor;
 
 import dagger.Lazy;
@@ -199,7 +195,6 @@
     private static final int RESET = 3;
     private static final int VERIFY_UNLOCK = 4;
     private static final int NOTIFY_FINISHED_GOING_TO_SLEEP = 5;
-    private static final int NOTIFY_SCREEN_TURNING_ON = 6;
     private static final int KEYGUARD_DONE = 7;
     private static final int KEYGUARD_DONE_DRAWING = 8;
     private static final int SET_OCCLUDED = 9;
@@ -208,8 +203,6 @@
     private static final int START_KEYGUARD_EXIT_ANIM = 12;
     private static final int KEYGUARD_DONE_PENDING_TIMEOUT = 13;
     private static final int NOTIFY_STARTED_WAKING_UP = 14;
-    private static final int NOTIFY_SCREEN_TURNED_ON = 15;
-    private static final int NOTIFY_SCREEN_TURNED_OFF = 16;
     private static final int NOTIFY_STARTED_GOING_TO_SLEEP = 17;
     private static final int SYSTEM_READY = 18;
     private static final int CANCEL_KEYGUARD_EXIT_ANIM = 19;
@@ -433,8 +426,6 @@
     private boolean mShowHomeOverLockscreen;
     private boolean mInGestureNavigationMode;
 
-    private boolean mWakeAndUnlocking;
-    private Runnable mWakeAndUnlockingDrawnCallback;
     private CharSequence mCustomMessage;
 
     /**
@@ -811,14 +802,11 @@
     private DeviceConfigProxy mDeviceConfig;
     private DozeParameters mDozeParameters;
 
-    private final Optional<UnfoldLightRevealOverlayAnimation> mUnfoldLightRevealAnimation;
-    private final Optional<FoldAodAnimationController> mFoldAodAnimationController;
-    private final PendingDrawnTasksContainer mPendingDrawnTasks = new PendingDrawnTasksContainer();
-
     private final KeyguardStateController mKeyguardStateController;
     private final Lazy<KeyguardUnlockAnimationController> mKeyguardUnlockAnimationControllerLazy;
     private final InteractionJankMonitor mInteractionJankMonitor;
     private boolean mWallpaperSupportsAmbientMode;
+    private ScreenOnCoordinator mScreenOnCoordinator;
 
     /**
      * Injected constructor. See {@link KeyguardModule}.
@@ -838,12 +826,12 @@
             NavigationModeController navigationModeController,
             KeyguardDisplayManager keyguardDisplayManager,
             DozeParameters dozeParameters,
-            Optional<SysUIUnfoldComponent> unfoldComponent,
             SysuiStatusBarStateController statusBarStateController,
             KeyguardStateController keyguardStateController,
             Lazy<KeyguardUnlockAnimationController> keyguardUnlockAnimationControllerLazy,
             ScreenOffAnimationController screenOffAnimationController,
             Lazy<NotificationShadeDepthController> notificationShadeDepthController,
+            ScreenOnCoordinator screenOnCoordinator,
             InteractionJankMonitor interactionJankMonitor) {
         super(context);
         mFalsingCollector = falsingCollector;
@@ -860,6 +848,7 @@
         mKeyguardDisplayManager = keyguardDisplayManager;
         dumpManager.registerDumpable(getClass().getName(), this);
         mDeviceConfig = deviceConfig;
+        mScreenOnCoordinator = screenOnCoordinator;
         mShowHomeOverLockscreen = mDeviceConfig.getBoolean(
                 DeviceConfig.NAMESPACE_SYSTEMUI,
                 NAV_BAR_HANDLE_SHOW_OVER_LOCKSCREEN,
@@ -874,11 +863,6 @@
                 }));
         mDozeParameters = dozeParameters;
 
-        mUnfoldLightRevealAnimation = unfoldComponent
-                .map(SysUIUnfoldComponent::getUnfoldLightRevealOverlayAnimation);
-        mFoldAodAnimationController = unfoldComponent
-                .map(SysUIUnfoldComponent::getFoldAodAnimationController);
-
         mStatusBarStateController = statusBarStateController;
         statusBarStateController.addCallback(this);
 
@@ -1071,7 +1055,7 @@
         synchronized (this) {
             mDeviceInteractive = false;
             mGoingToSleep = false;
-            mWakeAndUnlocking = false;
+            mScreenOnCoordinator.setWakeAndUnlocking(false);
             mAnimatingScreenOff = mDozeParameters.shouldAnimateDozingChange();
 
             resetKeyguardDonePendingLocked();
@@ -1249,20 +1233,7 @@
         Trace.endSection();
     }
 
-    public void onScreenTurningOn(IKeyguardDrawnCallback callback) {
-        Trace.beginSection("KeyguardViewMediator#onScreenTurningOn");
-        notifyScreenOn(callback);
-        Trace.endSection();
-    }
-
-    public void onScreenTurnedOn() {
-        Trace.beginSection("KeyguardViewMediator#onScreenTurnedOn");
-        notifyScreenTurnedOn();
-        Trace.endSection();
-    }
-
     public void onScreenTurnedOff() {
-        notifyScreenTurnedOff();
         mUpdateMonitor.dispatchScreenTurnedOff();
     }
 
@@ -1654,24 +1625,6 @@
         mHandler.sendEmptyMessage(NOTIFY_STARTED_WAKING_UP);
     }
 
-    private void notifyScreenOn(IKeyguardDrawnCallback callback) {
-        if (DEBUG) Log.d(TAG, "notifyScreenOn");
-        Message msg = mHandler.obtainMessage(NOTIFY_SCREEN_TURNING_ON, callback);
-        mHandler.sendMessage(msg);
-    }
-
-    private void notifyScreenTurnedOn() {
-        if (DEBUG) Log.d(TAG, "notifyScreenTurnedOn");
-        Message msg = mHandler.obtainMessage(NOTIFY_SCREEN_TURNED_ON);
-        mHandler.sendMessage(msg);
-    }
-
-    private void notifyScreenTurnedOff() {
-        if (DEBUG) Log.d(TAG, "notifyScreenTurnedOff");
-        Message msg = mHandler.obtainMessage(NOTIFY_SCREEN_TURNED_OFF);
-        mHandler.sendMessage(msg);
-    }
-
     /**
      * Send message to keyguard telling it to show itself
      * @see #handleShow
@@ -1832,21 +1785,6 @@
                 case NOTIFY_FINISHED_GOING_TO_SLEEP:
                     handleNotifyFinishedGoingToSleep();
                     break;
-                case NOTIFY_SCREEN_TURNING_ON:
-                    Trace.beginSection(
-                            "KeyguardViewMediator#handleMessage NOTIFY_SCREEN_TURNING_ON");
-                    handleNotifyScreenTurningOn((IKeyguardDrawnCallback) msg.obj);
-                    Trace.endSection();
-                    break;
-                case NOTIFY_SCREEN_TURNED_ON:
-                    Trace.beginSection(
-                            "KeyguardViewMediator#handleMessage NOTIFY_SCREEN_TURNED_ON");
-                    handleNotifyScreenTurnedOn();
-                    Trace.endSection();
-                    break;
-                case NOTIFY_SCREEN_TURNED_OFF:
-                    handleNotifyScreenTurnedOff();
-                    break;
                 case NOTIFY_STARTED_WAKING_UP:
                     Trace.beginSection(
                             "KeyguardViewMediator#handleMessage NOTIFY_STARTED_WAKING_UP");
@@ -2075,7 +2013,7 @@
 
             mHiding = false;
             mKeyguardExitAnimationRunner = null;
-            mWakeAndUnlocking = false;
+            mScreenOnCoordinator.setWakeAndUnlocking(false);
             mPendingLock = false;
             setShowingLocked(true);
             mKeyguardViewControllerLazy.get().show(options);
@@ -2107,12 +2045,14 @@
 
             int flags = 0;
             if (mKeyguardViewControllerLazy.get().shouldDisableWindowAnimationsForUnlock()
-                    || mWakeAndUnlocking && !mWallpaperSupportsAmbientMode) {
+                    || mScreenOnCoordinator.getWakeAndUnlocking()
+                            && !mWallpaperSupportsAmbientMode) {
                 flags |= WindowManagerPolicyConstants
                         .KEYGUARD_GOING_AWAY_FLAG_NO_WINDOW_ANIMATIONS;
             }
             if (mKeyguardViewControllerLazy.get().isGoingToNotificationShade()
-                    || mWakeAndUnlocking && mWallpaperSupportsAmbientMode) {
+                    || mScreenOnCoordinator.getWakeAndUnlocking()
+                            && mWallpaperSupportsAmbientMode) {
                 // When the wallpaper supports ambient mode, the scrim isn't fully opaque during
                 // wake and unlock and we should fade in the app on top of the wallpaper
                 flags |= WindowManagerPolicyConstants.KEYGUARD_GOING_AWAY_FLAG_TO_SHADE;
@@ -2223,14 +2163,13 @@
             IRemoteAnimationRunner runner = mKeyguardExitAnimationRunner;
             mKeyguardExitAnimationRunner = null;
 
-            if (mWakeAndUnlocking && mWakeAndUnlockingDrawnCallback != null) {
+            if (mScreenOnCoordinator.getWakeAndUnlocking()) {
 
                 // Hack level over 9000: To speed up wake-and-unlock sequence, force it to report
                 // the next draw from here so we don't have to wait for window manager to signal
                 // this to our ViewRootImpl.
                 mKeyguardViewControllerLazy.get().getViewRootImpl().setReportNextDraw();
-                mWakeAndUnlockingDrawnCallback.run();
-                mWakeAndUnlockingDrawnCallback = null;
+                mScreenOnCoordinator.setWakeAndUnlocking(false);
             }
 
             LatencyTracker.getInstance(mContext)
@@ -2354,7 +2293,7 @@
         }
 
         setShowingLocked(false);
-        mWakeAndUnlocking = false;
+        mScreenOnCoordinator.setWakeAndUnlocking(false);
         mDismissCallbackRegistry.notifyDismissSucceeded();
         resetKeyguardDonePendingLocked();
         mHideAnimationRun = false;
@@ -2561,64 +2500,6 @@
         Trace.endSection();
     }
 
-    private void handleNotifyScreenTurningOn(IKeyguardDrawnCallback callback) {
-        Trace.beginSection("KeyguardViewMediator#handleNotifyScreenTurningOn");
-        synchronized (KeyguardViewMediator.this) {
-            if (DEBUG) Log.d(TAG, "handleNotifyScreenTurningOn");
-
-            mPendingDrawnTasks.reset();
-
-            if (mUnfoldLightRevealAnimation.isPresent()) {
-                mUnfoldLightRevealAnimation.get()
-                        .onScreenTurningOn(mPendingDrawnTasks.registerTask("unfold-reveal"));
-            }
-
-            if (mFoldAodAnimationController.isPresent()) {
-                mFoldAodAnimationController.get()
-                        .onScreenTurningOn(mPendingDrawnTasks.registerTask("fold-to-aod"));
-            }
-
-            if (callback != null) {
-                if (mWakeAndUnlocking) {
-                    mWakeAndUnlockingDrawnCallback =
-                            mPendingDrawnTasks.registerTask("wake-and-unlocking");
-                }
-            }
-
-            mPendingDrawnTasks.onTasksComplete(() -> notifyDrawn(callback));
-        }
-        Trace.endSection();
-    }
-
-    private void handleNotifyScreenTurnedOn() {
-        Trace.beginSection("KeyguardViewMediator#handleNotifyScreenTurnedOn");
-        synchronized (this) {
-            if (DEBUG) Log.d(TAG, "handleNotifyScreenTurnedOn");
-
-            mPendingDrawnTasks.reset();
-        }
-        Trace.endSection();
-    }
-
-    private void handleNotifyScreenTurnedOff() {
-        synchronized (this) {
-            if (DEBUG) Log.d(TAG, "handleNotifyScreenTurnedOff");
-            mWakeAndUnlockingDrawnCallback = null;
-        }
-    }
-
-    private void notifyDrawn(final IKeyguardDrawnCallback callback) {
-        Trace.beginSection("KeyguardViewMediator#notifyDrawn");
-        try {
-            if (callback != null) {
-                callback.onDrawn();
-            }
-        } catch (RemoteException e) {
-            Slog.w(TAG, "Exception calling onDrawn():", e);
-        }
-        Trace.endSection();
-    }
-
     private void resetKeyguardDonePendingLocked() {
         mKeyguardDonePending = false;
         mHandler.removeMessages(KEYGUARD_DONE_PENDING_TIMEOUT);
@@ -2642,7 +2523,7 @@
 
     public void onWakeAndUnlocking() {
         Trace.beginSection("KeyguardViewMediator#onWakeAndUnlocking");
-        mWakeAndUnlocking = true;
+        mScreenOnCoordinator.setWakeAndUnlocking(true);
         keyguardDone();
         Trace.endSection();
     }
@@ -2777,9 +2658,7 @@
         pw.print("  mHideAnimationRun: "); pw.println(mHideAnimationRun);
         pw.print("  mPendingReset: "); pw.println(mPendingReset);
         pw.print("  mPendingLock: "); pw.println(mPendingLock);
-        pw.print("  mPendingDrawnTasks: "); pw.println(mPendingDrawnTasks.getPendingCount());
-        pw.print("  mWakeAndUnlocking: "); pw.println(mWakeAndUnlocking);
-        pw.print("  mDrawnCallback: "); pw.println(mWakeAndUnlockingDrawnCallback);
+        pw.print("  wakeAndUnlocking: "); pw.println(mScreenOnCoordinator.getWakeAndUnlocking());
     }
 
     /**
@@ -2853,7 +2732,7 @@
     }
 
     private void setShowingLocked(boolean showing, boolean forceCallbacks) {
-        final boolean aodShowing = mDozing && !mWakeAndUnlocking;
+        final boolean aodShowing = mDozing && !mScreenOnCoordinator.getWakeAndUnlocking();
         final boolean notifyDefaultDisplayCallbacks = showing != mShowing
                 || aodShowing != mAodShowing || forceCallbacks;
         mShowing = showing;
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/Lifecycle.java b/packages/SystemUI/src/com/android/systemui/keyguard/Lifecycle.java
index 3da6caf..b870f58 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/Lifecycle.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/Lifecycle.java
@@ -20,6 +20,7 @@
 
 import java.util.ArrayList;
 import java.util.Objects;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 
 /**
@@ -42,4 +43,13 @@
             consumer.accept(mObservers.get(i));
         }
     }
+
+    /**
+     * Will dispatch the consumer to the observer, along with a single argument of type<U>.
+     */
+    public <U> void dispatch(BiConsumer<T, U> biConsumer, U arg) {
+        for (int i = 0; i < mObservers.size(); i++) {
+            biConsumer.accept(mObservers.get(i), arg);
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ScreenLifecycle.java b/packages/SystemUI/src/com/android/systemui/keyguard/ScreenLifecycle.java
index d17c39a..e71aa85 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ScreenLifecycle.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ScreenLifecycle.java
@@ -18,6 +18,8 @@
 
 import android.os.Trace;
 
+import androidx.annotation.NonNull;
+
 import com.android.systemui.Dumpable;
 import com.android.systemui.dump.DumpManager;
 
@@ -49,9 +51,14 @@
         return mScreenState;
     }
 
-    public void dispatchScreenTurningOn() {
+    /**
+     * Dispatch screen turning on events to the registered observers
+     *
+     * @param onDrawn Invoke to notify the caller that the event has been processed
+     */
+    public void dispatchScreenTurningOn(@NonNull Runnable onDrawn) {
         setScreenState(SCREEN_TURNING_ON);
-        dispatch(Observer::onScreenTurningOn);
+        dispatch(Observer::onScreenTurningOn, onDrawn);
     }
 
     public void dispatchScreenTurnedOn() {
@@ -81,7 +88,12 @@
     }
 
     public interface Observer {
-        default void onScreenTurningOn() {}
+        /**
+         * Receive the screen turning on event
+         *
+         * @param onDrawn Invoke to notify the caller that the event has been processed
+         */
+        default void onScreenTurningOn(@NonNull Runnable onDrawn) {}
         default void onScreenTurnedOn() {}
         default void onScreenTurningOff() {}
         default void onScreenTurnedOff() {}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/dagger/KeyguardModule.java b/packages/SystemUI/src/com/android/systemui/keyguard/dagger/KeyguardModule.java
index eecb55b..dd844e8 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/dagger/KeyguardModule.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/dagger/KeyguardModule.java
@@ -31,6 +31,7 @@
 import com.android.keyguard.dagger.KeyguardStatusBarViewComponent;
 import com.android.keyguard.dagger.KeyguardStatusViewComponent;
 import com.android.keyguard.dagger.KeyguardUserSwitcherComponent;
+import com.android.keyguard.mediator.ScreenOnCoordinator;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.classifier.FalsingCollector;
 import com.android.systemui.classifier.FalsingModule;
@@ -50,11 +51,9 @@
 import com.android.systemui.statusbar.phone.StatusBar;
 import com.android.systemui.statusbar.policy.KeyguardStateController;
 import com.android.systemui.statusbar.policy.UserSwitcherController;
-import com.android.systemui.unfold.SysUIUnfoldComponent;
 import com.android.systemui.util.DeviceConfigProxy;
 import com.android.systemui.util.sensors.AsyncSensorManager;
 
-import java.util.Optional;
 import java.util.concurrent.Executor;
 
 import dagger.Lazy;
@@ -93,12 +92,12 @@
             NavigationModeController navigationModeController,
             KeyguardDisplayManager keyguardDisplayManager,
             DozeParameters dozeParameters,
-            Optional<SysUIUnfoldComponent> unfoldComponent,
             SysuiStatusBarStateController statusBarStateController,
             KeyguardStateController keyguardStateController,
             Lazy<KeyguardUnlockAnimationController> keyguardUnlockAnimationController,
             ScreenOffAnimationController screenOffAnimationController,
             Lazy<NotificationShadeDepthController> notificationShadeDepthController,
+            ScreenOnCoordinator screenOnCoordinator,
             InteractionJankMonitor interactionJankMonitor) {
         return new KeyguardViewMediator(
                 context,
@@ -117,12 +116,12 @@
                 navigationModeController,
                 keyguardDisplayManager,
                 dozeParameters,
-                unfoldComponent,
                 statusBarStateController,
                 keyguardStateController,
                 keyguardUnlockAnimationController,
                 screenOffAnimationController,
                 notificationShadeDepthController,
+                screenOnCoordinator,
                 interactionJankMonitor
         );
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
index ec95d57..6f8a334 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
@@ -3592,7 +3592,7 @@
 
     final ScreenLifecycle.Observer mScreenObserver = new ScreenLifecycle.Observer() {
         @Override
-        public void onScreenTurningOn() {
+        public void onScreenTurningOn(Runnable onDrawn) {
             mFalsingCollector.onScreenTurningOn();
             mNotificationPanelViewController.onScreenTurningOn();
         }
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/FoldAodAnimationController.kt b/packages/SystemUI/src/com/android/systemui/unfold/FoldAodAnimationController.kt
index fb9df01..52c416ba 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/FoldAodAnimationController.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/FoldAodAnimationController.kt
@@ -19,7 +19,6 @@
 import android.os.PowerManager
 import android.provider.Settings
 import com.android.systemui.keyguard.KeyguardViewMediator
-import com.android.systemui.keyguard.ScreenLifecycle
 import com.android.systemui.keyguard.WakefulnessLifecycle
 import com.android.systemui.statusbar.LightRevealScrim
 import com.android.systemui.statusbar.phone.ScreenOffAnimation
@@ -36,12 +35,10 @@
  */
 @SysUIUnfoldScope
 class FoldAodAnimationController @Inject constructor(
-    private val screenLifecycle: ScreenLifecycle,
     private val keyguardViewMediatorLazy: Lazy<KeyguardViewMediator>,
     private val wakefulnessLifecycle: WakefulnessLifecycle,
     private val globalSettings: GlobalSettings
-) : ScreenLifecycle.Observer,
-    CallbackController<FoldAodAnimationStatus>,
+) : CallbackController<FoldAodAnimationStatus>,
     ScreenOffAnimation,
     WakefulnessLifecycle.Observer {
 
@@ -58,7 +55,6 @@
     override fun initialize(statusBar: StatusBar, lightRevealScrim: LightRevealScrim) {
         this.statusBar = statusBar
 
-        screenLifecycle.addObserver(this)
         wakefulnessLifecycle.addObserver(this)
     }
 
@@ -123,7 +119,7 @@
         }
     }
 
-    override fun onScreenTurnedOn() {
+    fun onScreenTurnedOn() {
         if (shouldPlayAnimation) {
             statusBar.notificationPanelViewController.startFoldToAodAnimation {
                 // End action
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/PendingDrawnTasksContainer.kt b/packages/SystemUI/src/com/android/systemui/util/concurrency/PendingTasksContainer.kt
similarity index 77%
rename from packages/SystemUI/src/com/android/systemui/keyguard/PendingDrawnTasksContainer.kt
rename to packages/SystemUI/src/com/android/systemui/util/concurrency/PendingTasksContainer.kt
index a60033c..6cd384f 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/PendingDrawnTasksContainer.kt
+++ b/packages/SystemUI/src/com/android/systemui/util/concurrency/PendingTasksContainer.kt
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.keyguard
+package com.android.systemui.util.concurrency
 
 import android.os.Trace
 import java.util.concurrent.atomic.AtomicInteger
@@ -23,9 +23,9 @@
 /**
  * Allows to wait for multiple callbacks and notify when the last one is executed
  */
-class PendingDrawnTasksContainer {
+class PendingTasksContainer {
 
-    private var pendingDrawnTasksCount: AtomicInteger = AtomicInteger(0)
+    private var pendingTasksCount: AtomicInteger = AtomicInteger(0)
     private var completionCallback: AtomicReference<Runnable> = AtomicReference()
 
     /**
@@ -33,19 +33,19 @@
      * @return a runnable that should be invoked when the task is finished
      */
     fun registerTask(name: String): Runnable {
-        pendingDrawnTasksCount.incrementAndGet()
+        pendingTasksCount.incrementAndGet()
 
         if (ENABLE_TRACE) {
-            Trace.beginAsyncSection("PendingDrawnTasksContainer#$name", 0)
+            Trace.beginAsyncSection("PendingTasksContainer#$name", 0)
         }
 
         return Runnable {
-            if (pendingDrawnTasksCount.decrementAndGet() == 0) {
+            if (pendingTasksCount.decrementAndGet() == 0) {
                 val onComplete = completionCallback.getAndSet(null)
                 onComplete?.run()
 
                 if (ENABLE_TRACE) {
-                    Trace.endAsyncSection("PendingDrawnTasksContainer#$name", 0)
+                    Trace.endAsyncSection("PendingTasksContainer#$name", 0)
                 }
             }
         }
@@ -57,7 +57,7 @@
     fun reset() {
         // Create new objects in case if there are pending callbacks from the previous invocations
         completionCallback = AtomicReference()
-        pendingDrawnTasksCount = AtomicInteger(0)
+        pendingTasksCount = AtomicInteger(0)
     }
 
     /**
@@ -67,7 +67,7 @@
     fun onTasksComplete(onComplete: Runnable) {
         completionCallback.set(onComplete)
 
-        if (pendingDrawnTasksCount.get() == 0) {
+        if (pendingTasksCount.get() == 0) {
             val currentOnComplete = completionCallback.getAndSet(null)
             currentOnComplete?.run()
         }
@@ -76,7 +76,7 @@
     /**
      * Returns current pending tasks count
      */
-    fun getPendingCount(): Int = pendingDrawnTasksCount.get()
+    fun getPendingCount(): Int = pendingTasksCount.get()
 }
 
 private const val ENABLE_TRACE = false
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/mediator/ScreenOnCoordinatorTest.kt b/packages/SystemUI/tests/src/com/android/keyguard/mediator/ScreenOnCoordinatorTest.kt
new file mode 100644
index 0000000..96e6bd1
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/keyguard/mediator/ScreenOnCoordinatorTest.kt
@@ -0,0 +1,136 @@
+/*
+ * Copyright (C) 2021 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.keyguard.mediator
+
+import android.testing.AndroidTestingRunner
+import androidx.test.filters.SmallTest
+
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.keyguard.ScreenLifecycle
+import com.android.systemui.unfold.FoldAodAnimationController
+import com.android.systemui.unfold.SysUIUnfoldComponent
+import com.android.systemui.unfold.UnfoldLightRevealOverlayAnimation
+import com.android.systemui.util.concurrency.FakeExecution
+import com.android.systemui.util.mockito.capture
+
+import java.util.Optional
+
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+import org.mockito.ArgumentCaptor
+import org.mockito.Captor
+import org.mockito.Mock
+import org.mockito.Mockito.`when`
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+class ScreenOnCoordinatorTest : SysuiTestCase() {
+
+    @Mock
+    private lateinit var runnable: Runnable
+    @Mock
+    private lateinit var unfoldComponent: SysUIUnfoldComponent
+    @Mock
+    private lateinit var foldAodAnimationController: FoldAodAnimationController
+    @Mock
+    private lateinit var unfoldAnimation: UnfoldLightRevealOverlayAnimation
+    @Mock
+    private lateinit var screenLifecycle: ScreenLifecycle
+    @Captor
+    private lateinit var readyCaptor: ArgumentCaptor<Runnable>
+
+    private lateinit var screenOnCoordinator: ScreenOnCoordinator
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+
+        `when`(unfoldComponent.getUnfoldLightRevealOverlayAnimation())
+                .thenReturn(unfoldAnimation)
+        `when`(unfoldComponent.getFoldAodAnimationController())
+                .thenReturn(foldAodAnimationController)
+
+        screenOnCoordinator = ScreenOnCoordinator(
+            screenLifecycle,
+            Optional.of(unfoldComponent),
+            FakeExecution()
+        )
+
+        // Make sure screen events are registered to observe
+        verify(screenLifecycle).addObserver(screenOnCoordinator)
+    }
+
+    @Test
+    fun testUnfoldTransitionEnabledDrawnTasksReady_onScreenTurningOn_callsDrawnCallback() {
+        screenOnCoordinator.onScreenTurningOn(runnable)
+
+        onUnfoldOverlayReady()
+        onFoldAodReady()
+
+        // Should be called when both unfold overlay and keyguard drawn ready
+        verify(runnable).run()
+    }
+
+    @Test
+    fun testUnfoldTransitionDisabledDrawnTasksReady_onScreenTurningOn_callsDrawnCallback() {
+        // Recreate with empty unfoldComponent
+        screenOnCoordinator = ScreenOnCoordinator(
+            screenLifecycle,
+            Optional.empty(),
+            FakeExecution()
+        )
+        screenOnCoordinator.onScreenTurningOn(runnable)
+
+        // Should be called when only keyguard drawn
+        verify(runnable).run()
+    }
+
+    @Test
+    fun testWakeAndUnlockDelaysRunnable() {
+        // GIVEN wakeAndUnlocking has been set to true
+        screenOnCoordinator.wakeAndUnlocking = true
+
+        // WHEN the screen turns on and two tasks have completed
+        screenOnCoordinator.onScreenTurningOn(runnable)
+        onUnfoldOverlayReady()
+        onFoldAodReady()
+
+        // THEN the runnable should not have run yet
+        verify(runnable, never()).run()
+
+        // WHEN the value of wakeAndUnlocking changes
+        screenOnCoordinator.wakeAndUnlocking = false
+
+        // THEN the runnable should have run, as it is the last task to complete
+        verify(runnable).run()
+    }
+
+    private fun onUnfoldOverlayReady() {
+        verify(unfoldAnimation).onScreenTurningOn(capture(readyCaptor))
+        readyCaptor.getValue().run()
+    }
+
+    private fun onFoldAodReady() {
+        verify(foldAodAnimationController).onScreenTurningOn(capture(readyCaptor))
+        readyCaptor.getValue().run()
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
index 27fcb11..9827d21 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
@@ -34,7 +34,6 @@
 import android.app.trust.TrustManager;
 import android.os.PowerManager;
 import android.os.PowerManager.WakeLock;
-import android.os.RemoteException;
 import android.telephony.TelephonyManager;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
@@ -42,10 +41,10 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.internal.jank.InteractionJankMonitor;
-import com.android.internal.policy.IKeyguardDrawnCallback;
 import com.android.internal.widget.LockPatternUtils;
 import com.android.keyguard.KeyguardDisplayManager;
 import com.android.keyguard.KeyguardUpdateMonitor;
+import com.android.keyguard.mediator.ScreenOnCoordinator;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.classifier.FalsingCollectorFake;
@@ -58,9 +57,6 @@
 import com.android.systemui.statusbar.phone.StatusBarKeyguardViewManager;
 import com.android.systemui.statusbar.policy.KeyguardStateController;
 import com.android.systemui.statusbar.policy.UserSwitcherController;
-import com.android.systemui.unfold.FoldAodAnimationController;
-import com.android.systemui.unfold.SysUIUnfoldComponent;
-import com.android.systemui.unfold.UnfoldLightRevealOverlayAnimation;
 import com.android.systemui.util.DeviceConfigProxy;
 import com.android.systemui.util.DeviceConfigProxyFake;
 import com.android.systemui.util.concurrency.FakeExecutor;
@@ -69,13 +65,9 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.Optional;
-import java.util.function.Function;
-
 @RunWith(AndroidTestingRunner.class)
 @TestableLooper.RunWithLooper
 @SmallTest
@@ -95,36 +87,25 @@
     private @Mock NavigationModeController mNavigationModeController;
     private @Mock KeyguardDisplayManager mKeyguardDisplayManager;
     private @Mock DozeParameters mDozeParameters;
-    private @Mock SysUIUnfoldComponent mSysUIUnfoldComponent;
-    private @Mock UnfoldLightRevealOverlayAnimation mUnfoldAnimation;
     private @Mock SysuiStatusBarStateController mStatusBarStateController;
     private @Mock KeyguardStateController mKeyguardStateController;
     private @Mock NotificationShadeDepthController mNotificationShadeDepthController;
     private @Mock KeyguardUnlockAnimationController mKeyguardUnlockAnimationController;
     private @Mock ScreenOffAnimationController mScreenOffAnimationController;
-    private @Mock FoldAodAnimationController mFoldAodAnimationController;
-    private @Mock IKeyguardDrawnCallback mKeyguardDrawnCallback;
     private @Mock InteractionJankMonitor mInteractionJankMonitor;
+    private @Mock ScreenOnCoordinator mScreenOnCoordinator;
     private DeviceConfigProxy mDeviceConfig = new DeviceConfigProxyFake();
     private FakeExecutor mUiBgExecutor = new FakeExecutor(new FakeSystemClock());
 
-    private Optional<SysUIUnfoldComponent> mSysUiUnfoldComponentOptional;
-
     private FalsingCollectorFake mFalsingCollector;
 
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         mFalsingCollector = new FalsingCollectorFake();
-        mSysUiUnfoldComponentOptional = Optional.of(mSysUIUnfoldComponent);
 
         when(mLockPatternUtils.getDevicePolicyManager()).thenReturn(mDevicePolicyManager);
         when(mPowerManager.newWakeLock(anyInt(), any())).thenReturn(mock(WakeLock.class));
-        when(mSysUIUnfoldComponent.getUnfoldLightRevealOverlayAnimation())
-                .thenReturn(mUnfoldAnimation);
-        when(mSysUIUnfoldComponent.getFoldAodAnimationController())
-                .thenReturn(mFoldAodAnimationController);
-
         when(mInteractionJankMonitor.begin(any(), anyInt())).thenReturn(true);
         when(mInteractionJankMonitor.end(anyInt())).thenReturn(true);
 
@@ -151,33 +132,6 @@
     }
 
     @Test
-    @TestableLooper.RunWithLooper(setAsMainLooper = true)
-    public void testUnfoldTransitionEnabledDrawnTasksReady_onScreenTurningOn_callsDrawnCallback()
-            throws RemoteException {
-        mViewMediator.onScreenTurningOn(mKeyguardDrawnCallback);
-        TestableLooper.get(this).processAllMessages();
-        onUnfoldOverlayReady();
-        onFoldAodReady();
-
-        // Should be called when both unfold overlay and keyguard drawn ready
-        verify(mKeyguardDrawnCallback).onDrawn();
-    }
-
-    @Test
-    @TestableLooper.RunWithLooper(setAsMainLooper = true)
-    public void testUnfoldTransitionDisabledDrawnTasksReady_onScreenTurningOn_callsDrawnCallback()
-            throws RemoteException {
-        mSysUiUnfoldComponentOptional = Optional.empty();
-        createAndStartViewMediator();
-
-        mViewMediator.onScreenTurningOn(mKeyguardDrawnCallback);
-        TestableLooper.get(this).processAllMessages();
-
-        // Should be called when only keyguard drawn
-        verify(mKeyguardDrawnCallback).onDrawn();
-    }
-
-    @Test
     public void testIsAnimatingScreenOff() {
         when(mDozeParameters.shouldControlUnlockedScreenOff()).thenReturn(true);
         when(mDozeParameters.shouldAnimateDozingChange()).thenReturn(true);
@@ -219,20 +173,6 @@
         verify(mStatusBarKeyguardViewManager, atLeast(1)).show(null);
     }
 
-    private void onUnfoldOverlayReady() {
-        ArgumentCaptor<Runnable> overlayReadyCaptor = ArgumentCaptor.forClass(Runnable.class);
-        verify(mUnfoldAnimation).onScreenTurningOn(overlayReadyCaptor.capture());
-        overlayReadyCaptor.getValue().run();
-        TestableLooper.get(this).processAllMessages();
-    }
-
-    private void onFoldAodReady() {
-        ArgumentCaptor<Runnable> ready = ArgumentCaptor.forClass(Runnable.class);
-        verify(mFoldAodAnimationController).onScreenTurningOn(ready.capture());
-        ready.getValue().run();
-        TestableLooper.get(this).processAllMessages();
-    }
-
     private void createAndStartViewMediator() {
         mViewMediator = new KeyguardViewMediator(
                 mContext,
@@ -251,12 +191,12 @@
                 mNavigationModeController,
                 mKeyguardDisplayManager,
                 mDozeParameters,
-                mSysUiUnfoldComponentOptional,
                 mStatusBarStateController,
                 mKeyguardStateController,
                 () -> mKeyguardUnlockAnimationController,
                 mScreenOffAnimationController,
                 () -> mNotificationShadeDepthController,
+                mScreenOnCoordinator,
                 mInteractionJankMonitor);
         mViewMediator.start();
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ScreenLifecycleTest.java b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ScreenLifecycleTest.java
index 2d1b258..53d0dd8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ScreenLifecycleTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ScreenLifecycleTest.java
@@ -57,15 +57,16 @@
 
     @Test
     public void screenTurningOn() throws Exception {
-        mScreen.dispatchScreenTurningOn();
+        Runnable onDrawn = () -> {};
+        mScreen.dispatchScreenTurningOn(onDrawn);
 
         assertEquals(ScreenLifecycle.SCREEN_TURNING_ON, mScreen.getScreenState());
-        verify(mScreenObserverMock).onScreenTurningOn();
+        verify(mScreenObserverMock).onScreenTurningOn(onDrawn);
     }
 
     @Test
     public void screenTurnedOn() throws Exception {
-        mScreen.dispatchScreenTurningOn();
+        mScreen.dispatchScreenTurningOn(null);
         mScreen.dispatchScreenTurnedOn();
 
         assertEquals(ScreenLifecycle.SCREEN_ON, mScreen.getScreenState());
@@ -74,7 +75,7 @@
 
     @Test
     public void screenTurningOff() throws Exception {
-        mScreen.dispatchScreenTurningOn();
+        mScreen.dispatchScreenTurningOn(null);
         mScreen.dispatchScreenTurnedOn();
         mScreen.dispatchScreenTurningOff();
 
@@ -84,7 +85,7 @@
 
     @Test
     public void screenTurnedOff() throws Exception {
-        mScreen.dispatchScreenTurningOn();
+        mScreen.dispatchScreenTurningOn(null);
         mScreen.dispatchScreenTurnedOn();
         mScreen.dispatchScreenTurningOff();
         mScreen.dispatchScreenTurnedOff();
@@ -97,4 +98,4 @@
     public void dump() throws Exception {
         mScreen.dump(null, new PrintWriter(new ByteArrayOutputStream()), new String[0]);
     }
-}
\ No newline at end of file
+}