Merge changes I9f79851d,I91578f20 into main

* changes:
  Move asleepKeyguardState to KeyguardInteractor
  [Flexiglass] Fix SharedNotificationContainerViewModel

12 files changed, 327 insertions, 184 deletions

diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
index 236e202749b2..a82ee4c359a3 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
@@ -190,7 +190,7 @@ sealed interface ObservableTransitionState {
         }
     }
 
-    fun isIdle(scene: SceneKey?): Boolean {
+    fun isIdle(scene: SceneKey? = null): Boolean {
         return this is Idle && (scene == null || this.currentScene == scene)
     }
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
index 8f8d2e22e161..d3b51d1d17f7 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
@@ -51,7 +51,7 @@ import com.android.systemui.keyguard.data.repository.keyguardRepository
 import com.android.systemui.keyguard.data.repository.keyguardTransitionRepository
 import com.android.systemui.keyguard.dismissCallbackRegistry
 import com.android.systemui.keyguard.domain.interactor.keyguardEnabledInteractor
-import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
+import com.android.systemui.keyguard.domain.interactor.keyguardInteractor
 import com.android.systemui.keyguard.domain.interactor.scenetransition.lockscreenSceneTransitionInteractor
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.SuccessFingerprintAuthenticationStatus
@@ -551,8 +551,7 @@ class SceneContainerStartableTest : SysuiTestCase() {
     fun switchToAOD_whenAvailable_whenDeviceSleepsLocked() =
         testScope.runTest {
             kosmos.lockscreenSceneTransitionInteractor.start()
-            val asleepState by
-                collectLastValue(kosmos.keyguardTransitionInteractor.asleepKeyguardState)
+            val asleepState by collectLastValue(kosmos.keyguardInteractor.asleepKeyguardState)
             val currentTransitionInfo by
                 collectLastValue(kosmos.keyguardTransitionRepository.currentTransitionInfoInternal)
             val transitionState =
@@ -584,8 +583,7 @@ class SceneContainerStartableTest : SysuiTestCase() {
     fun switchToDozing_whenAodUnavailable_whenDeviceSleepsLocked() =
         testScope.runTest {
             kosmos.lockscreenSceneTransitionInteractor.start()
-            val asleepState by
-                collectLastValue(kosmos.keyguardTransitionInteractor.asleepKeyguardState)
+            val asleepState by collectLastValue(kosmos.keyguardInteractor.asleepKeyguardState)
             val currentTransitionInfo by
                 collectLastValue(kosmos.keyguardTransitionRepository.currentTransitionInfoInternal)
             val transitionState =
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
index 733cac99f4ec..3f97f0b7a67d 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
@@ -42,10 +42,14 @@ import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.domain.interactor.keyguardInteractor
 import com.android.systemui.keyguard.shared.model.BurnInModel
-import com.android.systemui.keyguard.shared.model.KeyguardState
+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.DREAMING
+import com.android.systemui.keyguard.shared.model.KeyguardState.GLANCEABLE_HUB
 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.KeyguardState.OCCLUDED
+import com.android.systemui.keyguard.shared.model.KeyguardState.PRIMARY_BOUNCER
 import com.android.systemui.keyguard.shared.model.StatusBarState
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
@@ -56,6 +60,10 @@ import com.android.systemui.keyguard.ui.viewmodel.aodBurnInViewModel
 import com.android.systemui.keyguard.ui.viewmodel.keyguardRootViewModel
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.res.R
+import com.android.systemui.scene.data.repository.Idle
+import com.android.systemui.scene.data.repository.Transition
+import com.android.systemui.scene.data.repository.setTransition
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shade.data.repository.fakeShadeRepository
 import com.android.systemui.shade.mockLargeScreenHeaderHelper
 import com.android.systemui.shade.shadeTestUtil
@@ -66,6 +74,7 @@ import com.google.common.collect.Range
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.advanceTimeBy
 import kotlinx.coroutines.test.runCurrent
@@ -295,34 +304,47 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
 
             // Start transitioning to glanceable hub
             val progress = 0.6f
-            keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(
-                    transitionState = TransitionState.STARTED,
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GLANCEABLE_HUB,
-                    value = 0f,
-                )
+            kosmos.setTransition(
+                sceneTransition = Transition(from = Scenes.Lockscreen, to = Scenes.Communal),
+                stateTransition =
+                    TransitionStep(
+                        transitionState = TransitionState.STARTED,
+                        from = LOCKSCREEN,
+                        to = GLANCEABLE_HUB,
+                        value = 0f,
+                    )
             )
+
             runCurrent()
-            keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(
-                    transitionState = TransitionState.RUNNING,
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GLANCEABLE_HUB,
-                    value = progress,
-                )
+            kosmos.setTransition(
+                sceneTransition =
+                    Transition(
+                        from = Scenes.Lockscreen,
+                        to = Scenes.Communal,
+                        progress = flowOf(progress)
+                    ),
+                stateTransition =
+                    TransitionStep(
+                        transitionState = TransitionState.RUNNING,
+                        from = LOCKSCREEN,
+                        to = GLANCEABLE_HUB,
+                        value = progress,
+                    )
             )
+
             runCurrent()
             assertThat(alpha).isIn(Range.closed(0f, 1f))
 
             // Finish transition to glanceable hub
-            keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(
-                    transitionState = TransitionState.FINISHED,
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GLANCEABLE_HUB,
-                    value = 1f,
-                )
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Communal),
+                stateTransition =
+                    TransitionStep(
+                        transitionState = TransitionState.FINISHED,
+                        from = LOCKSCREEN,
+                        to = GLANCEABLE_HUB,
+                        value = 1f,
+                    )
             )
             assertThat(alpha).isEqualTo(0f)
 
@@ -348,35 +370,46 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
 
             // Start transitioning to glanceable hub
             val progress = 0.6f
-            keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(
-                    transitionState = TransitionState.STARTED,
-                    from = KeyguardState.DREAMING,
-                    to = KeyguardState.GLANCEABLE_HUB,
-                    value = 0f,
-                )
+            kosmos.setTransition(
+                sceneTransition = Transition(from = Scenes.Lockscreen, to = Scenes.Communal),
+                stateTransition =
+                    TransitionStep(
+                        transitionState = TransitionState.STARTED,
+                        from = DREAMING,
+                        to = GLANCEABLE_HUB,
+                        value = 0f,
+                    )
             )
             runCurrent()
-            keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(
-                    transitionState = TransitionState.RUNNING,
-                    from = KeyguardState.DREAMING,
-                    to = KeyguardState.GLANCEABLE_HUB,
-                    value = progress,
-                )
+            kosmos.setTransition(
+                sceneTransition =
+                    Transition(
+                        from = Scenes.Lockscreen,
+                        to = Scenes.Communal,
+                        progress = flowOf(progress)
+                    ),
+                stateTransition =
+                    TransitionStep(
+                        transitionState = TransitionState.RUNNING,
+                        from = DREAMING,
+                        to = GLANCEABLE_HUB,
+                        value = progress,
+                    )
             )
             runCurrent()
             // Keep notifications hidden during the transition from dream to hub
             assertThat(alpha).isEqualTo(0)
 
             // Finish transition to glanceable hub
-            keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(
-                    transitionState = TransitionState.FINISHED,
-                    from = KeyguardState.DREAMING,
-                    to = KeyguardState.GLANCEABLE_HUB,
-                    value = 1f,
-                )
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Communal),
+                stateTransition =
+                    TransitionStep(
+                        transitionState = TransitionState.FINISHED,
+                        from = DREAMING,
+                        to = GLANCEABLE_HUB,
+                        value = 1f,
+                    )
             )
             assertThat(alpha).isEqualTo(0f)
         }
@@ -400,35 +433,47 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
         testScope.runTest {
             val isOnLockscreen by collectLastValue(underTest.isOnLockscreen)
 
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.LOCKSCREEN,
-                to = KeyguardState.GONE,
-                testScope,
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Gone),
+                stateTransition = TransitionStep(from = LOCKSCREEN, to = GONE)
             )
             assertThat(isOnLockscreen).isFalse()
 
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Lockscreen),
+                stateTransition = TransitionStep(from = GONE, to = LOCKSCREEN)
+            )
+            assertThat(isOnLockscreen).isTrue()
             // While progressing from lockscreen, should still be true
-            keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GONE,
-                    value = 0.8f,
-                    transitionState = TransitionState.RUNNING
-                )
+            kosmos.setTransition(
+                sceneTransition = Transition(from = Scenes.Lockscreen, to = Scenes.Gone),
+                stateTransition =
+                    TransitionStep(
+                        from = LOCKSCREEN,
+                        to = GONE,
+                        value = 0.8f,
+                        transitionState = TransitionState.RUNNING
+                    )
             )
             assertThat(isOnLockscreen).isTrue()
 
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.GONE,
-                to = KeyguardState.LOCKSCREEN,
-                testScope,
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Lockscreen),
+                stateTransition =
+                    TransitionStep(
+                        from = GONE,
+                        to = LOCKSCREEN,
+                    )
             )
             assertThat(isOnLockscreen).isTrue()
 
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.LOCKSCREEN,
-                to = KeyguardState.PRIMARY_BOUNCER,
-                testScope,
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Bouncer),
+                stateTransition =
+                    TransitionStep(
+                        from = LOCKSCREEN,
+                        to = PRIMARY_BOUNCER,
+                    )
             )
             assertThat(isOnLockscreen).isTrue()
         }
@@ -442,8 +487,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             shadeTestUtil.setLockscreenShadeExpansion(0f)
             shadeTestUtil.setQsExpansion(0f)
             keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.LOCKSCREEN,
-                to = KeyguardState.OCCLUDED,
+                from = LOCKSCREEN,
+                to = OCCLUDED,
                 testScope,
             )
             assertThat(isOnLockscreenWithoutShade).isFalse()
@@ -480,11 +525,15 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             assertThat(isOnGlanceableHubWithoutShade).isFalse()
 
             // Move to glanceable hub
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.LOCKSCREEN,
-                to = KeyguardState.GLANCEABLE_HUB,
-                testScope = this
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Communal),
+                stateTransition =
+                    TransitionStep(
+                        from = LOCKSCREEN,
+                        to = GLANCEABLE_HUB,
+                    )
             )
+
             assertThat(isOnGlanceableHubWithoutShade).isTrue()
 
             // While state is GLANCEABLE_HUB, validate variations of both shade and qs expansion
@@ -502,6 +551,14 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
 
             shadeTestUtil.setQsExpansion(0f)
             shadeTestUtil.setLockscreenShadeExpansion(0f)
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Communal),
+                stateTransition =
+                    TransitionStep(
+                        from = LOCKSCREEN,
+                        to = GLANCEABLE_HUB,
+                    )
+            )
             assertThat(isOnGlanceableHubWithoutShade).isTrue()
         }
 
@@ -808,8 +865,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             // GONE transition gets to 90% complete
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GONE,
+                    from = LOCKSCREEN,
+                    to = GONE,
                     transitionState = TransitionState.STARTED,
                     value = 0f,
                 )
@@ -817,8 +874,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             runCurrent()
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GONE,
+                    from = LOCKSCREEN,
+                    to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.9f,
                 )
@@ -843,8 +900,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             // OCCLUDED transition gets to 90% complete
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.OCCLUDED,
+                    from = LOCKSCREEN,
+                    to = OCCLUDED,
                     transitionState = TransitionState.STARTED,
                     value = 0f,
                 )
@@ -852,8 +909,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             runCurrent()
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.OCCLUDED,
+                    from = LOCKSCREEN,
+                    to = OCCLUDED,
                     transitionState = TransitionState.RUNNING,
                     value = 0.9f,
                 )
@@ -877,8 +934,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             showLockscreen()
 
             keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.LOCKSCREEN,
-                to = KeyguardState.GONE,
+                from = LOCKSCREEN,
+                to = GONE,
                 testScope
             )
             keyguardRepository.setStatusBarState(StatusBarState.SHADE)
@@ -922,8 +979,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
 
             // ... then user hits power to go to AOD
             keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.LOCKSCREEN,
-                to = KeyguardState.AOD,
+                from = LOCKSCREEN,
+                to = AOD,
                 testScope,
             )
             // ... followed by a shade collapse
@@ -945,7 +1002,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             // PRIMARY_BOUNCER->GONE transition is started
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.STARTED,
                     value = 0f,
@@ -956,7 +1013,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             // PRIMARY_BOUNCER->GONE transition running
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.1f,
@@ -967,7 +1024,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
 
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.9f,
@@ -979,7 +1036,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             hideCommunalScene()
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.FINISHED,
                     value = 1f
@@ -1003,7 +1060,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             // PRIMARY_BOUNCER->GONE transition is started
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.STARTED,
                 )
@@ -1013,7 +1070,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             // PRIMARY_BOUNCER->GONE transition running
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.1f,
@@ -1024,7 +1081,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
 
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.9f,
@@ -1035,7 +1092,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
 
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.FINISHED,
                     value = 1f
@@ -1058,7 +1115,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             // ALTERNATE_BOUNCER->GONE transition is started
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.STARTED,
                     value = 0f,
@@ -1069,7 +1126,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             // ALTERNATE_BOUNCER->GONE transition running
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.1f,
@@ -1080,7 +1137,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
 
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.9f,
@@ -1092,7 +1149,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             hideCommunalScene()
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.FINISHED,
                     value = 1f
@@ -1116,7 +1173,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             // ALTERNATE_BOUNCER->GONE transition is started
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.STARTED,
                 )
@@ -1126,7 +1183,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
             // ALTERNATE_BOUNCER->GONE transition running
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.1f,
@@ -1137,7 +1194,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
 
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.9f,
@@ -1148,7 +1205,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
 
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.FINISHED,
                     value = 1f
@@ -1165,8 +1222,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
         keyguardRepository.setStatusBarState(StatusBarState.KEYGUARD)
         runCurrent()
         keyguardTransitionRepository.sendTransitionSteps(
-            from = KeyguardState.AOD,
-            to = KeyguardState.LOCKSCREEN,
+            from = AOD,
+            to = LOCKSCREEN,
             testScope,
         )
     }
@@ -1178,8 +1235,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
         keyguardRepository.setDreaming(true)
         runCurrent()
         keyguardTransitionRepository.sendTransitionSteps(
-            from = KeyguardState.LOCKSCREEN,
-            to = KeyguardState.DREAMING,
+            from = LOCKSCREEN,
+            to = DREAMING,
             testScope,
         )
     }
@@ -1191,8 +1248,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
         keyguardRepository.setStatusBarState(StatusBarState.SHADE_LOCKED)
         runCurrent()
         keyguardTransitionRepository.sendTransitionSteps(
-            from = KeyguardState.AOD,
-            to = KeyguardState.LOCKSCREEN,
+            from = AOD,
+            to = LOCKSCREEN,
             testScope,
         )
     }
@@ -1204,8 +1261,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
         keyguardRepository.setStatusBarState(StatusBarState.SHADE_LOCKED)
         runCurrent()
         keyguardTransitionRepository.sendTransitionSteps(
-            from = KeyguardState.AOD,
-            to = KeyguardState.LOCKSCREEN,
+            from = AOD,
+            to = LOCKSCREEN,
             testScope,
         )
     }
@@ -1219,8 +1276,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
         kosmos.keyguardBouncerRepository.setPrimaryShow(true)
         runCurrent()
         keyguardTransitionRepository.sendTransitionSteps(
-            from = KeyguardState.GLANCEABLE_HUB,
-            to = KeyguardState.PRIMARY_BOUNCER,
+            from = GLANCEABLE_HUB,
+            to = PRIMARY_BOUNCER,
             testScope,
         )
     }
@@ -1234,8 +1291,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization) : S
         kosmos.keyguardBouncerRepository.setPrimaryShow(false)
         runCurrent()
         keyguardTransitionRepository.sendTransitionSteps(
-            from = KeyguardState.GLANCEABLE_HUB,
-            to = KeyguardState.ALTERNATE_BOUNCER,
+            from = GLANCEABLE_HUB,
+            to = ALTERNATE_BOUNCER,
             testScope,
         )
     }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt
index cd3df07eea55..fc70ea5cd526 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt
@@ -334,7 +334,7 @@ constructor(
             listenForSleepTransition(
                 modeOnCanceledFromStartedStep = { startedStep ->
                     if (
-                        transitionInteractor.asleepKeyguardState.value == KeyguardState.AOD &&
+                        keyguardInteractor.asleepKeyguardState.value == KeyguardState.AOD &&
                             startedStep.from == KeyguardState.AOD
                     ) {
                         TransitionModeOnCanceled.REVERSE
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
index a96d7a8f0997..f6f0cc58be71 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
@@ -36,7 +36,9 @@ import com.android.systemui.keyguard.shared.model.DozeStateModel
 import com.android.systemui.keyguard.shared.model.DozeStateModel.Companion.isDozeOff
 import com.android.systemui.keyguard.shared.model.DozeTransitionModel
 import com.android.systemui.keyguard.shared.model.Edge
+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.DOZING
 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.KeyguardState.OCCLUDED
@@ -409,6 +411,12 @@ constructor(
         }
     }
 
+    /** Which keyguard state to use when the device goes to sleep. */
+    val asleepKeyguardState: StateFlow<KeyguardState> =
+        repository.isAodAvailable
+            .map { aodAvailable -> if (aodAvailable) AOD else DOZING }
+            .stateIn(applicationScope, SharingStarted.Eagerly, DOZING)
+
     /**
      * Whether the primary authentication is required for the given user due to lockdown or
      * encryption after reboot.
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
index 6ff369ec8711..c10bacf32a5a 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
@@ -22,14 +22,16 @@ import android.util.Log
 import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.keyguard.data.repository.KeyguardRepository
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.Edge
 import com.android.systemui.keyguard.shared.model.KeyguardState
 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.DOZING
+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.KeyguardState.OCCLUDED
+import com.android.systemui.keyguard.shared.model.KeyguardState.OFF
 import com.android.systemui.keyguard.shared.model.KeyguardState.PRIMARY_BOUNCER
 import com.android.systemui.keyguard.shared.model.KeyguardState.UNDEFINED
 import com.android.systemui.keyguard.shared.model.TransitionState
@@ -66,7 +68,6 @@ class KeyguardTransitionInteractor
 @Inject
 constructor(
     @Application val scope: CoroutineScope,
-    private val keyguardRepository: KeyguardRepository,
     private val repository: KeyguardTransitionRepository,
     private val fromLockscreenTransitionInteractor: dagger.Lazy<FromLockscreenTransitionInteractor>,
     private val fromPrimaryBouncerTransitionInteractor:
@@ -126,8 +127,10 @@ constructor(
             repository.transitions
                 .filter { it.transitionState != TransitionState.CANCELED }
                 .collect { step ->
-                    getTransitionValueFlow(step.from).emit(1f - step.value)
-                    getTransitionValueFlow(step.to).emit(step.value)
+                    val value =
+                        if (step.transitionState == TransitionState.FINISHED) 1f else step.value
+                    getTransitionValueFlow(step.from).emit(1f - value)
+                    getTransitionValueFlow(step.to).emit(value)
                 }
         }
 
@@ -183,8 +186,14 @@ constructor(
         }
     }
 
-    fun transition(edge: Edge, edgeWithoutSceneContainer: Edge): Flow<TransitionStep> {
-        return transition(if (SceneContainerFlag.isEnabled) edge else edgeWithoutSceneContainer)
+    fun transition(edge: Edge, edgeWithoutSceneContainer: Edge? = null): Flow<TransitionStep> {
+        return transition(
+            if (SceneContainerFlag.isEnabled || edgeWithoutSceneContainer == null) {
+                edge
+            } else {
+                edgeWithoutSceneContainer
+            }
+        )
     }
 
     /** Given an [edge], return a Flow to collect only relevant [TransitionStep]s. */
@@ -250,10 +259,10 @@ constructor(
     }
 
     fun transitionValue(
-        scene: SceneKey,
+        scene: SceneKey? = null,
         stateWithoutSceneContainer: KeyguardState,
     ): Flow<Float> {
-        return if (SceneContainerFlag.isEnabled) {
+        return if (SceneContainerFlag.isEnabled && scene != null) {
             sceneInteractor.transitionProgress(scene)
         } else {
             transitionValue(stateWithoutSceneContainer)
@@ -297,12 +306,6 @@ constructor(
             .buffer(2, BufferOverflow.DROP_OLDEST)
             .shareIn(scope, SharingStarted.Eagerly, replay = 1)
 
-    /** Which keyguard state to use when the device goes to sleep. */
-    val asleepKeyguardState: StateFlow<KeyguardState> =
-        keyguardRepository.isAodAvailable
-            .map { aodAvailable -> if (aodAvailable) AOD else DOZING }
-            .stateIn(scope, SharingStarted.Eagerly, DOZING)
-
     /**
      * The last [KeyguardState] to which we [TransitionState.FINISHED] a transition.
      *
@@ -410,7 +413,7 @@ constructor(
                 }
             }
             .distinctUntilChanged()
-            .stateIn(scope, SharingStarted.Eagerly, KeyguardState.OFF)
+            .stateIn(scope, SharingStarted.Eagerly, OFF)
 
     val isInTransition =
         combine(
@@ -438,8 +441,8 @@ constructor(
                 fromAlternateBouncerTransitionInteractor.get().dismissAlternateBouncer()
             AOD -> fromAodTransitionInteractor.get().dismissAod()
             DOZING -> fromDozingTransitionInteractor.get().dismissFromDozing()
-            KeyguardState.OCCLUDED -> fromOccludedTransitionInteractor.get().dismissFromOccluded()
-            KeyguardState.GONE ->
+            OCCLUDED -> fromOccludedTransitionInteractor.get().dismissFromOccluded()
+            GONE ->
                 Log.i(
                     TAG,
                     "Already transitioning to GONE; ignoring startDismissKeyguardTransition."
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/TransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/TransitionInteractor.kt
index d06ee645652c..950eafa23043 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/TransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/TransitionInteractor.kt
@@ -211,7 +211,7 @@ sealed class TransitionInteractor(
             .map(modeOnCanceledFromStartedStep)
             .collect { modeOnCanceled ->
                 startTransitionTo(
-                    toState = transitionInteractor.asleepKeyguardState.value,
+                    toState = keyguardInteractor.asleepKeyguardState.value,
                     modeOnCanceled = modeOnCanceled,
                     ownerReason = "Sleep transition triggered"
                 )
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
index cc46216b6e43..7eb48d6a06ff 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
@@ -479,7 +479,7 @@ constructor(
                     switchToScene(
                         targetSceneKey = Scenes.Lockscreen,
                         loggingReason = "device is starting to sleep",
-                        sceneState = keyguardTransitionInteractor.asleepKeyguardState.value,
+                        sceneState = keyguardInteractor.asleepKeyguardState.value,
                     )
                 } else {
                     val canSwipeToEnter = deviceEntryInteractor.canSwipeToEnter.value
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
index f63ee7b9520d..aed00d8cd5be 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
@@ -20,6 +20,7 @@
 package com.android.systemui.statusbar.notification.stack.ui.viewmodel
 
 import androidx.annotation.VisibleForTesting
+import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.common.shared.model.NotificationContainerBounds
 import com.android.systemui.communal.domain.interactor.CommunalSceneInteractor
 import com.android.systemui.dagger.SysUISingleton
@@ -141,10 +142,6 @@ constructor(
     private val communalSceneInteractor: CommunalSceneInteractor,
     unfoldTransitionInteractor: UnfoldTransitionInteractor,
 ) : FlowDumperImpl(dumpManager) {
-    // TODO(b/349784682): Transform deprecated states for Flexiglass
-    private val statesForConstrainedNotifications: Set<KeyguardState> =
-        setOf(AOD, LOCKSCREEN, DOZING, ALTERNATE_BOUNCER, PRIMARY_BOUNCER)
-    private val statesForHiddenKeyguard: Set<KeyguardState> = setOf(GONE, OCCLUDED)
 
     /**
      * Is either shade/qs expanded? This intentionally does not use the [ShadeInteractor] version,
@@ -217,14 +214,16 @@ constructor(
 
     /** If the user is visually on one of the unoccluded lockscreen states. */
     val isOnLockscreen: Flow<Boolean> =
-        combine(
-                keyguardTransitionInteractor.finishedKeyguardState.map {
-                    statesForConstrainedNotifications.contains(it)
-                },
+        anyOf(
+                keyguardTransitionInteractor.isFinishedIn(AOD),
+                keyguardTransitionInteractor.isFinishedIn(DOZING),
+                keyguardTransitionInteractor.isFinishedIn(ALTERNATE_BOUNCER),
+                keyguardTransitionInteractor.isFinishedIn(
+                    scene = Scenes.Bouncer,
+                    stateWithoutSceneContainer = PRIMARY_BOUNCER
+                ),
                 keyguardTransitionInteractor.transitionValue(LOCKSCREEN).map { it > 0f },
-            ) { constrainedNotificationState, transitioningToOrFromLockscreen ->
-                constrainedNotificationState || transitioningToOrFromLockscreen
-            }
+            )
             .stateIn(
                 scope = applicationScope,
                 started = SharingStarted.Eagerly,
@@ -250,9 +249,10 @@ constructor(
     /** If the user is visually on the glanceable hub or transitioning to/from it */
     private val isOnGlanceableHub: Flow<Boolean> =
         combine(
-                keyguardTransitionInteractor.finishedKeyguardState.map { state ->
-                    state == GLANCEABLE_HUB
-                },
+                keyguardTransitionInteractor.isFinishedIn(
+                    scene = Scenes.Communal,
+                    stateWithoutSceneContainer = GLANCEABLE_HUB
+                ),
                 anyOf(
                     keyguardTransitionInteractor.isInTransition(
                         edge = Edge.create(to = Scenes.Communal),
@@ -424,32 +424,19 @@ constructor(
             .onStart { emit(1f) }
             .dumpWhileCollecting("alphaForShadeAndQsExpansion")
 
-    private fun toFlowArray(
-        states: Set<KeyguardState>,
-        flow: (KeyguardState) -> Flow<Boolean>
-    ): Array<Flow<Boolean>> {
-        return states.map { flow(it) }.toTypedArray()
-    }
-
     private val isTransitioningToHiddenKeyguard: Flow<Boolean> =
         flow {
                 while (currentCoroutineContext().isActive) {
                     emit(false)
                     // Ensure states are inactive to start
-                    allOf(
-                            *toFlowArray(statesForHiddenKeyguard) { state ->
-                                keyguardTransitionInteractor.transitionValue(state).map { it == 0f }
-                            }
-                        )
-                        .first { it }
+                    allOf(isNotOnState(OCCLUDED), isNotOnState(GONE, Scenes.Gone)).first { it }
                     // Wait for a qualifying transition to begin
                     anyOf(
-                            *toFlowArray(statesForHiddenKeyguard) { state ->
-                                keyguardTransitionInteractor
-                                    .transition(Edge.create(to = state))
-                                    .map { it.value > 0f && it.transitionState == RUNNING }
-                                    .onStart { emit(false) }
-                            }
+                            transitionToIsRunning(Edge.create(to = OCCLUDED)),
+                            transitionToIsRunning(
+                                edge = Edge.create(to = Scenes.Gone),
+                                edgeWithoutSceneContainer = Edge.create(to = GONE)
+                            )
                         )
                         .first { it }
                     emit(true)
@@ -458,13 +445,7 @@ constructor(
                     // it is considered safe to reset alpha to 1f for HUNs.
                     combine(
                             keyguardInteractor.statusBarState,
-                            allOf(
-                                *toFlowArray(statesForHiddenKeyguard) { state ->
-                                    keyguardTransitionInteractor.transitionValue(state).map {
-                                        it == 0f
-                                    }
-                                }
-                            )
+                            allOf(isNotOnState(OCCLUDED), isNotOnState(GONE, Scenes.Gone))
                         ) { statusBarState, stateIsReversed ->
                             statusBarState == SHADE || stateIsReversed
                         }
@@ -473,6 +454,17 @@ constructor(
             }
             .dumpWhileCollecting("isTransitioningToHiddenKeyguard")
 
+    private fun isNotOnState(stateWithoutSceneContainer: KeyguardState, scene: SceneKey? = null) =
+        keyguardTransitionInteractor
+            .transitionValue(scene = scene, stateWithoutSceneContainer = stateWithoutSceneContainer)
+            .map { it == 0f }
+
+    private fun transitionToIsRunning(edge: Edge, edgeWithoutSceneContainer: Edge? = null) =
+        keyguardTransitionInteractor
+            .transition(edge = edge, edgeWithoutSceneContainer = edgeWithoutSceneContainer)
+            .map { it.value > 0f && it.transitionState == RUNNING }
+            .onStart { emit(false) }
+
     val panelAlpha = keyguardInteractor.panelAlpha
 
     private fun bouncerToGoneNotificationAlpha(viewState: ViewStateAccessor): Flow<Float> =
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt
index 616f2b688746..a73c184a1ba8 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt
@@ -102,6 +102,45 @@ class FakeKeyguardTransitionRepository(
     }
 
     /**
+     * Sends the provided [step] and makes sure that all previous [TransitionState]'s are sent when
+     * [fillInSteps] is true. e.g. when a step FINISHED is provided, a step with STARTED and RUNNING
+     * is also sent.
+     */
+    suspend fun sendTransitionSteps(
+        step: TransitionStep,
+        testScope: TestScope,
+        fillInSteps: Boolean = true,
+    ) {
+        if (fillInSteps && step.transitionState != TransitionState.STARTED) {
+            sendTransitionStep(
+                step =
+                    TransitionStep(
+                        transitionState = TransitionState.STARTED,
+                        from = step.from,
+                        to = step.to,
+                        value = 0f,
+                    )
+            )
+            testScope.testScheduler.runCurrent()
+
+            if (step.transitionState != TransitionState.RUNNING) {
+                sendTransitionStep(
+                    step =
+                        TransitionStep(
+                            transitionState = TransitionState.RUNNING,
+                            from = step.from,
+                            to = step.to,
+                            value = 0.6f,
+                        )
+                )
+                testScope.testScheduler.runCurrent()
+            }
+        }
+        sendTransitionStep(step = step)
+        testScope.testScheduler.runCurrent()
+    }
+
+    /**
      * Sends TransitionSteps between [from] and [to], calling [runCurrent] after each step.
      *
      * By default, sends steps through FINISHED (STARTED, RUNNING, FINISHED) but can be halted part
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorKosmos.kt
index b68d6a0510d5..8e8f4b69e401 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorKosmos.kt
@@ -27,7 +27,6 @@ val Kosmos.keyguardTransitionInteractor: KeyguardTransitionInteractor by
         KeyguardTransitionInteractor(
             scope = applicationCoroutineScope,
             repository = keyguardTransitionRepository,
-            keyguardRepository = keyguardRepository,
             fromLockscreenTransitionInteractor = { fromLockscreenTransitionInteractor },
             fromPrimaryBouncerTransitionInteractor = { fromPrimaryBouncerTransitionInteractor },
             fromAodTransitionInteractor = { fromAodTransitionInteractor },
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryUtil.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryUtil.kt
index 53d3c0121b1f..59f2b9412413 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryUtil.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryUtil.kt
@@ -19,6 +19,7 @@ package com.android.systemui.scene.data.repository
 import com.android.compose.animation.scene.ObservableTransitionState
 import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
+import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testScope
@@ -36,20 +37,26 @@ private val mutableTransitionState =
 suspend fun Kosmos.setTransition(
     sceneTransition: ObservableTransitionState,
     stateTransition: TransitionStep? = null,
+    fillInStateSteps: Boolean = true,
     scope: TestScope = testScope,
     repository: SceneContainerRepository = sceneContainerRepository
 ) {
+    var state: TransitionStep? = stateTransition
     if (SceneContainerFlag.isEnabled) {
         setSceneTransition(sceneTransition, scope, repository)
-    } else {
-        if (stateTransition == null) throw IllegalArgumentException("No transitionStep provided")
-        fakeKeyguardTransitionRepository.sendTransitionSteps(
-            from = stateTransition.from,
-            to = stateTransition.to,
-            testScope = scope,
-            throughTransitionState = stateTransition.transitionState
-        )
+
+        if (state != null) {
+            state = getStateWithUndefined(sceneTransition, state)
+        }
     }
+
+    if (state == null) return
+    fakeKeyguardTransitionRepository.sendTransitionSteps(
+        step = state,
+        testScope = scope,
+        fillInSteps = fillInStateSteps,
+    )
+    scope.testScheduler.runCurrent()
 }
 
 fun Kosmos.setSceneTransition(
@@ -59,7 +66,7 @@ fun Kosmos.setSceneTransition(
 ) {
     repository.setTransitionState(mutableTransitionState)
     mutableTransitionState.value = transition
-    scope.runCurrent()
+    scope.testScheduler.runCurrent()
 }
 
 fun Transition(
@@ -87,3 +94,43 @@ fun Transition(
 fun Idle(currentScene: SceneKey): ObservableTransitionState.Idle {
     return ObservableTransitionState.Idle(currentScene)
 }
+
+private fun getStateWithUndefined(
+    sceneTransition: ObservableTransitionState,
+    state: TransitionStep
+): TransitionStep {
+    return when (sceneTransition) {
+        is ObservableTransitionState.Idle -> {
+            TransitionStep(
+                from = state.from,
+                to =
+                    if (sceneTransition.currentScene != Scenes.Lockscreen) {
+                        KeyguardState.UNDEFINED
+                    } else {
+                        state.to
+                    },
+                value = state.value,
+                transitionState = state.transitionState
+            )
+        }
+        is ObservableTransitionState.Transition -> {
+            TransitionStep(
+                from =
+                    if (sceneTransition.fromScene != Scenes.Lockscreen) {
+                        KeyguardState.UNDEFINED
+                    } else {
+                        state.from
+                    },
+                to =
+                    if (sceneTransition.toScene != Scenes.Lockscreen) {
+                        KeyguardState.UNDEFINED
+                    } else {
+                        state.from
+                    },
+                value = state.value,
+                transitionState = state.transitionState
+            )
+        }
+        else -> state
+    }
+}