Reset keyguard drawn tasks in onScreenTurnedOff

Resets keyguard drawn tasks in onScreenTurnedOff
callback instead of onScreenTurnedOn callback.

This is necessary to avoid clearing of the tasks
on the first boot when onScreenTurnedOn is called
immediately after onScreenTurningOn from the WM
side (without waiting for the SystemUI drawn
callback).

Bug: 268144550
Test: manual fold/unfold
Test: check boot time with/without fix
Test: atest com.android.keyguard.mediator.ScreenOnCoordinatorTest
Change-Id: I780e2738ed2fcf8ec013058a2212612fb6874288
(cherry picked from commit edc4cb8c7d93b2ede0590bcf881887c8ffdf5955)
Merged-In: I780e2738ed2fcf8ec013058a2212612fb6874288

3 files changed, 33 insertions, 0 deletions

diff --git a/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt b/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt
index 0f00a040b094..603471b1de41 100644
--- a/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt
+++ b/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt
@@ -73,6 +73,10 @@ class ScreenOnCoordinator @Inject constructor(
     @BinderThread
     fun onScreenTurnedOn() {
         foldAodAnimationController?.onScreenTurnedOn()
+    }
+
+    @BinderThread
+    fun onScreenTurnedOff() {
         pendingTasks.reset()
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
index 47872d2d68bb..4d40db0a0cfd 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
@@ -638,6 +638,7 @@ public class KeyguardService extends Service {
             checkPermission();
             mKeyguardViewMediator.onScreenTurnedOff();
             mKeyguardLifecyclesDispatcher.dispatch(KeyguardLifecyclesDispatcher.SCREEN_TURNED_OFF);
+            mScreenOnCoordinator.onScreenTurnedOff();
         }
 
         @Override // Binder interface
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/mediator/ScreenOnCoordinatorTest.kt b/packages/SystemUI/tests/src/com/android/keyguard/mediator/ScreenOnCoordinatorTest.kt
index e9a2789bb5c8..9fe32f1e378b 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/mediator/ScreenOnCoordinatorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/keyguard/mediator/ScreenOnCoordinatorTest.kt
@@ -32,6 +32,7 @@ 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
 import java.util.Optional
@@ -83,6 +84,33 @@ class ScreenOnCoordinatorTest : SysuiTestCase() {
     }
 
     @Test
+    fun testTasksReady_onScreenTurningOnAndTurnedOnEventsCalledTogether_callsDrawnCallback() {
+        screenOnCoordinator.onScreenTurningOn(runnable)
+        screenOnCoordinator.onScreenTurnedOn()
+
+        onUnfoldOverlayReady()
+        onFoldAodReady()
+        waitHandlerIdle(testHandler)
+
+        // Should be called when both unfold overlay and keyguard drawn ready
+        verify(runnable).run()
+    }
+
+    @Test
+    fun testTasksReady_onScreenTurnedOnAndTurnedOffBeforeCompletion_doesNotCallDrawnCallback() {
+        screenOnCoordinator.onScreenTurningOn(runnable)
+        screenOnCoordinator.onScreenTurnedOn()
+        screenOnCoordinator.onScreenTurnedOff()
+
+        onUnfoldOverlayReady()
+        onFoldAodReady()
+        waitHandlerIdle(testHandler)
+
+        // Should not be called because this screen turning on call is not valid anymore
+        verify(runnable, never()).run()
+    }
+
+    @Test
     fun testUnfoldTransitionDisabledDrawnTasksReady_onScreenTurningOn_callsDrawnCallback() {
         // Recreate with empty unfoldComponent
         screenOnCoordinator = ScreenOnCoordinator(