Rewire setDozing() to use repositories

Testing out the process of convert existing events, which is a mix of
both callbacks and direct calls to injected controllers, over to
flowable streams through the KeyguardRepository.

Test: atest DozeServiceHostTest KeyguardRepositoryImplTest
FoldAodAnimationControllerTest
Test: Manually fold/unfold the device
Bug: 242853098

Change-Id: I8864c413cd136e7cb831736331160bea04e98222

packages/SystemUI/src/com/android/systemui/doze/DozeHost.java

@@ -101,6 +101,11 @@ public interface DozeHost {
          * Called when the always on suppression state changes. See {@link #isAlwaysOnSuppressed()}.
          */
         default void onAlwaysOnSuppressedChanged(boolean suppressed) {}
+
+        /**
+         * Called when the dozing state may have been updated.
+         */
+        default void onDozingChanged(boolean isDozing) {}
     }
 
     interface PulseCallback {

packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt

@@ -20,6 +20,7 @@ import com.android.systemui.common.coroutine.ChannelExt.trySendWithFailureLoggin
 import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
 import com.android.systemui.common.shared.model.Position
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.doze.DozeHost
 import com.android.systemui.plugins.statusbar.StatusBarStateController
 import com.android.systemui.statusbar.policy.KeyguardStateController
 import javax.inject.Inject
@@ -28,6 +29,7 @@ import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.distinctUntilChanged
 
 /** Defines interface for classes that encapsulate application state for the keyguard. */
 interface KeyguardRepository {
@@ -102,6 +104,7 @@ class KeyguardRepositoryImpl
 constructor(
     statusBarStateController: StatusBarStateController,
     keyguardStateController: KeyguardStateController,
+    dozeHost: DozeHost,
 ) : KeyguardRepository {
     private val _animateBottomAreaDozingTransitions = MutableStateFlow(false)
     override val animateBottomAreaDozingTransitions =
@@ -136,19 +139,21 @@ constructor(
         awaitClose { keyguardStateController.removeCallback(callback) }
     }
 
-    override val isDozing: Flow<Boolean> = conflatedCallbackFlow {
-        val callback =
-            object : StatusBarStateController.StateListener {
-                override fun onDozingChanged(isDozing: Boolean) {
-                    trySendWithFailureLogging(isDozing, TAG, "updated isDozing")
-                }
+    override val isDozing: Flow<Boolean> =
+        conflatedCallbackFlow {
+                val callback =
+                    object : DozeHost.Callback {
+                        override fun onDozingChanged(isDozing: Boolean) {
+                            trySendWithFailureLogging(isDozing, TAG, "updated isDozing")
+                        }
+                    }
+                dozeHost.addCallback(callback)
+                trySendWithFailureLogging(false, TAG, "initial isDozing: false")
+
+                awaitClose { dozeHost.removeCallback(callback) }
             }
+            .distinctUntilChanged()
 
-        statusBarStateController.addCallback(callback)
-        trySendWithFailureLogging(statusBarStateController.isDozing, TAG, "initial isDozing")
-
-        awaitClose { statusBarStateController.removeCallback(callback) }
-    }
     override val dozeAmount: Flow<Float> = conflatedCallbackFlow {
         val callback =
             object : StatusBarStateController.StateListener {

packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeServiceHost.java

@@ -28,8 +28,6 @@ import android.util.Log;
 import android.view.MotionEvent;
 import android.view.View;
 
-import androidx.annotation.Nullable;
-
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.keyguard.KeyguardUpdateMonitor;
 import com.android.systemui.assist.AssistManager;
@@ -50,11 +48,9 @@ import com.android.systemui.statusbar.notification.NotificationWakeUpCoordinator
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.policy.BatteryController;
 import com.android.systemui.statusbar.policy.DeviceProvisionedController;
-import com.android.systemui.unfold.FoldAodAnimationController;
-import com.android.systemui.unfold.SysUIUnfoldComponent;
+import com.android.systemui.util.Assert;
 
 import java.util.ArrayList;
-import java.util.Optional;
 
 import javax.inject.Inject;
 
@@ -80,8 +76,6 @@ public final class DozeServiceHost implements DozeHost {
     private final WakefulnessLifecycle mWakefulnessLifecycle;
     private final SysuiStatusBarStateController mStatusBarStateController;
     private final DeviceProvisionedController mDeviceProvisionedController;
-    @Nullable
-    private final FoldAodAnimationController mFoldAodAnimationController;
     private final HeadsUpManagerPhone mHeadsUpManagerPhone;
     private final BatteryController mBatteryController;
     private final ScrimController mScrimController;
@@ -114,7 +108,6 @@ public final class DozeServiceHost implements DozeHost {
             Lazy<AssistManager> assistManagerLazy,
             DozeScrimController dozeScrimController, KeyguardUpdateMonitor keyguardUpdateMonitor,
             PulseExpansionHandler pulseExpansionHandler,
-            Optional<SysUIUnfoldComponent> sysUIUnfoldComponent,
             NotificationShadeWindowController notificationShadeWindowController,
             NotificationWakeUpCoordinator notificationWakeUpCoordinator,
             AuthController authController,
@@ -138,8 +131,6 @@ public final class DozeServiceHost implements DozeHost {
         mNotificationWakeUpCoordinator = notificationWakeUpCoordinator;
         mAuthController = authController;
         mNotificationIconAreaController = notificationIconAreaController;
-        mFoldAodAnimationController = sysUIUnfoldComponent
-                .map(SysUIUnfoldComponent::getFoldAodAnimationController).orElse(null);
     }
 
     // TODO: we should try to not pass status bar in here if we can avoid it.
@@ -167,6 +158,7 @@ public final class DozeServiceHost implements DozeHost {
     }
 
     void firePowerSaveChanged(boolean active) {
+        Assert.isMainThread();
         for (Callback callback : mCallbacks) {
             callback.onPowerSaveChanged(active);
         }
@@ -177,6 +169,7 @@ public final class DozeServiceHost implements DozeHost {
             entry.setPulseSuppressed(true);
             mNotificationIconAreaController.updateAodNotificationIcons();
         };
+        Assert.isMainThread();
         for (Callback callback : mCallbacks) {
             callback.onNotificationAlerted(pulseSuppressedListener);
         }
@@ -193,11 +186,13 @@ public final class DozeServiceHost implements DozeHost {
 
     @Override
     public void addCallback(@NonNull Callback callback) {
+        Assert.isMainThread();
         mCallbacks.add(callback);
     }
 
     @Override
     public void removeCallback(@NonNull Callback callback) {
+        Assert.isMainThread();
         mCallbacks.remove(callback);
     }
 
@@ -212,6 +207,8 @@ public final class DozeServiceHost implements DozeHost {
     }
 
     void updateDozing() {
+        Assert.isMainThread();
+
         // When in wake-and-unlock while pulsing, keep dozing state until fully unlocked.
         boolean
                 dozing =
@@ -225,10 +222,10 @@ public final class DozeServiceHost implements DozeHost {
             dozing = false;
         }
 
-        mStatusBarStateController.setIsDozing(dozing);
-        if (mFoldAodAnimationController != null) {
-            mFoldAodAnimationController.setIsDozing(dozing);
+        for (Callback callback : mCallbacks) {
+            callback.onDozingChanged(dozing);
         }
+        mStatusBarStateController.setIsDozing(dozing);
     }
 
     @Override
@@ -452,6 +449,7 @@ public final class DozeServiceHost implements DozeHost {
             return;
         }
         mAlwaysOnSuppressed = suppressed;
+        Assert.isMainThread();
         for (Callback callback : mCallbacks) {
             callback.onAlwaysOnSuppressedChanged(suppressed);
         }

packages/SystemUI/src/com/android/systemui/unfold/FoldAodAnimationController.kt

@@ -18,22 +18,31 @@ package com.android.systemui.unfold
 
 import android.content.Context
 import android.hardware.devicestate.DeviceStateManager
-import android.os.Handler
 import android.os.PowerManager
 import android.provider.Settings
+import androidx.annotation.VisibleForTesting
 import androidx.core.view.OneShotPreDrawListener
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.repeatOnLifecycle
 import com.android.internal.util.LatencyTracker
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.keyguard.WakefulnessLifecycle
+import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
+import com.android.systemui.lifecycle.repeatWhenAttached
 import com.android.systemui.statusbar.LightRevealScrim
 import com.android.systemui.statusbar.phone.CentralSurfaces
 import com.android.systemui.statusbar.phone.ScreenOffAnimation
 import com.android.systemui.statusbar.policy.CallbackController
 import com.android.systemui.unfold.FoldAodAnimationController.FoldAodAnimationStatus
+import com.android.systemui.util.concurrency.DelayableExecutor
 import com.android.systemui.util.settings.GlobalSettings
-import java.util.concurrent.Executor
+import dagger.Lazy
 import java.util.function.Consumer
 import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.launch
 
 /**
  * Controls folding to AOD animation: when AOD is enabled and foldable device is folded we play a
@@ -43,16 +52,16 @@ import javax.inject.Inject
 class FoldAodAnimationController
 @Inject
 constructor(
-    @Main private val handler: Handler,
-    @Main private val executor: Executor,
+    @Main private val executor: DelayableExecutor,
     private val context: Context,
     private val deviceStateManager: DeviceStateManager,
     private val wakefulnessLifecycle: WakefulnessLifecycle,
     private val globalSettings: GlobalSettings,
     private val latencyTracker: LatencyTracker,
+    private val keyguardInteractor: Lazy<KeyguardInteractor>,
 ) : CallbackController<FoldAodAnimationStatus>, ScreenOffAnimation, WakefulnessLifecycle.Observer {
 
-    private lateinit var mCentralSurfaces: CentralSurfaces
+    private lateinit var centralSurfaces: CentralSurfaces
 
     private var isFolded = false
     private var isFoldHandled = true
@@ -64,12 +73,13 @@ constructor(
 
     private var shouldPlayAnimation = false
     private var isAnimationPlaying = false
+    private var cancelAnimation: Runnable? = null
 
     private val statusListeners = arrayListOf<FoldAodAnimationStatus>()
     private val foldToAodLatencyTracker = FoldToAodLatencyTracker()
 
     private val startAnimationRunnable = Runnable {
-        mCentralSurfaces.notificationPanelViewController.startFoldToAodAnimation(
+        centralSurfaces.notificationPanelViewController.startFoldToAodAnimation(
             /* startAction= */ { foldToAodLatencyTracker.onAnimationStarted() },
             /* endAction= */ { setAnimationState(playing = false) },
             /* cancelAction= */ { setAnimationState(playing = false) },
@@ -77,10 +87,14 @@ constructor(
     }
 
     override fun initialize(centralSurfaces: CentralSurfaces, lightRevealScrim: LightRevealScrim) {
-        this.mCentralSurfaces = centralSurfaces
+        this.centralSurfaces = centralSurfaces
 
         deviceStateManager.registerCallback(executor, FoldListener())
         wakefulnessLifecycle.addObserver(this)
+
+        centralSurfaces.notificationPanelViewController.view.repeatWhenAttached {
+            repeatOnLifecycle(Lifecycle.State.STARTED) { listenForDozing(this) }
+        }
     }
 
     /** Returns true if we should run fold to AOD animation */
@@ -94,7 +108,7 @@ constructor(
     override fun startAnimation(): Boolean =
         if (shouldStartAnimation()) {
             setAnimationState(playing = true)
-            mCentralSurfaces.notificationPanelViewController.prepareFoldToAodAnimation()
+            centralSurfaces.notificationPanelViewController.prepareFoldToAodAnimation()
             true
         } else {
             setAnimationState(playing = false)
@@ -104,8 +118,8 @@ constructor(
     override fun onStartedWakingUp() {
         if (isAnimationPlaying) {
             foldToAodLatencyTracker.cancel()
-            handler.removeCallbacks(startAnimationRunnable)
-            mCentralSurfaces.notificationPanelViewController.cancelFoldToAodAnimation()
+            cancelAnimation?.run()
+            centralSurfaces.notificationPanelViewController.cancelFoldToAodAnimation()
         }
 
         setAnimationState(playing = false)
@@ -138,13 +152,13 @@ constructor(
             // We should play the folding to AOD animation
 
             setAnimationState(playing = true)
-            mCentralSurfaces.notificationPanelViewController.prepareFoldToAodAnimation()
+            centralSurfaces.notificationPanelViewController.prepareFoldToAodAnimation()
 
             // We don't need to wait for the scrim as it is already displayed
             // but we should wait for the initial animation preparations to be drawn
             // (setting initial alpha/translation)
             OneShotPreDrawListener.add(
-                mCentralSurfaces.notificationPanelViewController.view,
+                centralSurfaces.notificationPanelViewController.view,
                 onReady
             )
         } else {
@@ -165,18 +179,14 @@ constructor(
 
     fun onScreenTurnedOn() {
         if (shouldPlayAnimation) {
-            handler.removeCallbacks(startAnimationRunnable)
+            cancelAnimation?.run()
 
             // Post starting the animation to the next frame to avoid junk due to inset changes
-            handler.post(startAnimationRunnable)
+            cancelAnimation = executor.executeDelayed(startAnimationRunnable, /* delayMillis= */ 0)
             shouldPlayAnimation = false
         }
     }
 
-    fun setIsDozing(dozing: Boolean) {
-        isDozing = dozing
-    }
-
     override fun isAnimationPlaying(): Boolean = isAnimationPlaying
 
     override fun isKeyguardHideDelayed(): Boolean = isAnimationPlaying()
@@ -204,6 +214,11 @@ constructor(
         statusListeners.remove(listener)
     }
 
+    @VisibleForTesting
+    internal suspend fun listenForDozing(scope: CoroutineScope): Job {
+        return scope.launch { keyguardInteractor.get().isDozing.collect { isDozing = it } }
+    }
+
     interface FoldAodAnimationStatus {
         fun onFoldToAodAnimationChanged()
     }

packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardRepositoryImplTest.kt

@@ -19,6 +19,7 @@ package com.android.systemui.keyguard.data.repository
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.common.shared.model.Position
+import com.android.systemui.doze.DozeHost
 import com.android.systemui.plugins.statusbar.StatusBarStateController
 import com.android.systemui.statusbar.policy.KeyguardStateController
 import com.android.systemui.util.mockito.argumentCaptor
@@ -40,6 +41,7 @@ import org.mockito.MockitoAnnotations
 class KeyguardRepositoryImplTest : SysuiTestCase() {
 
     @Mock private lateinit var statusBarStateController: StatusBarStateController
+    @Mock private lateinit var dozeHost: DozeHost
     @Mock private lateinit var keyguardStateController: KeyguardStateController
 
     private lateinit var underTest: KeyguardRepositoryImpl
@@ -48,7 +50,12 @@ class KeyguardRepositoryImplTest : SysuiTestCase() {
     fun setUp() {
         MockitoAnnotations.initMocks(this)
 
-        underTest = KeyguardRepositoryImpl(statusBarStateController, keyguardStateController)
+        underTest =
+            KeyguardRepositoryImpl(
+                statusBarStateController,
+                keyguardStateController,
+                dozeHost,
+            )
     }
 
     @Test
@@ -129,8 +136,8 @@ class KeyguardRepositoryImplTest : SysuiTestCase() {
         var latest: Boolean? = null
         val job = underTest.isDozing.onEach { latest = it }.launchIn(this)
 
-        val captor = argumentCaptor<StatusBarStateController.StateListener>()
-        verify(statusBarStateController).addCallback(captor.capture())
+        val captor = argumentCaptor<DozeHost.Callback>()
+        verify(dozeHost).addCallback(captor.capture())
 
         captor.value.onDozingChanged(true)
         assertThat(latest).isTrue()
@@ -139,7 +146,7 @@ class KeyguardRepositoryImplTest : SysuiTestCase() {
         assertThat(latest).isFalse()
 
         job.cancel()
-        verify(statusBarStateController).removeCallback(captor.value)
+        verify(dozeHost).removeCallback(captor.value)
     }
 
     @Test

packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeServiceHostTest.java

@@ -29,6 +29,7 @@ import static org.mockito.Mockito.when;
 
 import android.os.PowerManager;
 import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper.RunWithLooper;
 import android.view.View;
 
 import androidx.test.filters.SmallTest;
@@ -61,10 +62,10 @@ import org.mockito.MockitoAnnotations;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashSet;
-import java.util.Optional;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
+@RunWithLooper(setAsMainLooper = true)
 public class DozeServiceHostTest extends SysuiTestCase {
 
     private DozeServiceHost mDozeServiceHost;
@@ -92,6 +93,7 @@ public class DozeServiceHostTest extends SysuiTestCase {
     @Mock private View mAmbientIndicationContainer;
     @Mock private BiometricUnlockController mBiometricUnlockController;
     @Mock private AuthController mAuthController;
+    @Mock private DozeHost.Callback mCallback;
 
     @Before
     public void setup() {
@@ -100,7 +102,7 @@ public class DozeServiceHostTest extends SysuiTestCase {
                 mStatusBarStateController, mDeviceProvisionedController, mHeadsUpManager,
                 mBatteryController, mScrimController, () -> mBiometricUnlockController,
                 mKeyguardViewMediator, () -> mAssistManager, mDozeScrimController,
-                mKeyguardUpdateMonitor, mPulseExpansionHandler, Optional.empty(),
+                mKeyguardUpdateMonitor, mPulseExpansionHandler,
                 mNotificationShadeWindowController, mNotificationWakeUpCoordinator,
                 mAuthController, mNotificationIconAreaController);
 
@@ -114,16 +116,19 @@ public class DozeServiceHostTest extends SysuiTestCase {
 
     @Test
     public void testStartStopDozing() {
+        mDozeServiceHost.addCallback(mCallback);
         when(mStatusBarStateController.getState()).thenReturn(StatusBarState.KEYGUARD);
         when(mStatusBarStateController.isKeyguardRequested()).thenReturn(true);
 
         assertFalse(mDozeServiceHost.getDozingRequested());
 
         mDozeServiceHost.startDozing();
+        verify(mCallback).onDozingChanged(eq(true));
         verify(mStatusBarStateController).setIsDozing(eq(true));
         verify(mCentralSurfaces).updateIsKeyguard();
 
         mDozeServiceHost.stopDozing();
+        verify(mCallback).onDozingChanged(eq(false));
         verify(mStatusBarStateController).setIsDozing(eq(false));
     }
 

packages/SystemUI/tests/src/com/android/systemui/unfold/FoldAodAnimationControllerTest.kt

@@ -18,24 +18,28 @@ package com.android.systemui.unfold
 
 import android.hardware.devicestate.DeviceStateManager
 import android.hardware.devicestate.DeviceStateManager.FoldStateListener
-import android.os.Handler
 import android.os.PowerManager
 import android.testing.AndroidTestingRunner
-import android.testing.TestableLooper
-import android.testing.TestableLooper.RunWithLooper
 import android.view.ViewGroup
 import android.view.ViewTreeObserver
 import androidx.test.filters.SmallTest
 import com.android.internal.util.LatencyTracker
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.keyguard.WakefulnessLifecycle
+import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
+import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
 import com.android.systemui.shade.NotificationPanelViewController
 import com.android.systemui.statusbar.LightRevealScrim
 import com.android.systemui.statusbar.phone.CentralSurfaces
 import com.android.systemui.unfold.util.FoldableDeviceStates
 import com.android.systemui.unfold.util.FoldableTestUtils
+import com.android.systemui.util.concurrency.FakeExecutor
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.settings.GlobalSettings
+import com.android.systemui.util.time.FakeSystemClock
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.yield
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -49,7 +53,6 @@ import org.mockito.MockitoAnnotations
 
 @RunWith(AndroidTestingRunner::class)
 @SmallTest
-@RunWithLooper
 class FoldAodAnimationControllerTest : SysuiTestCase() {
 
     @Mock lateinit var deviceStateManager: DeviceStateManager
@@ -74,26 +77,15 @@ class FoldAodAnimationControllerTest : SysuiTestCase() {
 
     private lateinit var deviceStates: FoldableDeviceStates
 
-    private lateinit var testableLooper: TestableLooper
+    lateinit var keyguardRepository: FakeKeyguardRepository
 
-    lateinit var foldAodAnimationController: FoldAodAnimationController
+    lateinit var underTest: FoldAodAnimationController
+    private val fakeExecutor = FakeExecutor(FakeSystemClock())
 
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
-        testableLooper = TestableLooper.get(this)
 
-        foldAodAnimationController =
-            FoldAodAnimationController(
-                    Handler(testableLooper.looper),
-                    context.mainExecutor,
-                    context,
-                    deviceStateManager,
-                    wakefulnessLifecycle,
-                    globalSettings,
-                    latencyTracker,
-                )
-                .apply { initialize(centralSurfaces, lightRevealScrim) }
         deviceStates = FoldableTestUtils.findDeviceStates(context)
 
         whenever(notificationPanelViewController.view).thenReturn(viewGroup)
@@ -107,60 +99,102 @@ class FoldAodAnimationControllerTest : SysuiTestCase() {
                 val onActionStarted = it.arguments[0] as Runnable
                 onActionStarted.run()
             }
-        verify(deviceStateManager).registerCallback(any(), foldStateListenerCaptor.capture())
 
-        foldAodAnimationController.setIsDozing(dozing = true)
-        setAodEnabled(enabled = true)
-        sendFoldEvent(folded = false)
+        keyguardRepository = FakeKeyguardRepository()
+        val keyguardInteractor = KeyguardInteractor(repository = keyguardRepository)
+
+        // Needs to be run on the main thread
+        runBlocking(IMMEDIATE) {
+            underTest =
+                FoldAodAnimationController(
+                        fakeExecutor,
+                        context,
+                        deviceStateManager,
+                        wakefulnessLifecycle,
+                        globalSettings,
+                        latencyTracker,
+                        { keyguardInteractor },
+                    )
+                    .apply { initialize(centralSurfaces, lightRevealScrim) }
+
+            verify(deviceStateManager).registerCallback(any(), foldStateListenerCaptor.capture())
+
+            setAodEnabled(enabled = true)
+            sendFoldEvent(folded = false)
+        }
     }
 
     @Test
-    fun onFolded_aodDisabled_doesNotLogLatency() {
-        setAodEnabled(enabled = false)
+    fun onFolded_aodDisabled_doesNotLogLatency() =
+        runBlocking(IMMEDIATE) {
+            val job = underTest.listenForDozing(this)
+            keyguardRepository.setDozing(true)
+            setAodEnabled(enabled = false)
 
-        fold()
-        simulateScreenTurningOn()
+            yield()
 
-        verifyNoMoreInteractions(latencyTracker)
-    }
+            fold()
+            simulateScreenTurningOn()
+
+            verifyNoMoreInteractions(latencyTracker)
+
+            job.cancel()
+        }
 
     @Test
-    fun onFolded_aodEnabled_logsLatency() {
-        setAodEnabled(enabled = true)
+    fun onFolded_aodEnabled_logsLatency() =
+        runBlocking(IMMEDIATE) {
+            val job = underTest.listenForDozing(this)
+            keyguardRepository.setDozing(true)
+            setAodEnabled(enabled = true)
 
-        fold()
-        simulateScreenTurningOn()
+            yield()
 
-        verify(latencyTracker).onActionStart(any())
-        verify(latencyTracker).onActionEnd(any())
-    }
+            fold()
+            simulateScreenTurningOn()
+
+            verify(latencyTracker).onActionStart(any())
+            verify(latencyTracker).onActionEnd(any())
+
+            job.cancel()
+        }
 
     @Test
-    fun onFolded_animationCancelled_doesNotLogLatency() {
-        setAodEnabled(enabled = true)
+    fun onFolded_animationCancelled_doesNotLogLatency() =
+        runBlocking(IMMEDIATE) {
+            val job = underTest.listenForDozing(this)
+            keyguardRepository.setDozing(true)
+            setAodEnabled(enabled = true)
 
-        fold()
-        foldAodAnimationController.onScreenTurningOn({})
-        foldAodAnimationController.onStartedWakingUp()
-        testableLooper.processAllMessages()
+            yield()
 
-        verify(latencyTracker).onActionStart(any())
-        verify(latencyTracker).onActionCancel(any())
-    }
+            fold()
+            underTest.onScreenTurningOn({})
+            underTest.onStartedWakingUp()
+            fakeExecutor.runAllReady()
+
+            verify(latencyTracker).onActionStart(any())
+            verify(latencyTracker).onActionCancel(any())
+
+            job.cancel()
+        }
 
     private fun simulateScreenTurningOn() {
-        foldAodAnimationController.onScreenTurningOn({})
-        foldAodAnimationController.onScreenTurnedOn()
-        testableLooper.processAllMessages()
+        underTest.onScreenTurningOn({})
+        underTest.onScreenTurnedOn()
+        fakeExecutor.runAllReady()
     }
 
     private fun fold() = sendFoldEvent(folded = true)
 
-    private fun setAodEnabled(enabled: Boolean) =
-        foldAodAnimationController.onAlwaysOnChanged(alwaysOn = enabled)
+    private fun setAodEnabled(enabled: Boolean) = underTest.onAlwaysOnChanged(alwaysOn = enabled)
 
     private fun sendFoldEvent(folded: Boolean) {
         val state = if (folded) deviceStates.folded else deviceStates.unfolded
         foldStateListenerCaptor.value.onStateChanged(state)
     }
+
+    companion object {
+        private val IMMEDIATE = Dispatchers.Main.immediate
+    }
 }