2 files changed, 22 insertions, 9 deletions

diff --git a/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt b/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt
index 98ac2c0bd026..0f00a040b094 100644
--- a/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt
+++ b/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt
@@ -17,8 +17,10 @@
 package com.android.keyguard.mediator
 
 import android.annotation.BinderThread
+import android.os.Handler
 import android.os.Trace
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.unfold.SysUIUnfoldComponent
 import com.android.systemui.util.concurrency.PendingTasksContainer
 import com.android.systemui.util.kotlin.getOrNull
@@ -33,7 +35,8 @@ import javax.inject.Inject
  */
 @SysUISingleton
 class ScreenOnCoordinator @Inject constructor(
-    unfoldComponent: Optional<SysUIUnfoldComponent>
+    unfoldComponent: Optional<SysUIUnfoldComponent>,
+    @Main private val mainHandler: Handler
 ) {
 
     private val unfoldLightRevealAnimation = unfoldComponent.map(
@@ -55,7 +58,11 @@ class ScreenOnCoordinator @Inject constructor(
         unfoldLightRevealAnimation?.onScreenTurningOn(pendingTasks.registerTask("unfold-reveal"))
         foldAodAnimationController?.onScreenTurningOn(pendingTasks.registerTask("fold-to-aod"))
 
-        pendingTasks.onTasksComplete { onDrawn.run() }
+        pendingTasks.onTasksComplete {
+            mainHandler.post {
+                onDrawn.run()
+            }
+        }
         Trace.endSection()
     }
 
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 34e78eb8c2eb..e9a2789bb5c8 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/mediator/ScreenOnCoordinatorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/keyguard/mediator/ScreenOnCoordinatorTest.kt
@@ -16,29 +16,25 @@
 
 package com.android.keyguard.mediator
 
+import android.os.Handler
+import android.os.Looper
 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.verify
 import org.mockito.MockitoAnnotations
+import java.util.Optional
 
 @SmallTest
 @RunWith(AndroidTestingRunner::class)
@@ -55,6 +51,8 @@ class ScreenOnCoordinatorTest : SysuiTestCase() {
     @Captor
     private lateinit var readyCaptor: ArgumentCaptor<Runnable>
 
+    private val testHandler = Handler(Looper.getMainLooper())
+
     private lateinit var screenOnCoordinator: ScreenOnCoordinator
 
     @Before
@@ -68,6 +66,7 @@ class ScreenOnCoordinatorTest : SysuiTestCase() {
 
         screenOnCoordinator = ScreenOnCoordinator(
             Optional.of(unfoldComponent),
+            testHandler
         )
     }
 
@@ -77,6 +76,7 @@ class ScreenOnCoordinatorTest : SysuiTestCase() {
 
         onUnfoldOverlayReady()
         onFoldAodReady()
+        waitHandlerIdle(testHandler)
 
         // Should be called when both unfold overlay and keyguard drawn ready
         verify(runnable).run()
@@ -87,8 +87,10 @@ class ScreenOnCoordinatorTest : SysuiTestCase() {
         // Recreate with empty unfoldComponent
         screenOnCoordinator = ScreenOnCoordinator(
             Optional.empty(),
+            testHandler
         )
         screenOnCoordinator.onScreenTurningOn(runnable)
+        waitHandlerIdle(testHandler)
 
         // Should be called when only keyguard drawn
         verify(runnable).run()
@@ -103,4 +105,8 @@ class ScreenOnCoordinatorTest : SysuiTestCase() {
         verify(foldAodAnimationController).onScreenTurningOn(capture(readyCaptor))
         readyCaptor.value.run()
     }
+
+    private fun waitHandlerIdle(handler: Handler) {
+        handler.runWithScissors({},  /* timeout= */ 0)
+    }
 }