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

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

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
         };
     }
 

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"
     }

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 {
-            /*
+            keyguardInteractor
-             * Listening to the startedKeyguardTransitionStep (last started step) allows this code
+                .dozeTransitionTo(DozeStateModel.DOZE_AOD)
-             * 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
                 .sample(
                     keyguardTransitionInteractor.startedKeyguardTransitionStep,
                     { a, b -> Pair(a, b) }
                 )
                 .collect { pair ->
-                    val (wakefulnessState, lastStartedStep) = pair
+                    val (dozeToAod, lastStartedStep) = pair
-                    if (
+                    if (lastStartedStep.to == KeyguardState.LOCKSCREEN) {
-                        isSleepingOrStartingToSleep(wakefulnessState) &&
-                            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,

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)

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
 

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) })
+                .sample(
-                .collect { pair ->
+                    combine(
-                    val (isDreaming, keyguardState) = pair
+                        keyguardInteractor.dozeTransitionModel,
-                    if (isDreaming && keyguardState == KeyguardState.LOCKSCREEN) {
+                        keyguardTransitionInteractor.finishedKeyguardState
-                        keyguardTransitionRepository.startTransition(
+                    ) { a, b -> Pair(a, b) },
-                            TransitionInfo(
+                    { a, bc -> Triple(a, bc.first, bc.second) }
-                                name,
+                )
-                                KeyguardState.LOCKSCREEN,
+                .collect { triple ->
-                                KeyguardState.DREAMING,
+                    val (isDreaming, dozeTransitionModel, keyguardState) = triple
-                                getAnimator(),
+                    // 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) {
-                    } else if (!isDreaming && keyguardState == KeyguardState.DREAMING) {
+                            keyguardTransitionRepository.startTransition(
-                        keyguardTransitionRepository.startTransition(
+                                TransitionInfo(
-                            TransitionInfo(
+                                    name,
-                                name,
+                                    KeyguardState.DREAMING,
-                                KeyguardState.DREAMING,
+                                    KeyguardState.LOCKSCREEN,
-                                KeyguardState.LOCKSCREEN,
+                                    getAnimator(),
-                                getAnimator(),
+                                )
                             )
-                        )
+                        }
                     }
                 }
         }

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

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
 

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
+}

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,
+)

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 {
+    fun animateBottomAreaDozingTransitions() =
-        assertThat(underTest.animateBottomAreaDozingTransitions.value).isEqualTo(false)
+        runTest(UnconfinedTestDispatcher()) {
+            assertThat(underTest.animateBottomAreaDozingTransitions.value).isEqualTo(false)
 
-        underTest.setAnimateDozingTransitions(true)
+            underTest.setAnimateDozingTransitions(true)
-        assertThat(underTest.animateBottomAreaDozingTransitions.value).isTrue()
+            assertThat(underTest.animateBottomAreaDozingTransitions.value).isTrue()
 
-        underTest.setAnimateDozingTransitions(false)
+            underTest.setAnimateDozingTransitions(false)
-        assertThat(underTest.animateBottomAreaDozingTransitions.value).isFalse()
+            assertThat(underTest.animateBottomAreaDozingTransitions.value).isFalse()
 
-        underTest.setAnimateDozingTransitions(true)
+            underTest.setAnimateDozingTransitions(true)
-        assertThat(underTest.animateBottomAreaDozingTransitions.value).isTrue()
+            assertThat(underTest.animateBottomAreaDozingTransitions.value).isTrue()
-    }
+        }
 
     @Test
-    fun bottomAreaAlpha() = runBlockingTest {
+    fun bottomAreaAlpha() =
-        assertThat(underTest.bottomAreaAlpha.value).isEqualTo(1f)
+        runTest(UnconfinedTestDispatcher()) {
+            assertThat(underTest.bottomAreaAlpha.value).isEqualTo(1f)
 
-        underTest.setBottomAreaAlpha(0.1f)
+            underTest.setBottomAreaAlpha(0.1f)
-        assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.1f)
+            assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.1f)
 
-        underTest.setBottomAreaAlpha(0.2f)
+            underTest.setBottomAreaAlpha(0.2f)
-        assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.2f)
+            assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.2f)
 
-        underTest.setBottomAreaAlpha(0.3f)
+            underTest.setBottomAreaAlpha(0.3f)
-        assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.3f)
+            assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.3f)
 
-        underTest.setBottomAreaAlpha(0.5f)
+            underTest.setBottomAreaAlpha(0.5f)
-        assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.5f)
+            assertThat(underTest.bottomAreaAlpha.value).isEqualTo(0.5f)
 
-        underTest.setBottomAreaAlpha(1.0f)
+            underTest.setBottomAreaAlpha(1.0f)
-        assertThat(underTest.bottomAreaAlpha.value).isEqualTo(1f)
+            assertThat(underTest.bottomAreaAlpha.value).isEqualTo(1f)
-    }
+        }
 
     @Test
-    fun clockPosition() = runBlockingTest {
+    fun clockPosition() =
-        assertThat(underTest.clockPosition.value).isEqualTo(Position(0, 0))
+        runTest(UnconfinedTestDispatcher()) {
+            assertThat(underTest.clockPosition.value).isEqualTo(Position(0, 0))
 
-        underTest.setClockPosition(0, 1)
+            underTest.setClockPosition(0, 1)
-        assertThat(underTest.clockPosition.value).isEqualTo(Position(0, 1))
+            assertThat(underTest.clockPosition.value).isEqualTo(Position(0, 1))
 
-        underTest.setClockPosition(1, 9)
+            underTest.setClockPosition(1, 9)
-        assertThat(underTest.clockPosition.value).isEqualTo(Position(1, 9))
+            assertThat(underTest.clockPosition.value).isEqualTo(Position(1, 9))
 
-        underTest.setClockPosition(1, 0)
+            underTest.setClockPosition(1, 0)
-        assertThat(underTest.clockPosition.value).isEqualTo(Position(1, 0))
+            assertThat(underTest.clockPosition.value).isEqualTo(Position(1, 0))
 
-        underTest.setClockPosition(3, 1)
+            underTest.setClockPosition(3, 1)
-        assertThat(underTest.clockPosition.value).isEqualTo(Position(3, 1))
+            assertThat(underTest.clockPosition.value).isEqualTo(Position(3, 1))
-    }
+        }
 
     @Test
-    fun isKeyguardShowing() = runBlockingTest {
+    fun isKeyguardShowing() =
-        whenever(keyguardStateController.isShowing).thenReturn(false)
+        runTest(UnconfinedTestDispatcher()) {
-        var latest: Boolean? = null
+            whenever(keyguardStateController.isShowing).thenReturn(false)
-        val job = underTest.isKeyguardShowing.onEach { latest = it }.launchIn(this)
+            var latest: Boolean? = null
+            val job = underTest.isKeyguardShowing.onEach { latest = it }.launchIn(this)
 
-        assertThat(latest).isFalse()
+            assertThat(latest).isFalse()
-        assertThat(underTest.isKeyguardShowing()).isFalse()
+            assertThat(underTest.isKeyguardShowing()).isFalse()
 
-        val captor = argumentCaptor<KeyguardStateController.Callback>()
+            val captor = argumentCaptor<KeyguardStateController.Callback>()
-        verify(keyguardStateController).addCallback(captor.capture())
+            verify(keyguardStateController).addCallback(captor.capture())
 
-        whenever(keyguardStateController.isShowing).thenReturn(true)
+            whenever(keyguardStateController.isShowing).thenReturn(true)
-        captor.value.onKeyguardShowingChanged()
+            captor.value.onKeyguardShowingChanged()
-        assertThat(latest).isTrue()
+            assertThat(latest).isTrue()
-        assertThat(underTest.isKeyguardShowing()).isTrue()
+            assertThat(underTest.isKeyguardShowing()).isTrue()
 
-        whenever(keyguardStateController.isShowing).thenReturn(false)
+            whenever(keyguardStateController.isShowing).thenReturn(false)
-        captor.value.onKeyguardShowingChanged()
+            captor.value.onKeyguardShowingChanged()
-        assertThat(latest).isFalse()
+            assertThat(latest).isFalse()
-        assertThat(underTest.isKeyguardShowing()).isFalse()
+            assertThat(underTest.isKeyguardShowing()).isFalse()
 
-        job.cancel()
+            job.cancel()
-    }
+        }
 
     @Test
-    fun isDozing() = runBlockingTest {
+    fun isDozing() =
-        var latest: Boolean? = null
+        runTest(UnconfinedTestDispatcher()) {
-        val job = underTest.isDozing.onEach { latest = it }.launchIn(this)
+            var latest: Boolean? = null
+            val job = underTest.isDozing.onEach { latest = it }.launchIn(this)
 
-        val captor = argumentCaptor<DozeHost.Callback>()
+            val captor = argumentCaptor<DozeHost.Callback>()
-        verify(dozeHost).addCallback(captor.capture())
+            verify(dozeHost).addCallback(captor.capture())
 
-        captor.value.onDozingChanged(true)
+            captor.value.onDozingChanged(true)
-        assertThat(latest).isTrue()
+            assertThat(latest).isTrue()
 
-        captor.value.onDozingChanged(false)
+            captor.value.onDozingChanged(false)
-        assertThat(latest).isFalse()
+            assertThat(latest).isFalse()
 
-        job.cancel()
+            job.cancel()
-        verify(dozeHost).removeCallback(captor.value)
+            verify(dozeHost).removeCallback(captor.value)
-    }
+        }
 
     @Test
-    fun `isDozing - starts with correct initial value for isDozing`() = runBlockingTest {
+    fun `isDozing - starts with correct initial value for isDozing`() =
-        var latest: Boolean? = null
+        runTest(UnconfinedTestDispatcher()) {
+            var latest: Boolean? = null
 
-        whenever(statusBarStateController.isDozing).thenReturn(true)
+            whenever(statusBarStateController.isDozing).thenReturn(true)
-        var job = underTest.isDozing.onEach { latest = it }.launchIn(this)
+            var job = underTest.isDozing.onEach { latest = it }.launchIn(this)
-        assertThat(latest).isTrue()
+            assertThat(latest).isTrue()
-        job.cancel()
+            job.cancel()
 
-        whenever(statusBarStateController.isDozing).thenReturn(false)
+            whenever(statusBarStateController.isDozing).thenReturn(false)
-        job = underTest.isDozing.onEach { latest = it }.launchIn(this)
+            job = underTest.isDozing.onEach { latest = it }.launchIn(this)
-        assertThat(latest).isFalse()
+            assertThat(latest).isFalse()
-        job.cancel()
+            job.cancel()
-    }
+        }
 
     @Test
-    fun dozeAmount() = runBlockingTest {
+    fun dozeAmount() =
-        val values = mutableListOf<Float>()
+        runTest(UnconfinedTestDispatcher()) {
-        val job = underTest.dozeAmount.onEach(values::add).launchIn(this)
+            val values = mutableListOf<Float>()
+            val job = underTest.dozeAmount.onEach(values::add).launchIn(this)
 
-        val captor = argumentCaptor<StatusBarStateController.StateListener>()
+            val captor = argumentCaptor<StatusBarStateController.StateListener>()
-        verify(statusBarStateController).addCallback(captor.capture())
+            verify(statusBarStateController).addCallback(captor.capture())
 
-        captor.value.onDozeAmountChanged(0.433f, 0.4f)
+            captor.value.onDozeAmountChanged(0.433f, 0.4f)
-        captor.value.onDozeAmountChanged(0.498f, 0.5f)
+            captor.value.onDozeAmountChanged(0.498f, 0.5f)
-        captor.value.onDozeAmountChanged(0.661f, 0.65f)
+            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()
+            job.cancel()
-        verify(statusBarStateController).removeCallback(captor.value)
+            verify(statusBarStateController).removeCallback(captor.value)
-    }
+        }
 
     @Test
-    fun wakefulness() = runBlockingTest {
+    fun wakefulness() =
-        val values = mutableListOf<WakefulnessModel>()
+        runTest(UnconfinedTestDispatcher()) {
-        val job = underTest.wakefulnessState.onEach(values::add).launchIn(this)
+            val values = mutableListOf<WakefulnessModel>()
+            val job = underTest.wakefulnessState.onEach(values::add).launchIn(this)
 
-        val captor = argumentCaptor<WakefulnessLifecycle.Observer>()
+            val captor = argumentCaptor<WakefulnessLifecycle.Observer>()
-        verify(wakefulnessLifecycle).addObserver(captor.capture())
+            verify(wakefulnessLifecycle).addObserver(captor.capture())
 
-        captor.value.onStartedWakingUp()
+            captor.value.onStartedWakingUp()
-        captor.value.onFinishedWakingUp()
+            captor.value.onFinishedWakingUp()
-        captor.value.onStartedGoingToSleep()
+            captor.value.onStartedGoingToSleep()
-        captor.value.onFinishedGoingToSleep()
+            captor.value.onFinishedGoingToSleep()
 
-        assertThat(values)
+            assertThat(values)
-            .isEqualTo(
+                .isEqualTo(
-                listOf(
+                    listOf(
-                    // Initial value will be ASLEEP
+                        // Initial value will be ASLEEP
-                    WakefulnessModel.ASLEEP,
+                        WakefulnessModel.ASLEEP,
-                    WakefulnessModel.STARTING_TO_WAKE,
+                        WakefulnessModel.STARTING_TO_WAKE,
-                    WakefulnessModel.AWAKE,
+                        WakefulnessModel.AWAKE,
-                    WakefulnessModel.STARTING_TO_SLEEP,
+                        WakefulnessModel.STARTING_TO_SLEEP,
-                    WakefulnessModel.ASLEEP,
+                        WakefulnessModel.ASLEEP,
+                    )
                 )
-            )
 
-        job.cancel()
+            job.cancel()
-        verify(wakefulnessLifecycle).removeObserver(captor.value)
+            verify(wakefulnessLifecycle).removeObserver(captor.value)
-    }
+        }
 
     @Test
-    fun isUdfpsSupported() = runBlockingTest {
+    fun isUdfpsSupported() =
-        whenever(keyguardUpdateMonitor.isUdfpsSupported).thenReturn(true)
+        runTest(UnconfinedTestDispatcher()) {
-        assertThat(underTest.isUdfpsSupported()).isTrue()
+            whenever(keyguardUpdateMonitor.isUdfpsSupported).thenReturn(true)
+            assertThat(underTest.isUdfpsSupported()).isTrue()
 
-        whenever(keyguardUpdateMonitor.isUdfpsSupported).thenReturn(false)
+            whenever(keyguardUpdateMonitor.isUdfpsSupported).thenReturn(false)
-        assertThat(underTest.isUdfpsSupported()).isFalse()
+            assertThat(underTest.isUdfpsSupported()).isFalse()
-    }
+        }
 
     @Test
-    fun isBouncerShowing() = runBlockingTest {
+    fun isBouncerShowing() =
-        whenever(keyguardStateController.isBouncerShowing).thenReturn(false)
+        runTest(UnconfinedTestDispatcher()) {
-        var latest: Boolean? = null
+            whenever(keyguardStateController.isBouncerShowing).thenReturn(false)
-        val job = underTest.isBouncerShowing.onEach { latest = it }.launchIn(this)
+            var latest: Boolean? = null
+            val job = underTest.isBouncerShowing.onEach { latest = it }.launchIn(this)
 
-        assertThat(latest).isFalse()
+            assertThat(latest).isFalse()
 
-        val captor = argumentCaptor<KeyguardStateController.Callback>()
+            val captor = argumentCaptor<KeyguardStateController.Callback>()
-        verify(keyguardStateController).addCallback(captor.capture())
+            verify(keyguardStateController).addCallback(captor.capture())
 
-        whenever(keyguardStateController.isBouncerShowing).thenReturn(true)
+            whenever(keyguardStateController.isBouncerShowing).thenReturn(true)
-        captor.value.onBouncerShowingChanged()
+            captor.value.onBouncerShowingChanged()
-        assertThat(latest).isTrue()
+            assertThat(latest).isTrue()
 
-        whenever(keyguardStateController.isBouncerShowing).thenReturn(false)
+            whenever(keyguardStateController.isBouncerShowing).thenReturn(false)
-        captor.value.onBouncerShowingChanged()
+            captor.value.onBouncerShowingChanged()
-        assertThat(latest).isFalse()
+            assertThat(latest).isFalse()
 
-        job.cancel()
+            job.cancel()
-    }
+        }
 
     @Test
-    fun isKeyguardGoingAway() = runBlockingTest {
+    fun isKeyguardGoingAway() =
-        whenever(keyguardStateController.isKeyguardGoingAway).thenReturn(false)
+        runTest(UnconfinedTestDispatcher()) {
-        var latest: Boolean? = null
+            whenever(keyguardStateController.isKeyguardGoingAway).thenReturn(false)
-        val job = underTest.isKeyguardGoingAway.onEach { latest = it }.launchIn(this)
+            var latest: Boolean? = null
+            val job = underTest.isKeyguardGoingAway.onEach { latest = it }.launchIn(this)
 
-        assertThat(latest).isFalse()
+            assertThat(latest).isFalse()
 
-        val captor = argumentCaptor<KeyguardStateController.Callback>()
+            val captor = argumentCaptor<KeyguardStateController.Callback>()
-        verify(keyguardStateController).addCallback(captor.capture())
+            verify(keyguardStateController).addCallback(captor.capture())
 
-        whenever(keyguardStateController.isKeyguardGoingAway).thenReturn(true)
+            whenever(keyguardStateController.isKeyguardGoingAway).thenReturn(true)
-        captor.value.onKeyguardGoingAwayChanged()
+            captor.value.onKeyguardGoingAwayChanged()
-        assertThat(latest).isTrue()
+            assertThat(latest).isTrue()
 
-        whenever(keyguardStateController.isKeyguardGoingAway).thenReturn(false)
+            whenever(keyguardStateController.isKeyguardGoingAway).thenReturn(false)
-        captor.value.onKeyguardGoingAwayChanged()
+            captor.value.onKeyguardGoingAwayChanged()
-        assertThat(latest).isFalse()
+            assertThat(latest).isFalse()
 
-        job.cancel()
+            job.cancel()
-    }
+        }
 
     @Test
-    fun isDreaming() = runBlockingTest {
+    fun isDreaming() =
-        whenever(keyguardUpdateMonitor.isDreaming()).thenReturn(false)
+        runTest(UnconfinedTestDispatcher()) {
-        var latest: Boolean? = null
+            whenever(keyguardUpdateMonitor.isDreaming()).thenReturn(false)
-        val job = underTest.isDreaming.onEach { latest = it }.launchIn(this)
+            var latest: Boolean? = null
+            val job = underTest.isDreaming.onEach { latest = it }.launchIn(this)
 
-        assertThat(latest).isFalse()
+            assertThat(latest).isFalse()
 
-        val captor = argumentCaptor<KeyguardUpdateMonitorCallback>()
+            val captor = argumentCaptor<KeyguardUpdateMonitorCallback>()
-        verify(keyguardUpdateMonitor).registerCallback(captor.capture())
+            verify(keyguardUpdateMonitor).registerCallback(captor.capture())
 
-        captor.value.onDreamingStateChanged(true)
+            captor.value.onDreamingStateChanged(true)
-        assertThat(latest).isTrue()
+            assertThat(latest).isTrue()
 
-        captor.value.onDreamingStateChanged(false)
+            captor.value.onDreamingStateChanged(false)
-        assertThat(latest).isFalse()
+            assertThat(latest).isFalse()
 
-        job.cancel()
+            job.cancel()
-    }
+        }
 
     @Test
-    fun biometricUnlockState() = runBlockingTest {
+    fun biometricUnlockState() =
-        val values = mutableListOf<BiometricUnlockModel>()
+        runTest(UnconfinedTestDispatcher()) {
-        val job = underTest.biometricUnlockState.onEach(values::add).launchIn(this)
+            val values = mutableListOf<BiometricUnlockModel>()
+            val job = underTest.biometricUnlockState.onEach(values::add).launchIn(this)
 
-        val captor = argumentCaptor<BiometricUnlockController.BiometricModeListener>()
+            val captor = argumentCaptor<BiometricUnlockController.BiometricModeListener>()
-        verify(biometricUnlockController).addBiometricModeListener(captor.capture())
+            verify(biometricUnlockController).addBiometricModeListener(captor.capture())
 
-        captor.value.onModeChanged(BiometricUnlockController.MODE_NONE)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_NONE)
-        captor.value.onModeChanged(BiometricUnlockController.MODE_WAKE_AND_UNLOCK)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_WAKE_AND_UNLOCK)
-        captor.value.onModeChanged(BiometricUnlockController.MODE_WAKE_AND_UNLOCK_PULSING)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_WAKE_AND_UNLOCK_PULSING)
-        captor.value.onModeChanged(BiometricUnlockController.MODE_SHOW_BOUNCER)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_SHOW_BOUNCER)
-        captor.value.onModeChanged(BiometricUnlockController.MODE_ONLY_WAKE)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_ONLY_WAKE)
-        captor.value.onModeChanged(BiometricUnlockController.MODE_UNLOCK_COLLAPSING)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_UNLOCK_COLLAPSING)
-        captor.value.onModeChanged(BiometricUnlockController.MODE_DISMISS_BOUNCER)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_DISMISS_BOUNCER)
-        captor.value.onModeChanged(BiometricUnlockController.MODE_WAKE_AND_UNLOCK_FROM_DREAM)
+            captor.value.onModeChanged(BiometricUnlockController.MODE_WAKE_AND_UNLOCK_FROM_DREAM)
 
-        assertThat(values)
+            assertThat(values)
-            .isEqualTo(
+                .isEqualTo(
-                listOf(
+                    listOf(
-                    // Initial value will be NONE, followed by onModeChanged() call
+                        // Initial value will be NONE, followed by onModeChanged() call
-                    BiometricUnlockModel.NONE,
+                        BiometricUnlockModel.NONE,
-                    BiometricUnlockModel.NONE,
+                        BiometricUnlockModel.NONE,
-                    BiometricUnlockModel.WAKE_AND_UNLOCK,
+                        BiometricUnlockModel.WAKE_AND_UNLOCK,
-                    BiometricUnlockModel.WAKE_AND_UNLOCK_PULSING,
+                        BiometricUnlockModel.WAKE_AND_UNLOCK_PULSING,
-                    BiometricUnlockModel.SHOW_BOUNCER,
+                        BiometricUnlockModel.SHOW_BOUNCER,
-                    BiometricUnlockModel.ONLY_WAKE,
+                        BiometricUnlockModel.ONLY_WAKE,
-                    BiometricUnlockModel.UNLOCK_COLLAPSING,
+                        BiometricUnlockModel.UNLOCK_COLLAPSING,
-                    BiometricUnlockModel.DISMISS_BOUNCER,
+                        BiometricUnlockModel.DISMISS_BOUNCER,
-                    BiometricUnlockModel.WAKE_AND_UNLOCK_FROM_DREAM,
+                        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()
+            assertThat(values)
-        verify(biometricUnlockController).removeBiometricModeListener(captor.value)
+                .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)
+        }
 }

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