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

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()
     }
 }

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

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
@@ -82,6 +83,33 @@ class ScreenOnCoordinatorTest : SysuiTestCase() {
         verify(runnable).run()
     }
 
+    @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