Merge "Transitions - Add doze transition information" into tm-qpr-dev

2022-11-18 12:25:23 +00:00
parent 7592051299 7dee55afa3
commit 2850e3301a
13 changed files with 527 additions and 231 deletions
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeTransitionListener.kt
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeTransitionListener.kt
@@ -0,0 +1,48 @@
+/*
+ * 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.systemui.doze
+
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.statusbar.policy.CallbackController
+import javax.inject.Inject
+
+/** Receives doze transition events, and passes those events to registered callbacks. */
+@SysUISingleton
+class DozeTransitionListener @Inject constructor() :
+    DozeMachine.Part, CallbackController<DozeTransitionCallback> {
+    val callbacks = mutableSetOf<DozeTransitionCallback>()
+    var oldState = DozeMachine.State.UNINITIALIZED
+    var newState = DozeMachine.State.UNINITIALIZED
+
+    override fun transitionTo(oldState: DozeMachine.State, newState: DozeMachine.State) {
+        this.oldState = oldState
+        this.newState = newState
+        callbacks.forEach { it.onDozeTransition(oldState, newState) }
+    }
+
+    override fun addCallback(callback: DozeTransitionCallback) {
+        callbacks.add(callback)
+    }
+
+    override fun removeCallback(callback: DozeTransitionCallback) {
+        callbacks.remove(callback)
+    }
+}
+
+interface DozeTransitionCallback {
+    fun onDozeTransition(oldState: DozeMachine.State, newState: DozeMachine.State)
+}
--- a/packages/SystemUI/src/com/android/systemui/doze/dagger/DozeModule.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/dagger/DozeModule.java
@@ -35,6 +35,7 @@ import com.android.systemui.doze.DozeScreenStatePreventingAdapter;
 import com.android.systemui.doze.DozeSensors;
 import com.android.systemui.doze.DozeSuppressor;
 import com.android.systemui.doze.DozeSuspendScreenStatePreventingAdapter;
+import com.android.systemui.doze.DozeTransitionListener;
 import com.android.systemui.doze.DozeTriggers;
 import com.android.systemui.doze.DozeUi;
 import com.android.systemui.doze.DozeWallpaperState;
@@ -83,7 +84,7 @@ public abstract class DozeModule {
            DozeUi dozeUi, DozeScreenState dozeScreenState,
            DozeScreenBrightness dozeScreenBrightness, DozeWallpaperState dozeWallpaperState,
            DozeDockHandler dozeDockHandler, DozeAuthRemover dozeAuthRemover,
-            DozeSuppressor dozeSuppressor) {
+            DozeSuppressor dozeSuppressor, DozeTransitionListener dozeTransitionListener) {
        return new DozeMachine.Part[]{
                dozePauser,
                dozeFalsingManagerAdapter,
@@ -94,7 +95,8 @@ public abstract class DozeModule {
                dozeWallpaperState,
                dozeDockHandler,
                dozeAuthRemover,
-                dozeSuppressor
+                dozeSuppressor,
+                dozeTransitionListener
        };
    }

--- a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
@@ -23,9 +23,14 @@ import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCall
 import com.android.systemui.common.shared.model.Position
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.doze.DozeHost
+import com.android.systemui.doze.DozeMachine
+import com.android.systemui.doze.DozeTransitionCallback
+import com.android.systemui.doze.DozeTransitionListener
 import com.android.systemui.keyguard.WakefulnessLifecycle
 import com.android.systemui.keyguard.WakefulnessLifecycle.Wakefulness
 import com.android.systemui.keyguard.shared.model.BiometricUnlockModel
+import com.android.systemui.keyguard.shared.model.DozeStateModel
+import com.android.systemui.keyguard.shared.model.DozeTransitionModel
 import com.android.systemui.keyguard.shared.model.StatusBarState
 import com.android.systemui.keyguard.shared.model.WakefulnessModel
 import com.android.systemui.plugins.statusbar.StatusBarStateController
@@ -108,6 +113,9 @@ interface KeyguardRepository {
     */
    val dozeAmount: Flow<Float>

+    /** Doze state information, as it transitions */
+    val dozeTransitionModel: Flow<DozeTransitionModel>
+
    /** Observable for the [StatusBarState] */
    val statusBarState: Flow<StatusBarState>

@@ -154,6 +162,7 @@ constructor(
    biometricUnlockController: BiometricUnlockController,
    private val keyguardStateController: KeyguardStateController,
    private val keyguardUpdateMonitor: KeyguardUpdateMonitor,
+    private val dozeTransitionListener: DozeTransitionListener,
 ) : KeyguardRepository {
    private val _animateBottomAreaDozingTransitions = MutableStateFlow(false)
    override val animateBottomAreaDozingTransitions =
@@ -286,6 +295,37 @@ constructor(
        awaitClose { statusBarStateController.removeCallback(callback) }
    }

+    override val dozeTransitionModel: Flow<DozeTransitionModel> = conflatedCallbackFlow {
+        val callback =
+            object : DozeTransitionCallback {
+                override fun onDozeTransition(
+                    oldState: DozeMachine.State,
+                    newState: DozeMachine.State
+                ) {
+                    trySendWithFailureLogging(
+                        DozeTransitionModel(
+                            from = dozeMachineStateToModel(oldState),
+                            to = dozeMachineStateToModel(newState),
+                        ),
+                        TAG,
+                        "doze transition model"
+                    )
+                }
+            }
+
+        dozeTransitionListener.addCallback(callback)
+        trySendWithFailureLogging(
+            DozeTransitionModel(
+                from = dozeMachineStateToModel(dozeTransitionListener.oldState),
+                to = dozeMachineStateToModel(dozeTransitionListener.newState),
+            ),
+            TAG,
+            "initial doze transition model"
+        )
+
+        awaitClose { dozeTransitionListener.removeCallback(callback) }
+    }
+
    override fun isKeyguardShowing(): Boolean {
        return keyguardStateController.isShowing
    }
@@ -407,6 +447,25 @@ constructor(
        }
    }

+    private fun dozeMachineStateToModel(state: DozeMachine.State): DozeStateModel {
+        return when (state) {
+            DozeMachine.State.UNINITIALIZED -> DozeStateModel.UNINITIALIZED
+            DozeMachine.State.INITIALIZED -> DozeStateModel.INITIALIZED
+            DozeMachine.State.DOZE -> DozeStateModel.DOZE
+            DozeMachine.State.DOZE_SUSPEND_TRIGGERS -> DozeStateModel.DOZE_SUSPEND_TRIGGERS
+            DozeMachine.State.DOZE_AOD -> DozeStateModel.DOZE_AOD
+            DozeMachine.State.DOZE_REQUEST_PULSE -> DozeStateModel.DOZE_REQUEST_PULSE
+            DozeMachine.State.DOZE_PULSING -> DozeStateModel.DOZE_PULSING
+            DozeMachine.State.DOZE_PULSING_BRIGHT -> DozeStateModel.DOZE_PULSING_BRIGHT
+            DozeMachine.State.DOZE_PULSE_DONE -> DozeStateModel.DOZE_PULSE_DONE
+            DozeMachine.State.FINISH -> DozeStateModel.FINISH
+            DozeMachine.State.DOZE_AOD_PAUSED -> DozeStateModel.DOZE_AOD_PAUSED
+            DozeMachine.State.DOZE_AOD_PAUSING -> DozeStateModel.DOZE_AOD_PAUSING
+            DozeMachine.State.DOZE_AOD_DOCKED -> DozeStateModel.DOZE_AOD_DOCKED
+            else -> throw IllegalArgumentException("Invalid DozeMachine.State: state")
+        }
+    }
+
    companion object {
        private const val TAG = "KeyguardRepositoryImpl"
    }
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/AodLockscreenTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/AodLockscreenTransitionInteractor.kt
@@ -21,10 +21,9 @@ import com.android.systemui.animation.Interpolators
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
+import com.android.systemui.keyguard.shared.model.DozeStateModel
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.TransitionInfo
-import com.android.systemui.keyguard.shared.model.WakefulnessModel.Companion.isSleepingOrStartingToSleep
-import com.android.systemui.keyguard.shared.model.WakefulnessModel.Companion.isWakingOrStartingToWake
 import com.android.systemui.util.kotlin.sample
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
@@ -39,27 +38,24 @@ constructor(
    private val keyguardInteractor: KeyguardInteractor,
    private val keyguardTransitionRepository: KeyguardTransitionRepository,
    private val keyguardTransitionInteractor: KeyguardTransitionInteractor,
-) : TransitionInteractor("AOD<->LOCKSCREEN") {
+) : TransitionInteractor(AodLockscreenTransitionInteractor::class.simpleName!!) {

    override fun start() {
+        listenForTransitionToAodFromLockscreen()
+        listenForTransitionToLockscreenFromAod()
+    }
+
+    private fun listenForTransitionToAodFromLockscreen() {
        scope.launch {
-            /*
-             * Listening to the startedKeyguardTransitionStep (last started step) allows this code
-             * to interrupt an active transition, as long as they were either going to LOCKSCREEN or
-             * AOD state. One example is when the user presses the power button in the middle of an
-             * active transition.
-             */
-            keyguardInteractor.wakefulnessState
+            keyguardInteractor
+                .dozeTransitionTo(DozeStateModel.DOZE_AOD)
                .sample(
                    keyguardTransitionInteractor.startedKeyguardTransitionStep,
                    { a, b -> Pair(a, b) }
                )
                .collect { pair ->
-                    val (wakefulnessState, lastStartedStep) = pair
-                    if (
-                        isSleepingOrStartingToSleep(wakefulnessState) &&
-                            lastStartedStep.to == KeyguardState.LOCKSCREEN
-                    ) {
+                    val (dozeToAod, lastStartedStep) = pair
+                    if (lastStartedStep.to == KeyguardState.LOCKSCREEN) {
                        keyguardTransitionRepository.startTransition(
                            TransitionInfo(
                                name,
@@ -68,10 +64,22 @@ constructor(
                                getAnimator(),
                            )
                        )
-                    } else if (
-                        isWakingOrStartingToWake(wakefulnessState) &&
-                            lastStartedStep.to == KeyguardState.AOD
-                    ) {
+                    }
+                }
+        }
+    }
+
+    private fun listenForTransitionToLockscreenFromAod() {
+        scope.launch {
+            keyguardInteractor
+                .dozeTransitionTo(DozeStateModel.FINISH)
+                .sample(
+                    keyguardTransitionInteractor.startedKeyguardTransitionStep,
+                    { a, b -> Pair(a, b) }
+                )
+                .collect { pair ->
+                    val (dozeToAod, lastStartedStep) = pair
+                    if (lastStartedStep.to == KeyguardState.AOD) {
                        keyguardTransitionRepository.startTransition(
                            TransitionInfo(
                                name,
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/AodToGoneTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/AodToGoneTransitionInteractor.kt
@@ -40,7 +40,7 @@ constructor(
    private val keyguardInteractor: KeyguardInteractor,
    private val keyguardTransitionRepository: KeyguardTransitionRepository,
    private val keyguardTransitionInteractor: KeyguardTransitionInteractor,
-) : TransitionInteractor("AOD->GONE") {
+) : TransitionInteractor(AodToGoneTransitionInteractor::class.simpleName!!) {

    private val wakeAndUnlockModes =
        setOf(WAKE_AND_UNLOCK, WAKE_AND_UNLOCK_FROM_DREAM, WAKE_AND_UNLOCK_PULSING)
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/BouncerToGoneTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/BouncerToGoneTransitionInteractor.kt
@@ -40,7 +40,7 @@ constructor(
    private val shadeRepository: ShadeRepository,
    private val keyguardTransitionRepository: KeyguardTransitionRepository,
    private val keyguardTransitionInteractor: KeyguardTransitionInteractor
-) : TransitionInteractor("BOUNCER->GONE") {
+) : TransitionInteractor(BouncerToGoneTransitionInteractor::class.simpleName!!) {

    private var transitionId: UUID? = null

--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/DreamingLockscreenTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/DreamingLockscreenTransitionInteractor.kt
@@ -21,12 +21,14 @@ import com.android.systemui.animation.Interpolators
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
+import com.android.systemui.keyguard.shared.model.DozeStateModel
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.TransitionInfo
 import com.android.systemui.util.kotlin.sample
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.launch

@SysUISingleton
@@ -37,32 +39,43 @@ constructor(
    private val keyguardInteractor: KeyguardInteractor,
    private val keyguardTransitionRepository: KeyguardTransitionRepository,
    private val keyguardTransitionInteractor: KeyguardTransitionInteractor,
-) : TransitionInteractor("DREAMING<->LOCKSCREEN") {
+) : TransitionInteractor(DreamingLockscreenTransitionInteractor::class.simpleName!!) {

    override fun start() {
        scope.launch {
            keyguardInteractor.isDreaming
-                .sample(keyguardTransitionInteractor.finishedKeyguardState, { a, b -> Pair(a, b) })
-                .collect { pair ->
-                    val (isDreaming, keyguardState) = pair
-                    if (isDreaming && keyguardState == KeyguardState.LOCKSCREEN) {
-                        keyguardTransitionRepository.startTransition(
-                            TransitionInfo(
-                                name,
-                                KeyguardState.LOCKSCREEN,
-                                KeyguardState.DREAMING,
-                                getAnimator(),
+                .sample(
+                    combine(
+                        keyguardInteractor.dozeTransitionModel,
+                        keyguardTransitionInteractor.finishedKeyguardState
+                    ) { a, b -> Pair(a, b) },
+                    { a, bc -> Triple(a, bc.first, bc.second) }
+                )
+                .collect { triple ->
+                    val (isDreaming, dozeTransitionModel, keyguardState) = triple
+                    // Dozing/AOD and dreaming have overlapping events. If the state remains in
+                    // FINISH, it means that doze mode is not running and DREAMING is ok to
+                    // commence.
+                    if (dozeTransitionModel.to == DozeStateModel.FINISH) {
+                        if (isDreaming && keyguardState == KeyguardState.LOCKSCREEN) {
+                            keyguardTransitionRepository.startTransition(
+                                TransitionInfo(
+                                    name,
+                                    KeyguardState.LOCKSCREEN,
+                                    KeyguardState.DREAMING,
+                                    getAnimator(),
+                                )
                            )
-                        )
-                    } else if (!isDreaming && keyguardState == KeyguardState.DREAMING) {
-                        keyguardTransitionRepository.startTransition(
-                            TransitionInfo(
-                                name,
-                                KeyguardState.DREAMING,
-                                KeyguardState.LOCKSCREEN,
-                                getAnimator(),
+                        } else if (!isDreaming && keyguardState == KeyguardState.DREAMING) {
+                            keyguardTransitionRepository.startTransition(
+                                TransitionInfo(
+                                    name,
+                                    KeyguardState.DREAMING,
+                                    KeyguardState.LOCKSCREEN,
+                                    getAnimator(),
+                                )
                            )
-                        )
+                        }
                    }
                }
        }
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
@@ -20,10 +20,13 @@ package com.android.systemui.keyguard.domain.interactor
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.keyguard.data.repository.KeyguardRepository
 import com.android.systemui.keyguard.shared.model.BiometricUnlockModel
+import com.android.systemui.keyguard.shared.model.DozeStateModel
+import com.android.systemui.keyguard.shared.model.DozeTransitionModel
 import com.android.systemui.keyguard.shared.model.StatusBarState
 import com.android.systemui.keyguard.shared.model.WakefulnessModel
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.filter

 /**
 * Encapsulates business-logic related to the keyguard but not to a more specific part within it.
@@ -41,6 +44,8 @@ constructor(
    val dozeAmount: Flow<Float> = repository.dozeAmount
    /** Whether the system is in doze mode. */
    val isDozing: Flow<Boolean> = repository.isDozing
+    /** Doze transition information. */
+    val dozeTransitionModel: Flow<DozeTransitionModel> = repository.dozeTransitionModel
    /**
     * Whether the system is dreaming. [isDreaming] will be always be true when [isDozing] is true,
     * but not vice-versa.
@@ -62,6 +67,10 @@ constructor(
     */
    val biometricUnlockState: Flow<BiometricUnlockModel> = repository.biometricUnlockState

+    fun dozeTransitionTo(state: DozeStateModel): Flow<DozeTransitionModel> {
+        return dozeTransitionModel.filter { it.to == state }
+    }
+
    fun isKeyguardShowing(): Boolean {
        return repository.isKeyguardShowing()
    }
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/LockscreenBouncerTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/LockscreenBouncerTransitionInteractor.kt
@@ -44,7 +44,7 @@ constructor(
    private val shadeRepository: ShadeRepository,
    private val keyguardTransitionRepository: KeyguardTransitionRepository,
    private val keyguardTransitionInteractor: KeyguardTransitionInteractor
-) : TransitionInteractor("LOCKSCREEN<->BOUNCER") {
+) : TransitionInteractor(LockscreenBouncerTransitionInteractor::class.simpleName!!) {

    private var transitionId: UUID? = null

--- a/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/DozeStateModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/DozeStateModel.kt
@@ -0,0 +1,46 @@
+/*
+ * 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.systemui.keyguard.shared.model
+
+/** Model device doze states. */
+enum class DozeStateModel {
+    /** Default state. Transition to INITIALIZED to get Doze going. */
+    UNINITIALIZED,
+    /** Doze components are set up. Followed by transition to DOZE or DOZE_AOD. */
+    INITIALIZED,
+    /** Regular doze. Device is asleep and listening for pulse triggers. */
+    DOZE,
+    /** Deep doze. Device is asleep and is not listening for pulse triggers. */
+    DOZE_SUSPEND_TRIGGERS,
+    /** Always-on doze. Device is asleep, showing UI and listening for pulse triggers. */
+    DOZE_AOD,
+    /** Pulse has been requested. Device is awake and preparing UI */
+    DOZE_REQUEST_PULSE,
+    /** Pulse is showing. Device is awake and showing UI. */
+    DOZE_PULSING,
+    /** Pulse is showing with bright wallpaper. Device is awake and showing UI. */
+    DOZE_PULSING_BRIGHT,
+    /** Pulse is done showing. Followed by transition to DOZE or DOZE_AOD. */
+    DOZE_PULSE_DONE,
+    /** Doze is done. DozeService is finished. */
+    FINISH,
+    /** AOD, but the display is temporarily off. */
+    DOZE_AOD_PAUSED,
+    /** AOD, prox is near, transitions to DOZE_AOD_PAUSED after a timeout. */
+    DOZE_AOD_PAUSING,
+    /** Always-on doze. Device is awake, showing docking UI and listening for pulse triggers. */
+    DOZE_AOD_DOCKED
+}
--- a/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/DozeTransitionModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/DozeTransitionModel.kt
@@ -0,0 +1,22 @@
+/*
+ * 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.systemui.keyguard.shared.model
+
+/** Doze transition information. */
+data class DozeTransitionModel(
+    val from: DozeStateModel = DozeStateModel.UNINITIALIZED,
+    val to: DozeStateModel = DozeStateModel.UNINITIALIZED,
+)
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardRepositoryImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardRepositoryImplTest.kt
@@ -22,18 +22,25 @@ import com.android.keyguard.KeyguardUpdateMonitorCallback
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.common.shared.model.Position
 import com.android.systemui.doze.DozeHost
+import com.android.systemui.doze.DozeMachine
+import com.android.systemui.doze.DozeTransitionCallback
+import com.android.systemui.doze.DozeTransitionListener
 import com.android.systemui.keyguard.WakefulnessLifecycle
 import com.android.systemui.keyguard.shared.model.BiometricUnlockModel
+import com.android.systemui.keyguard.shared.model.DozeStateModel
+import com.android.systemui.keyguard.shared.model.DozeTransitionModel
 import com.android.systemui.keyguard.shared.model.WakefulnessModel
 import com.android.systemui.plugins.statusbar.StatusBarStateController
 import com.android.systemui.statusbar.phone.BiometricUnlockController
 import com.android.systemui.statusbar.policy.KeyguardStateController
 import com.android.systemui.util.mockito.argumentCaptor
 import com.android.systemui.util.mockito.whenever
+import com.android.systemui.util.mockito.withArgCaptor
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.test.runBlockingTest
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -52,6 +59,7 @@ class KeyguardRepositoryImplTest : SysuiTestCase() {
    @Mock private lateinit var wakefulnessLifecycle: WakefulnessLifecycle
    @Mock private lateinit var keyguardUpdateMonitor: KeyguardUpdateMonitor
    @Mock private lateinit var biometricUnlockController: BiometricUnlockController
+    @Mock private lateinit var dozeTransitionListener: DozeTransitionListener

    private lateinit var underTest: KeyguardRepositoryImpl

@@ -67,272 +75,349 @@ class KeyguardRepositoryImplTest : SysuiTestCase() {
                biometricUnlockController,
                keyguardStateController,
                keyguardUpdateMonitor,
+                dozeTransitionListener,
            )
    }

    @Test
-    fun animateBottomAreaDozingTransitions() = runBlockingTest {
-        assertThat(underTest.animateBottomAreaDozingTransitions.value).isEqualTo(false)
+    fun animateBottomAreaDozingTransitions() =
+        runTest(UnconfinedTestDispatcher()) {
+            assertThat(underTest.animateBottomAreaDozingTransitions.value).isEqualTo(false)

-        underTest.setAnimateDozingTransitions(true)
-        assertThat(underTest.animateBottomAreaDozingTransitions.value).isTrue()
+            underTest.setAnimateDozingTransitions(true)
+            assertThat(underTest.animateBottomAreaDozingTransitions.value).isTrue()

-        underTest.setAnimateDozingTransitions(false)
-        assertThat(underTest.animateBottomAreaDozingTransitions.value).isFalse()
+            underTest.setAnimateDozingTransitions(false)
+            assertThat(underTest.animateBottomAreaDozingTransitions.value).isFalse()

-        underTest.setAnimateDozingTransitions(true)
-        assertThat(underTest.animateBottomAreaDozingTransitions.value).isTrue()
-    }
+            underTest.setAnimateDozingTransitions(true)
+            assertThat(underTest.animateBottomAreaDozingTransitions.value).isTrue()
+        }

    @Test
-    fun bottomAreaAlpha() = runBlockingTest {
-        assertThat(underTest.bottomAreaAlpha.value).isEqualTo(1f)
+    fun bottomAreaAlpha() =
+        runTest(UnconfinedTestDispatcher()) {
+            assertThat(underTest.bottomAreaAlpha.value).isEqualTo(1f)

-        underTest.setBottomAreaAlpha(0.1f)
-        assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.1f)
+            underTest.setBottomAreaAlpha(0.1f)
+            assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.1f)

-        underTest.setBottomAreaAlpha(0.2f)
-        assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.2f)
+            underTest.setBottomAreaAlpha(0.2f)
+            assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.2f)

-        underTest.setBottomAreaAlpha(0.3f)
-        assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.3f)
+            underTest.setBottomAreaAlpha(0.3f)
+            assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.3f)

-        underTest.setBottomAreaAlpha(0.5f)
-        assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.5f)
+            underTest.setBottomAreaAlpha(0.5f)
+            assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.5f)

-        underTest.setBottomAreaAlpha(1.0f)
-        assertThat(underTest.bottomAreaAlpha.value).isEqualTo(1f)
-    }
+            underTest.setBottomAreaAlpha(1.0f)
+            assertThat(underTest.bottomAreaAlpha.value).isEqualTo(1f)
+        }

    @Test
-    fun clockPosition() = runBlockingTest {
-        assertThat(underTest.clockPosition.value).isEqualTo(Position(0, 0))
+    fun clockPosition() =
+        runTest(UnconfinedTestDispatcher()) {
+            assertThat(underTest.clockPosition.value).isEqualTo(Position(0, 0))

-        underTest.setClockPosition(0, 1)
-        assertThat(underTest.clockPosition.value).isEqualTo(Position(0, 1))
+            underTest.setClockPosition(0, 1)
+            assertThat(underTest.clockPosition.value).isEqualTo(Position(0, 1))

-        underTest.setClockPosition(1, 9)
-        assertThat(underTest.clockPosition.value).isEqualTo(Position(1, 9))
+            underTest.setClockPosition(1, 9)
+            assertThat(underTest.clockPosition.value).isEqualTo(Position(1, 9))

-        underTest.setClockPosition(1, 0)
-        assertThat(underTest.clockPosition.value).isEqualTo(Position(1, 0))
+            underTest.setClockPosition(1, 0)
+            assertThat(underTest.clockPosition.value).isEqualTo(Position(1, 0))

-        underTest.setClockPosition(3, 1)
-        assertThat(underTest.clockPosition.value).isEqualTo(Position(3, 1))
-    }
+            underTest.setClockPosition(3, 1)
+            assertThat(underTest.clockPosition.value).isEqualTo(Position(3, 1))
+        }

    @Test
-    fun isKeyguardShowing() = runBlockingTest {
-        whenever(keyguardStateController.isShowing).thenReturn(false)
-        var latest: Boolean? = null
-        val job = underTest.isKeyguardShowing.onEach { latest = it }.launchIn(this)
+    fun isKeyguardShowing() =
+        runTest(UnconfinedTestDispatcher()) {
+            whenever(keyguardStateController.isShowing).thenReturn(false)
+            var latest: Boolean? = null
+            val job = underTest.isKeyguardShowing.onEach { latest = it }.launchIn(this)

-        assertThat(latest).isFalse()
-        assertThat(underTest.isKeyguardShowing()).isFalse()
+            assertThat(latest).isFalse()
+            assertThat(underTest.isKeyguardShowing()).isFalse()

-        val captor = argumentCaptor<KeyguardStateController.Callback>()
-        verify(keyguardStateController).addCallback(captor.capture())
+            val captor = argumentCaptor<KeyguardStateController.Callback>()
+            verify(keyguardStateController).addCallback(captor.capture())

-        whenever(keyguardStateController.isShowing).thenReturn(true)
-        captor.value.onKeyguardShowingChanged()
-        assertThat(latest).isTrue()
-        assertThat(underTest.isKeyguardShowing()).isTrue()
+            whenever(keyguardStateController.isShowing).thenReturn(true)
+            captor.value.onKeyguardShowingChanged()
+            assertThat(latest).isTrue()
+            assertThat(underTest.isKeyguardShowing()).isTrue()

-        whenever(keyguardStateController.isShowing).thenReturn(false)
-        captor.value.onKeyguardShowingChanged()
-        assertThat(latest).isFalse()
-        assertThat(underTest.isKeyguardShowing()).isFalse()
+            whenever(keyguardStateController.isShowing).thenReturn(false)
+            captor.value.onKeyguardShowingChanged()
+            assertThat(latest).isFalse()
+            assertThat(underTest.isKeyguardShowing()).isFalse()

-        job.cancel()
-    }
+            job.cancel()
+        }

    @Test
-    fun isDozing() = runBlockingTest {
-        var latest: Boolean? = null
-        val job = underTest.isDozing.onEach { latest = it }.launchIn(this)
+    fun isDozing() =
+        runTest(UnconfinedTestDispatcher()) {
+            var latest: Boolean? = null
+            val job = underTest.isDozing.onEach { latest = it }.launchIn(this)

-        val captor = argumentCaptor<DozeHost.Callback>()
-        verify(dozeHost).addCallback(captor.capture())
+            val captor = argumentCaptor<DozeHost.Callback>()
+            verify(dozeHost).addCallback(captor.capture())

-        captor.value.onDozingChanged(true)
-        assertThat(latest).isTrue()
+            captor.value.onDozingChanged(true)
+            assertThat(latest).isTrue()

-        captor.value.onDozingChanged(false)
-        assertThat(latest).isFalse()
+            captor.value.onDozingChanged(false)
+            assertThat(latest).isFalse()

-        job.cancel()
-        verify(dozeHost).removeCallback(captor.value)
-    }
+            job.cancel()
+            verify(dozeHost).removeCallback(captor.value)
+        }

    @Test
-    fun `isDozing - starts with correct initial value for isDozing`() = runBlockingTest {
-        var latest: Boolean? = null
+    fun `isDozing - starts with correct initial value for isDozing`() =
+        runTest(UnconfinedTestDispatcher()) {
+            var latest: Boolean? = null

-        whenever(statusBarStateController.isDozing).thenReturn(true)
-        var job = underTest.isDozing.onEach { latest = it }.launchIn(this)
-        assertThat(latest).isTrue()
-        job.cancel()
+            whenever(statusBarStateController.isDozing).thenReturn(true)
+            var job = underTest.isDozing.onEach { latest = it }.launchIn(this)
+            assertThat(latest).isTrue()
+            job.cancel()

-        whenever(statusBarStateController.isDozing).thenReturn(false)
-        job = underTest.isDozing.onEach { latest = it }.launchIn(this)
-        assertThat(latest).isFalse()
-        job.cancel()
-    }
+            whenever(statusBarStateController.isDozing).thenReturn(false)
+            job = underTest.isDozing.onEach { latest = it }.launchIn(this)
+            assertThat(latest).isFalse()
+            job.cancel()
+        }

    @Test
-    fun dozeAmount() = runBlockingTest {
-        val values = mutableListOf<Float>()
-        val job = underTest.dozeAmount.onEach(values::add).launchIn(this)
+    fun dozeAmount() =
+        runTest(UnconfinedTestDispatcher()) {
+            val values = mutableListOf<Float>()
+            val job = underTest.dozeAmount.onEach(values::add).launchIn(this)

-        val captor = argumentCaptor<StatusBarStateController.StateListener>()
-        verify(statusBarStateController).addCallback(captor.capture())
+            val captor = argumentCaptor<StatusBarStateController.StateListener>()
+            verify(statusBarStateController).addCallback(captor.capture())

-        captor.value.onDozeAmountChanged(0.433f, 0.4f)
-        captor.value.onDozeAmountChanged(0.498f, 0.5f)
-        captor.value.onDozeAmountChanged(0.661f, 0.65f)
+            captor.value.onDozeAmountChanged(0.433f, 0.4f)
+            captor.value.onDozeAmountChanged(0.498f, 0.5f)
+            captor.value.onDozeAmountChanged(0.661f, 0.65f)

-        assertThat(values).isEqualTo(listOf(0f, 0.4f, 0.5f, 0.65f))
+            assertThat(values).isEqualTo(listOf(0f, 0.4f, 0.5f, 0.65f))

-        job.cancel()
-        verify(statusBarStateController).removeCallback(captor.value)
-    }
+            job.cancel()
+            verify(statusBarStateController).removeCallback(captor.value)
+        }

    @Test
-    fun wakefulness() = runBlockingTest {
-        val values = mutableListOf<WakefulnessModel>()
-        val job = underTest.wakefulnessState.onEach(values::add).launchIn(this)
+    fun wakefulness() =
+        runTest(UnconfinedTestDispatcher()) {
+            val values = mutableListOf<WakefulnessModel>()
+            val job = underTest.wakefulnessState.onEach(values::add).launchIn(this)

-        val captor = argumentCaptor<WakefulnessLifecycle.Observer>()
-        verify(wakefulnessLifecycle).addObserver(captor.capture())
+            val captor = argumentCaptor<WakefulnessLifecycle.Observer>()
+            verify(wakefulnessLifecycle).addObserver(captor.capture())

-        captor.value.onStartedWakingUp()
-        captor.value.onFinishedWakingUp()
-        captor.value.onStartedGoingToSleep()
-        captor.value.onFinishedGoingToSleep()
+            captor.value.onStartedWakingUp()
+            captor.value.onFinishedWakingUp()
+            captor.value.onStartedGoingToSleep()
+            captor.value.onFinishedGoingToSleep()

-        assertThat(values)
-            .isEqualTo(
-                listOf(
-                    // Initial value will be ASLEEP
-                    WakefulnessModel.ASLEEP,
-                    WakefulnessModel.STARTING_TO_WAKE,
-                    WakefulnessModel.AWAKE,
-                    WakefulnessModel.STARTING_TO_SLEEP,
-                    WakefulnessModel.ASLEEP,
+            assertThat(values)
+                .isEqualTo(
+                    listOf(
+                        // Initial value will be ASLEEP
+                        WakefulnessModel.ASLEEP,
+                        WakefulnessModel.STARTING_TO_WAKE,
+                        WakefulnessModel.AWAKE,
+                        WakefulnessModel.STARTING_TO_SLEEP,
+                        WakefulnessModel.ASLEEP,
+                    )
                )
-            )

-        job.cancel()
-        verify(wakefulnessLifecycle).removeObserver(captor.value)
-    }
+            job.cancel()
+            verify(wakefulnessLifecycle).removeObserver(captor.value)
+        }

    @Test
-    fun isUdfpsSupported() = runBlockingTest {
-        whenever(keyguardUpdateMonitor.isUdfpsSupported).thenReturn(true)
-        assertThat(underTest.isUdfpsSupported()).isTrue()
+    fun isUdfpsSupported() =
+        runTest(UnconfinedTestDispatcher()) {
+            whenever(keyguardUpdateMonitor.isUdfpsSupported).thenReturn(true)
+            assertThat(underTest.isUdfpsSupported()).isTrue()

-        whenever(keyguardUpdateMonitor.isUdfpsSupported).thenReturn(false)
-        assertThat(underTest.isUdfpsSupported()).isFalse()
-    }
+            whenever(keyguardUpdateMonitor.isUdfpsSupported).thenReturn(false)
+            assertThat(underTest.isUdfpsSupported()).isFalse()
+        }

    @Test
-    fun isBouncerShowing() = runBlockingTest {
-        whenever(keyguardStateController.isBouncerShowing).thenReturn(false)
-        var latest: Boolean? = null
-        val job = underTest.isBouncerShowing.onEach { latest = it }.launchIn(this)
+    fun isBouncerShowing() =
+        runTest(UnconfinedTestDispatcher()) {
+            whenever(keyguardStateController.isBouncerShowing).thenReturn(false)
+            var latest: Boolean? = null
+            val job = underTest.isBouncerShowing.onEach { latest = it }.launchIn(this)

-        assertThat(latest).isFalse()
+            assertThat(latest).isFalse()

-        val captor = argumentCaptor<KeyguardStateController.Callback>()
-        verify(keyguardStateController).addCallback(captor.capture())
+            val captor = argumentCaptor<KeyguardStateController.Callback>()
+            verify(keyguardStateController).addCallback(captor.capture())

-        whenever(keyguardStateController.isBouncerShowing).thenReturn(true)
-        captor.value.onBouncerShowingChanged()
-        assertThat(latest).isTrue()
+            whenever(keyguardStateController.isBouncerShowing).thenReturn(true)
+            captor.value.onBouncerShowingChanged()
+            assertThat(latest).isTrue()

-        whenever(keyguardStateController.isBouncerShowing).thenReturn(false)
-        captor.value.onBouncerShowingChanged()
-        assertThat(latest).isFalse()
+            whenever(keyguardStateController.isBouncerShowing).thenReturn(false)
+            captor.value.onBouncerShowingChanged()
+            assertThat(latest).isFalse()

-        job.cancel()
-    }
+            job.cancel()
+        }

    @Test
-    fun isKeyguardGoingAway() = runBlockingTest {
-        whenever(keyguardStateController.isKeyguardGoingAway).thenReturn(false)
-        var latest: Boolean? = null
-        val job = underTest.isKeyguardGoingAway.onEach { latest = it }.launchIn(this)
+    fun isKeyguardGoingAway() =
+        runTest(UnconfinedTestDispatcher()) {
+            whenever(keyguardStateController.isKeyguardGoingAway).thenReturn(false)
+            var latest: Boolean? = null
+            val job = underTest.isKeyguardGoingAway.onEach { latest = it }.launchIn(this)

-        assertThat(latest).isFalse()
+            assertThat(latest).isFalse()

-        val captor = argumentCaptor<KeyguardStateController.Callback>()
-        verify(keyguardStateController).addCallback(captor.capture())
+            val captor = argumentCaptor<KeyguardStateController.Callback>()
+            verify(keyguardStateController).addCallback(captor.capture())

-        whenever(keyguardStateController.isKeyguardGoingAway).thenReturn(true)
-        captor.value.onKeyguardGoingAwayChanged()
-        assertThat(latest).isTrue()
+            whenever(keyguardStateController.isKeyguardGoingAway).thenReturn(true)
+            captor.value.onKeyguardGoingAwayChanged()
+            assertThat(latest).isTrue()

-        whenever(keyguardStateController.isKeyguardGoingAway).thenReturn(false)
-        captor.value.onKeyguardGoingAwayChanged()
-        assertThat(latest).isFalse()
+            whenever(keyguardStateController.isKeyguardGoingAway).thenReturn(false)
+            captor.value.onKeyguardGoingAwayChanged()
+            assertThat(latest).isFalse()

-        job.cancel()
-    }
+            job.cancel()
+        }

    @Test
-    fun isDreaming() = runBlockingTest {
-        whenever(keyguardUpdateMonitor.isDreaming()).thenReturn(false)
-        var latest: Boolean? = null
-        val job = underTest.isDreaming.onEach { latest = it }.launchIn(this)
+    fun isDreaming() =
+        runTest(UnconfinedTestDispatcher()) {
+            whenever(keyguardUpdateMonitor.isDreaming()).thenReturn(false)
+            var latest: Boolean? = null
+            val job = underTest.isDreaming.onEach { latest = it }.launchIn(this)

-        assertThat(latest).isFalse()
+            assertThat(latest).isFalse()

-        val captor = argumentCaptor<KeyguardUpdateMonitorCallback>()
-        verify(keyguardUpdateMonitor).registerCallback(captor.capture())
+            val captor = argumentCaptor<KeyguardUpdateMonitorCallback>()
+            verify(keyguardUpdateMonitor).registerCallback(captor.capture())

-        captor.value.onDreamingStateChanged(true)
-        assertThat(latest).isTrue()
+            captor.value.onDreamingStateChanged(true)
+            assertThat(latest).isTrue()

-        captor.value.onDreamingStateChanged(false)
-        assertThat(latest).isFalse()
+            captor.value.onDreamingStateChanged(false)
+            assertThat(latest).isFalse()

-        job.cancel()
-    }
+            job.cancel()
+        }

    @Test
-    fun biometricUnlockState() = runBlockingTest {
-        val values = mutableListOf<BiometricUnlockModel>()
-        val job = underTest.biometricUnlockState.onEach(values::add).launchIn(this)
+    fun biometricUnlockState() =
+        runTest(UnconfinedTestDispatcher()) {
+            val values = mutableListOf<BiometricUnlockModel>()
+            val job = underTest.biometricUnlockState.onEach(values::add).launchIn(this)

-        val captor = argumentCaptor<BiometricUnlockController.BiometricModeListener>()
-        verify(biometricUnlockController).addBiometricModeListener(captor.capture())
+            val captor = argumentCaptor<BiometricUnlockController.BiometricModeListener>()
+            verify(biometricUnlockController).addBiometricModeListener(captor.capture())

-        captor.value.onModeChanged(BiometricUnlockController.MODE_NONE)
-        captor.value.onModeChanged(BiometricUnlockController.MODE_WAKE_AND_UNLOCK)
-        captor.value.onModeChanged(BiometricUnlockController.MODE_WAKE_AND_UNLOCK_PULSING)
-        captor.value.onModeChanged(BiometricUnlockController.MODE_SHOW_BOUNCER)
-        captor.value.onModeChanged(BiometricUnlockController.MODE_ONLY_WAKE)
-        captor.value.onModeChanged(BiometricUnlockController.MODE_UNLOCK_COLLAPSING)
-        captor.value.onModeChanged(BiometricUnlockController.MODE_DISMISS_BOUNCER)
-        captor.value.onModeChanged(BiometricUnlockController.MODE_WAKE_AND_UNLOCK_FROM_DREAM)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_NONE)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_WAKE_AND_UNLOCK)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_WAKE_AND_UNLOCK_PULSING)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_SHOW_BOUNCER)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_ONLY_WAKE)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_UNLOCK_COLLAPSING)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_DISMISS_BOUNCER)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_WAKE_AND_UNLOCK_FROM_DREAM)

-        assertThat(values)
-            .isEqualTo(
-                listOf(
-                    // Initial value will be NONE, followed by onModeChanged() call
-                    BiometricUnlockModel.NONE,
-                    BiometricUnlockModel.NONE,
-                    BiometricUnlockModel.WAKE_AND_UNLOCK,
-                    BiometricUnlockModel.WAKE_AND_UNLOCK_PULSING,
-                    BiometricUnlockModel.SHOW_BOUNCER,
-                    BiometricUnlockModel.ONLY_WAKE,
-                    BiometricUnlockModel.UNLOCK_COLLAPSING,
-                    BiometricUnlockModel.DISMISS_BOUNCER,
-                    BiometricUnlockModel.WAKE_AND_UNLOCK_FROM_DREAM,
+            assertThat(values)
+                .isEqualTo(
+                    listOf(
+                        // Initial value will be NONE, followed by onModeChanged() call
+                        BiometricUnlockModel.NONE,
+                        BiometricUnlockModel.NONE,
+                        BiometricUnlockModel.WAKE_AND_UNLOCK,
+                        BiometricUnlockModel.WAKE_AND_UNLOCK_PULSING,
+                        BiometricUnlockModel.SHOW_BOUNCER,
+                        BiometricUnlockModel.ONLY_WAKE,
+                        BiometricUnlockModel.UNLOCK_COLLAPSING,
+                        BiometricUnlockModel.DISMISS_BOUNCER,
+                        BiometricUnlockModel.WAKE_AND_UNLOCK_FROM_DREAM,
+                    )
                )
+
+            job.cancel()
+            verify(biometricUnlockController).removeBiometricModeListener(captor.value)
+        }
+
+    @Test
+    fun dozeTransitionModel() =
+        runTest(UnconfinedTestDispatcher()) {
+            // For the initial state
+            whenever(dozeTransitionListener.oldState).thenReturn(DozeMachine.State.UNINITIALIZED)
+            whenever(dozeTransitionListener.newState).thenReturn(DozeMachine.State.UNINITIALIZED)
+
+            val values = mutableListOf<DozeTransitionModel>()
+            val job = underTest.dozeTransitionModel.onEach(values::add).launchIn(this)
+
+            val listener =
+                withArgCaptor<DozeTransitionCallback> {
+                    verify(dozeTransitionListener).addCallback(capture())
+                }
+
+            // These don't have to reflect real transitions from the DozeMachine. Only that the
+            // transitions are properly emitted
+            listener.onDozeTransition(DozeMachine.State.INITIALIZED, DozeMachine.State.DOZE)
+            listener.onDozeTransition(DozeMachine.State.DOZE, DozeMachine.State.DOZE_AOD)
+            listener.onDozeTransition(DozeMachine.State.DOZE_AOD_DOCKED, DozeMachine.State.FINISH)
+            listener.onDozeTransition(
+                DozeMachine.State.DOZE_REQUEST_PULSE,
+                DozeMachine.State.DOZE_PULSING
+            )
+            listener.onDozeTransition(
+                DozeMachine.State.DOZE_SUSPEND_TRIGGERS,
+                DozeMachine.State.DOZE_PULSE_DONE
+            )
+            listener.onDozeTransition(
+                DozeMachine.State.DOZE_AOD_PAUSING,
+                DozeMachine.State.DOZE_AOD_PAUSED
            )

-        job.cancel()
-        verify(biometricUnlockController).removeBiometricModeListener(captor.value)
-    }
+            assertThat(values)
+                .isEqualTo(
+                    listOf(
+                        // Initial value will be UNINITIALIZED
+                        DozeTransitionModel(
+                            DozeStateModel.UNINITIALIZED,
+                            DozeStateModel.UNINITIALIZED
+                        ),
+                        DozeTransitionModel(DozeStateModel.INITIALIZED, DozeStateModel.DOZE),
+                        DozeTransitionModel(DozeStateModel.DOZE, DozeStateModel.DOZE_AOD),
+                        DozeTransitionModel(DozeStateModel.DOZE_AOD_DOCKED, DozeStateModel.FINISH),
+                        DozeTransitionModel(
+                            DozeStateModel.DOZE_REQUEST_PULSE,
+                            DozeStateModel.DOZE_PULSING
+                        ),
+                        DozeTransitionModel(
+                            DozeStateModel.DOZE_SUSPEND_TRIGGERS,
+                            DozeStateModel.DOZE_PULSE_DONE
+                        ),
+                        DozeTransitionModel(
+                            DozeStateModel.DOZE_AOD_PAUSING,
+                            DozeStateModel.DOZE_AOD_PAUSED
+                        ),
+                    )
+                )
+
+            job.cancel()
+            verify(dozeTransitionListener).removeCallback(listener)
+        }
 }
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
@@ -19,6 +19,7 @@ package com.android.systemui.keyguard.data.repository

 import com.android.systemui.common.shared.model.Position
 import com.android.systemui.keyguard.shared.model.BiometricUnlockModel
+import com.android.systemui.keyguard.shared.model.DozeTransitionModel
 import com.android.systemui.keyguard.shared.model.StatusBarState
 import com.android.systemui.keyguard.shared.model.WakefulnessModel
 import kotlinx.coroutines.flow.Flow
@@ -53,6 +54,9 @@ class FakeKeyguardRepository : KeyguardRepository {
    private val _statusBarState = MutableStateFlow(StatusBarState.SHADE)
    override val statusBarState: Flow<StatusBarState> = _statusBarState

+    private val _dozeTransitionModel = MutableStateFlow(DozeTransitionModel())
+    override val dozeTransitionModel: Flow<DozeTransitionModel> = _dozeTransitionModel
+
    private val _wakefulnessState = MutableStateFlow(WakefulnessModel.ASLEEP)
    override val wakefulnessState: Flow<WakefulnessModel> = _wakefulnessState