Merge changes If9501500,I321d8023 into main

* changes:
  Smooth burn-in transitions
  Fix AOD icon flicker during GONE->AOD

13 files changed, 171 insertions, 33 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelTest.kt
index 0543bc257440..d52696a0bd87 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelTest.kt
@@ -20,6 +20,7 @@ package com.android.systemui.keyguard.ui.viewmodel
 
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.systemui.Flags as AConfigFlags
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
@@ -68,6 +69,8 @@ class AodBurnInViewModelTest : SysuiTestCase() {
 
     @Before
     fun setUp() {
+        mSetFlagsRule.disableFlags(AConfigFlags.FLAG_MIGRATE_CLOCKS_TO_BLUEPRINT)
+
         MockitoAnnotations.initMocks(this)
         whenever(burnInInteractor.keyguardBurnIn).thenReturn(burnInFlow)
         kosmos.burnInInteractor = burnInInteractor
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelTest.kt
index 2de013bc7abc..c23ec2290d6a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelTest.kt
@@ -116,6 +116,23 @@ class KeyguardRootViewModelTest : SysuiTestCase() {
         }
 
     @Test
+    fun iconContainer_isNotVisible_onKeyguard_dontShowWhenGoneToAodTransitionRunning() =
+        testScope.runTest {
+            val isVisible by collectLastValue(underTest.isNotifIconContainerVisible)
+            runCurrent()
+            keyguardTransitionRepository.sendTransitionSteps(
+                from = KeyguardState.GONE,
+                to = KeyguardState.AOD,
+                testScope,
+            )
+            whenever(screenOffAnimationController.shouldShowAodIconsWhenShade()).thenReturn(false)
+            runCurrent()
+
+            assertThat(isVisible?.value).isFalse()
+            assertThat(isVisible?.isAnimating).isFalse()
+        }
+
+    @Test
     fun iconContainer_isVisible_bypassEnabled() =
         testScope.runTest {
             val isVisible by collectLastValue(underTest.isNotifIconContainerVisible)
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/BurnInInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/BurnInInteractor.kt
index cc1cf911f1c1..7ae70a9a3e7c 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/BurnInInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/BurnInInteractor.kt
@@ -74,6 +74,7 @@ constructor(
                 BurnInModel(translationX, translationY, burnInHelperWrapper.burnInScale())
             }
             .distinctUntilChanged()
+            .stateIn(scope, SharingStarted.Lazily, BurnInModel())
 
     /**
      * Use for max burn-in offsets that are NOT specified in pixels. This flow will recalculate the
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt
index 5606d4301cfa..e0b5c0e1f4c6 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt
@@ -98,6 +98,8 @@ constructor(
                         val modeOnCanceled =
                             if (lastStartedStep.from == KeyguardState.LOCKSCREEN) {
                                 TransitionModeOnCanceled.REVERSE
+                            } else if (lastStartedStep.from == KeyguardState.GONE) {
+                                TransitionModeOnCanceled.RESET
                             } else {
                                 TransitionModeOnCanceled.LAST_VALUE
                             }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlow.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlow.kt
index 00b798901352..8b278cdb9a6c 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlow.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlow.kt
@@ -112,6 +112,38 @@ constructor(
             interpolator: Interpolator = LINEAR,
             name: String? = null
         ): Flow<Float> {
+            return sharedFlowWithState(
+                    duration = duration,
+                    onStep = onStep,
+                    startTime = startTime,
+                    onStart = onStart,
+                    onCancel = onCancel,
+                    onFinish = onFinish,
+                    interpolator = interpolator,
+                    name = name,
+                )
+                .mapNotNull { stateToValue -> stateToValue.value }
+        }
+
+        /**
+         * Transitions will occur over a [transitionDuration] with [TransitionStep]s being emitted
+         * in the range of [0, 1]. View animations should begin and end within a subset of this
+         * range. This function maps the [startTime] and [duration] into [0, 1], when this subset is
+         * valid.
+         *
+         * Will return a [StateToValue], which encompasses the calculated value as well as the
+         * transitionState that is associated with it.
+         */
+        fun sharedFlowWithState(
+            duration: Duration,
+            onStep: (Float) -> Float,
+            startTime: Duration = 0.milliseconds,
+            onStart: (() -> Unit)? = null,
+            onCancel: (() -> Float)? = null,
+            onFinish: (() -> Float)? = null,
+            interpolator: Interpolator = LINEAR,
+            name: String? = null
+        ): Flow<StateToValue> {
             if (!duration.isPositive()) {
                 throw IllegalArgumentException("duration must be a positive number: $duration")
             }
@@ -164,7 +196,6 @@ constructor(
                         .also { logger.logTransitionStep(name, step, it.value) }
                 }
                 .distinctUntilChanged()
-                .mapNotNull { stateToValue -> stateToValue.value }
         }
 
         /**
@@ -174,9 +205,9 @@ constructor(
             return sharedFlow(duration = 1.milliseconds, onStep = { value }, onFinish = { value })
         }
     }
-
-    data class StateToValue(
-        val transitionState: TransitionState,
-        val value: Float?,
-    )
 }
+
+data class StateToValue(
+    val transitionState: TransitionState = TransitionState.FINISHED,
+    val value: Float? = 0f,
+)
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt
index 789d30ff1a31..9e7c70d67156 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt
@@ -145,9 +145,7 @@ object KeyguardRootViewBinder {
                         launch {
                             viewModel.burnInLayerVisibility.collect { visibility ->
                                 childViews[burnInLayerId]?.visibility = visibility
-                                // Reset alpha only for the icons, as they currently have their
-                                // own animator
-                                childViews[aodNotificationIconContainerId]?.alpha = 0f
+                                childViews[aodNotificationIconContainerId]?.visibility = visibility
                             }
                         }
 
@@ -313,6 +311,12 @@ object KeyguardRootViewBinder {
             }
         }
 
+        if (KeyguardShadeMigrationNssl.isEnabled) {
+            burnInParams.update { current ->
+                current.copy(translationY = { childViews[burnInLayerId]?.translationY })
+            }
+        }
+
         onLayoutChangeListener = OnLayoutChange(viewModel, burnInParams)
         view.addOnLayoutChangeListener(onLayoutChangeListener)
 
@@ -435,11 +439,17 @@ object KeyguardRootViewBinder {
             }
         when {
             !isVisible.isAnimating -> {
-                alpha = 1f
                 if (!KeyguardShadeMigrationNssl.isEnabled) {
                     translationY = 0f
                 }
-                visibility = if (isVisible.value) View.VISIBLE else View.INVISIBLE
+                visibility =
+                    if (isVisible.value) {
+                        alpha = 1f
+                        View.VISIBLE
+                    } else {
+                        alpha = 0f
+                        View.INVISIBLE
+                    }
             }
             newAodTransition() -> {
                 animateInIconTranslation()
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModel.kt
index 828e03301b8e..8110de23be13 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModel.kt
@@ -31,6 +31,10 @@ import com.android.systemui.keyguard.shared.model.BurnInModel
 import com.android.systemui.keyguard.shared.model.KeyguardState.ALTERNATE_BOUNCER
 import com.android.systemui.keyguard.shared.model.KeyguardState.AOD
 import com.android.systemui.keyguard.shared.model.KeyguardState.GONE
+import com.android.systemui.keyguard.shared.model.TransitionState
+import com.android.systemui.keyguard.shared.model.TransitionState.RUNNING
+import com.android.systemui.keyguard.shared.model.TransitionState.STARTED
+import com.android.systemui.keyguard.ui.StateToValue
 import com.android.systemui.plugins.clocks.ClockController
 import com.android.systemui.res.R
 import javax.inject.Inject
@@ -58,6 +62,7 @@ constructor(
     private val keyguardInteractor: KeyguardInteractor,
     private val keyguardTransitionInteractor: KeyguardTransitionInteractor,
     private val goneToAodTransitionViewModel: GoneToAodTransitionViewModel,
+    private val aodToLockscreenTransitionViewModel: AodToLockscreenTransitionViewModel,
     private val occludedToLockscreenTransitionViewModel: OccludedToLockscreenTransitionViewModel,
     private val keyguardClockViewModel: KeyguardClockViewModel,
 ) {
@@ -83,21 +88,22 @@ constructor(
                     burnIn(params).map { it.translationY.toFloat() }.onStart { emit(0f) },
                     goneToAodTransitionViewModel
                         .enterFromTopTranslationY(enterFromTopAmount)
-                        .onStart { emit(0f) },
+                        .onStart { emit(StateToValue()) },
                     occludedToLockscreenTransitionViewModel.lockscreenTranslationY.onStart {
                         emit(0f)
                     },
-                ) {
-                    keyguardTransitionY,
-                    burnInTranslationY,
-                    goneToAodTransitionTranslationY,
-                    occludedToLockscreenTransitionTranslationY ->
-
-                    // All values need to be combined for a smooth translation
-                    keyguardTransitionY +
-                        burnInTranslationY +
-                        goneToAodTransitionTranslationY +
-                        occludedToLockscreenTransitionTranslationY
+                    aodToLockscreenTransitionViewModel.translationY(params.translationY).onStart {
+                        emit(StateToValue())
+                    },
+                ) { keyguardTranslationY, burnInY, goneToAod, occludedToLockscreen, aodToLockscreen
+                    ->
+                    if (isInTransition(aodToLockscreen.transitionState)) {
+                        aodToLockscreen.value ?: 0f
+                    } else if (isInTransition(goneToAod.transitionState)) {
+                        (goneToAod.value ?: 0f) + burnInY
+                    } else {
+                        burnInY + occludedToLockscreen + keyguardTranslationY
+                    }
                 }
             }
             .distinctUntilChanged()
@@ -115,6 +121,10 @@ constructor(
         }
     }
 
+    private fun isInTransition(state: TransitionState): Boolean {
+        return state == STARTED || state == RUNNING
+    }
+
     private fun burnIn(
         params: BurnInParameters,
     ): Flow<BurnInModel> {
@@ -185,6 +195,8 @@ data class BurnInParameters(
     val topInset: Int = 0,
     /** Status view top, without translation added in */
     val statusViewTop: Int = 0,
+    /** The current y translation of the view */
+    val translationY: () -> Float? = { null }
 )
 
 /**
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodToLockscreenTransitionViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodToLockscreenTransitionViewModel.kt
index 266fd02d5bbf..6d1d3cb00a68 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodToLockscreenTransitionViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodToLockscreenTransitionViewModel.kt
@@ -16,11 +16,14 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import android.util.MathUtils
+import com.android.app.animation.Interpolators.FAST_OUT_SLOW_IN
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryUdfpsInteractor
 import com.android.systemui.keyguard.domain.interactor.FromAodTransitionInteractor.Companion.TO_LOCKSCREEN_DURATION
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.ui.KeyguardTransitionAnimationFlow
+import com.android.systemui.keyguard.ui.StateToValue
 import com.android.systemui.keyguard.ui.transitions.DeviceEntryIconTransition
 import javax.inject.Inject
 import kotlin.time.Duration.Companion.milliseconds
@@ -48,6 +51,22 @@ constructor(
             to = KeyguardState.LOCKSCREEN,
         )
 
+    /**
+     * Begin the transition from wherever the y-translation value is currently. This helps ensure a
+     * smooth transition if a transition in canceled.
+     */
+    fun translationY(currentTranslationY: () -> Float?): Flow<StateToValue> {
+        var startValue = 0f
+        return transitionAnimation.sharedFlowWithState(
+            duration = 500.milliseconds,
+            onStart = {
+                startValue = currentTranslationY() ?: 0f
+                startValue
+            },
+            onStep = { MathUtils.lerp(startValue, 0f, FAST_OUT_SLOW_IN.getInterpolation(it)) },
+        )
+    }
+
     /** Ensure alpha is set to be visible */
     val lockscreenAlpha: Flow<Float> =
         transitionAnimation.sharedFlow(
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModel.kt
index f5e61355df37..85885b065264 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModel.kt
@@ -22,6 +22,7 @@ import com.android.systemui.deviceentry.domain.interactor.DeviceEntryUdfpsIntera
 import com.android.systemui.keyguard.domain.interactor.FromGoneTransitionInteractor.Companion.TO_AOD_DURATION
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.ui.KeyguardTransitionAnimationFlow
+import com.android.systemui.keyguard.ui.StateToValue
 import com.android.systemui.keyguard.ui.transitions.DeviceEntryIconTransition
 import javax.inject.Inject
 import kotlin.time.Duration.Companion.milliseconds
@@ -48,8 +49,8 @@ constructor(
         )
 
     /** y-translation from the top of the screen for AOD */
-    fun enterFromTopTranslationY(translatePx: Int): Flow<Float> {
-        return transitionAnimation.sharedFlow(
+    fun enterFromTopTranslationY(translatePx: Int): Flow<StateToValue> {
+        return transitionAnimation.sharedFlowWithState(
             startTime = 600.milliseconds,
             duration = 500.milliseconds,
             onStart = { translatePx },
@@ -63,8 +64,8 @@ constructor(
     /** alpha animation upon entering AOD */
     val enterFromTopAnimationAlpha: Flow<Float> =
         transitionAnimation.sharedFlow(
-            startTime = 600.milliseconds,
-            duration = 500.milliseconds,
+            startTime = 700.milliseconds,
+            duration = 400.milliseconds,
             onStart = { 0f },
             onStep = { it },
             onFinish = { 1f },
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
index f8a12bd226ad..ec13228c6216 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
@@ -30,6 +30,8 @@ import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.KeyguardState.AOD
 import com.android.systemui.keyguard.shared.model.KeyguardState.GONE
 import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
+import com.android.systemui.keyguard.shared.model.TransitionState.RUNNING
+import com.android.systemui.keyguard.shared.model.TransitionState.STARTED
 import com.android.systemui.statusbar.notification.domain.interactor.NotificationsKeyguardInteractor
 import com.android.systemui.statusbar.phone.DozeParameters
 import com.android.systemui.statusbar.phone.ScreenOffAnimationController
@@ -48,6 +50,7 @@ import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.merge
+import kotlinx.coroutines.flow.onStart
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SysUISingleton
@@ -78,6 +81,12 @@ constructor(
 
     val goneToAodTransition = keyguardTransitionInteractor.transition(from = GONE, to = AOD)
 
+    private val goneToAodTransitionRunning: Flow<Boolean> =
+        goneToAodTransition
+            .map { it.transitionState == STARTED || it.transitionState == RUNNING }
+            .onStart { emit(false) }
+            .distinctUntilChanged()
+
     /** Last point that the root view was tapped */
     val lastRootViewTapPosition: Flow<Point?> = keyguardInteractor.lastRootViewTapPosition
 
@@ -138,6 +147,7 @@ constructor(
     /** Is the notification icon container visible? */
     val isNotifIconContainerVisible: Flow<AnimatedValue<Boolean>> =
         combine(
+                goneToAodTransitionRunning,
                 keyguardTransitionInteractor.finishedKeyguardState.map {
                     KeyguardState.lockscreenVisibleInState(it)
                 },
@@ -145,6 +155,7 @@ constructor(
                 areNotifsFullyHiddenAnimated(),
                 isPulseExpandingAnimated(),
             ) {
+                goneToAodTransitionRunning: Boolean,
                 onKeyguard: Boolean,
                 isBypassEnabled: Boolean,
                 notifsFullyHidden: AnimatedValue<Boolean>,
@@ -154,7 +165,9 @@ constructor(
                     // Hide the AOD icons if we're not in the KEYGUARD state unless the screen off
                     // animation is playing, in which case we want them to be visible if we're
                     // animating in the AOD UI and will be switching to KEYGUARD shortly.
-                    !onKeyguard && !screenOffAnimationController.shouldShowAodIconsWhenShade() ->
+                    goneToAodTransitionRunning ||
+                        (!onKeyguard &&
+                            !screenOffAnimationController.shouldShowAodIconsWhenShade()) ->
                         AnimatedValue.NotAnimating(false)
                     else ->
                         zip(notifsFullyHidden, pulseExpanding) {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlowTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlowTest.kt
index 2d9d5ed2b5e1..0e9197ef8ac1 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlowTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlowTest.kt
@@ -181,6 +181,31 @@ class KeyguardTransitionAnimationFlowTest : SysuiTestCase() {
         }
 
     @Test
+    fun usesOnStepToDoubleValueWithState() =
+        testScope.runTest {
+            val flow =
+                underTest.sharedFlowWithState(
+                    duration = 1000.milliseconds,
+                    onStep = { it * 2 },
+                )
+            val animationValues by collectLastValue(flow)
+            runCurrent()
+
+            repository.sendTransitionStep(step(0f, TransitionState.STARTED))
+            assertThat(animationValues).isEqualTo(StateToValue(TransitionState.STARTED, 0f))
+            repository.sendTransitionStep(step(0.3f, TransitionState.RUNNING))
+            assertThat(animationValues).isEqualTo(StateToValue(TransitionState.RUNNING, 0.6f))
+            repository.sendTransitionStep(step(0.6f, TransitionState.RUNNING))
+            assertThat(animationValues).isEqualTo(StateToValue(TransitionState.RUNNING, 1.2f))
+            repository.sendTransitionStep(step(0.8f, TransitionState.RUNNING))
+            assertThat(animationValues).isEqualTo(StateToValue(TransitionState.RUNNING, 1.6f))
+            repository.sendTransitionStep(step(1f, TransitionState.RUNNING))
+            assertThat(animationValues).isEqualTo(StateToValue(TransitionState.RUNNING, 2f))
+            repository.sendTransitionStep(step(1f, TransitionState.FINISHED))
+            assertThat(animationValues).isEqualTo(StateToValue(TransitionState.FINISHED, null))
+        }
+
+    @Test
     fun sameFloatValueWithTheSameTransitionStateDoesNotEmitTwice() =
         testScope.runTest {
             val flow =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModelTest.kt
index 1c9c942eafc6..bfa84335d670 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModelTest.kt
@@ -26,6 +26,7 @@ import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepos
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
+import com.android.systemui.keyguard.ui.StateToValue
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.testKosmos
 import com.google.common.collect.Range
@@ -57,17 +58,19 @@ class GoneToAodTransitionViewModelTest : SysuiTestCase() {
 
             // The animation should only start > .4f way through
             repository.sendTransitionStep(step(0f, TransitionState.STARTED))
-            assertThat(enterFromTopTranslationY).isEqualTo(pixels)
+            assertThat(enterFromTopTranslationY)
+                .isEqualTo(StateToValue(TransitionState.STARTED, pixels))
 
-            repository.sendTransitionStep(step(0.4f))
-            assertThat(enterFromTopTranslationY).isEqualTo(pixels)
+            repository.sendTransitionStep(step(.55f))
+            assertThat(enterFromTopTranslationY!!.value ?: -1f).isIn(Range.closed(pixels, 0f))
 
             repository.sendTransitionStep(step(.85f))
-            assertThat(enterFromTopTranslationY).isIn(Range.closed(pixels, 0f))
+            assertThat(enterFromTopTranslationY!!.value ?: -1f).isIn(Range.closed(pixels, 0f))
 
             // At the end, the translation should be complete and set to zero
             repository.sendTransitionStep(step(1f))
-            assertThat(enterFromTopTranslationY).isEqualTo(0f)
+            assertThat(enterFromTopTranslationY)
+                .isEqualTo(StateToValue(TransitionState.RUNNING, 0f))
         }
 
     @Test
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelKosmos.kt
index a8f45b0974c4..6f168d47038d 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelKosmos.kt
@@ -33,6 +33,7 @@ var Kosmos.aodBurnInViewModel by Fixture {
         keyguardInteractor = keyguardInteractor,
         keyguardTransitionInteractor = keyguardTransitionInteractor,
         goneToAodTransitionViewModel = goneToAodTransitionViewModel,
+        aodToLockscreenTransitionViewModel = aodToLockscreenTransitionViewModel,
         occludedToLockscreenTransitionViewModel = occludedToLockscreenTransitionViewModel,
         keyguardClockViewModel = keyguardClockViewModel,
     )