diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerScene.kt
index 6d9497d..b9baa793 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerScene.kt
@@ -46,6 +46,7 @@
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.unit.dp
+import com.android.compose.animation.scene.SceneScope
 import com.android.systemui.R
 import com.android.systemui.bouncer.ui.viewmodel.AuthMethodBouncerViewModel
 import com.android.systemui.bouncer.ui.viewmodel.BouncerViewModel
@@ -81,7 +82,7 @@
             .asStateFlow()
 
     @Composable
-    override fun Content(
+    override fun SceneScope.Content(
         modifier: Modifier,
     ) = BouncerScene(viewModel, dialogFactory, modifier)
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenScene.kt
index ab7bc26..ca7352e 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenScene.kt
@@ -29,6 +29,7 @@
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.unit.dp
+import com.android.compose.animation.scene.SceneScope
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.common.ui.compose.Icon
 import com.android.systemui.dagger.SysUISingleton
@@ -66,7 +67,7 @@
             )
 
     @Composable
-    override fun Content(
+    override fun SceneScope.Content(
         modifier: Modifier,
     ) {
         LockscreenScene(
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt
index 130395a..29763c2 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt
@@ -27,6 +27,7 @@
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.unit.dp
+import com.android.compose.animation.scene.SceneScope
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.qs.ui.viewmodel.QuickSettingsSceneViewModel
 import com.android.systemui.scene.shared.model.Direction
@@ -57,7 +58,7 @@
             .asStateFlow()
 
     @Composable
-    override fun Content(
+    override fun SceneScope.Content(
         modifier: Modifier,
     ) {
         QuickSettingsScene(
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposableScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposableScene.kt
index a213666..3da6a02 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposableScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposableScene.kt
@@ -18,9 +18,10 @@
 
 import androidx.compose.runtime.Composable
 import androidx.compose.ui.Modifier
+import com.android.compose.animation.scene.SceneScope
 import com.android.systemui.scene.shared.model.Scene
 
 /** Compose-capable extension of [Scene]. */
 interface ComposableScene : Scene {
-    @Composable fun Content(modifier: Modifier)
+    @Composable fun SceneScope.Content(modifier: Modifier)
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/GoneScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/GoneScene.kt
index 0070552..774c409 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/GoneScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/GoneScene.kt
@@ -23,6 +23,7 @@
 import androidx.compose.runtime.Composable
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
+import com.android.compose.animation.scene.SceneScope
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.scene.shared.model.Direction
 import com.android.systemui.scene.shared.model.SceneKey
@@ -50,7 +51,7 @@
             .asStateFlow()
 
     @Composable
-    override fun Content(
+    override fun SceneScope.Content(
         modifier: Modifier,
     ) {
         /*
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
index 49e2bf9..3dfdbba 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
@@ -14,32 +14,31 @@
  * limitations under the License.
  */
 
-@file:OptIn(ExperimentalAnimationApi::class)
-
 package com.android.systemui.scene.ui.composable
 
-import androidx.activity.compose.BackHandler
-import androidx.compose.animation.AnimatedContent
-import androidx.compose.animation.ExperimentalAnimationApi
-import androidx.compose.foundation.layout.Arrangement
-import androidx.compose.foundation.layout.Box
-import androidx.compose.foundation.layout.Column
-import androidx.compose.foundation.layout.Row
 import androidx.compose.foundation.layout.fillMaxSize
-import androidx.compose.material3.Button
-import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
-import androidx.compose.ui.Alignment
+import androidx.compose.runtime.remember
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.unit.dp
+import com.android.compose.animation.scene.Back
+import com.android.compose.animation.scene.ObservableTransitionState as SceneTransitionObservableTransitionState
+import com.android.compose.animation.scene.SceneKey as SceneTransitionSceneKey
+import com.android.compose.animation.scene.SceneTransitionLayout
+import com.android.compose.animation.scene.SceneTransitionLayoutState
+import com.android.compose.animation.scene.Swipe
+import com.android.compose.animation.scene.UserAction as SceneTransitionUserAction
+import com.android.compose.animation.scene.observableTransitionState
+import com.android.compose.animation.scene.transitions
 import com.android.systemui.scene.shared.model.Direction
+import com.android.systemui.scene.shared.model.ObservableTransitionState
 import com.android.systemui.scene.shared.model.SceneKey
 import com.android.systemui.scene.shared.model.SceneModel
 import com.android.systemui.scene.shared.model.UserAction
 import com.android.systemui.scene.ui.viewmodel.SceneContainerViewModel
-import java.util.Locale
+import kotlinx.coroutines.flow.map
 
 /**
  * Renders a container of a collection of "scenes" that the user can switch between using certain
@@ -64,77 +63,94 @@
     sceneByKey: Map<SceneKey, ComposableScene>,
     modifier: Modifier = Modifier,
 ) {
-    val currentScene: SceneModel by viewModel.currentScene.collectAsState()
+    val currentSceneModel: SceneModel by viewModel.currentScene.collectAsState()
+    val currentSceneKey = currentSceneModel.key
+    val currentScene = checkNotNull(sceneByKey[currentSceneKey])
+    val currentDestinations: Map<UserAction, SceneModel> by
+        currentScene.destinationScenes().collectAsState()
+    val state = remember { SceneTransitionLayoutState(currentSceneKey.toTransitionSceneKey()) }
 
-    AnimatedContent(
-        targetState = currentScene.key,
-        label = "scene container",
-        modifier = modifier,
-    ) { currentSceneKey ->
-        sceneByKey.forEach { (key, composableScene) ->
-            if (key == currentSceneKey) {
-                Scene(
-                    scene = composableScene,
-                    onSceneChanged = viewModel::setCurrentScene,
-                    modifier = Modifier.fillMaxSize(),
-                )
-            }
-        }
+    DisposableEffect(viewModel, state) {
+        viewModel.setTransitionState(state.observableTransitionState().map { it.toModel() })
+        onDispose { viewModel.setTransitionState(null) }
     }
-}
 
-/** Renders the given [ComposableScene]. */
-@Composable
-private fun Scene(
-    scene: ComposableScene,
-    onSceneChanged: (SceneModel) -> Unit,
-    modifier: Modifier = Modifier,
-) {
-    val destinationScenes: Map<UserAction, SceneModel> by scene.destinationScenes().collectAsState()
-    val swipeLeftDestinationScene = destinationScenes[UserAction.Swipe(Direction.LEFT)]
-    val swipeUpDestinationScene = destinationScenes[UserAction.Swipe(Direction.UP)]
-    val swipeRightDestinationScene = destinationScenes[UserAction.Swipe(Direction.RIGHT)]
-    val swipeDownDestinationScene = destinationScenes[UserAction.Swipe(Direction.DOWN)]
-    val backDestinationScene = destinationScenes[UserAction.Back]
-
-    // TODO(b/280880714): replace with the real UI and make sure to call onTransitionProgress.
-    Box(modifier) {
-        Column(
-            horizontalAlignment = Alignment.CenterHorizontally,
-            modifier = Modifier.align(Alignment.Center),
-        ) {
-            scene.Content(
-                modifier = Modifier,
-            )
-
-            Row(
-                horizontalArrangement = Arrangement.spacedBy(8.dp),
-            ) {
-                DirectionalButton(Direction.LEFT, swipeLeftDestinationScene, onSceneChanged)
-                DirectionalButton(Direction.UP, swipeUpDestinationScene, onSceneChanged)
-                DirectionalButton(Direction.RIGHT, swipeRightDestinationScene, onSceneChanged)
-                DirectionalButton(Direction.DOWN, swipeDownDestinationScene, onSceneChanged)
-            }
-
-            if (backDestinationScene != null) {
-                BackHandler { onSceneChanged.invoke(backDestinationScene) }
-            }
-        }
-    }
-}
-
-@Composable
-private fun DirectionalButton(
-    direction: Direction,
-    destinationScene: SceneModel?,
-    onSceneChanged: (SceneModel) -> Unit,
-    modifier: Modifier = Modifier,
-) {
-    Button(
-        onClick = { destinationScene?.let { onSceneChanged.invoke(it) } },
-        enabled = destinationScene != null,
-        modifier = modifier,
+    SceneTransitionLayout(
+        currentScene = currentSceneKey.toTransitionSceneKey(),
+        onChangeScene = { sceneKey -> viewModel.setCurrentScene(sceneKey.toModel()) },
+        transitions = transitions {},
+        state = state,
+        modifier = modifier.fillMaxSize(),
     ) {
-        Text(direction.name.lowercase(Locale.getDefault()))
+        sceneByKey.forEach { (sceneKey, composableScene) ->
+            scene(
+                key = sceneKey.toTransitionSceneKey(),
+                userActions =
+                    if (sceneKey == currentSceneKey) {
+                            currentDestinations
+                        } else {
+                            composableScene.destinationScenes().value
+                        }
+                        .map { (userAction, destinationSceneModel) ->
+                            toTransitionModels(userAction, destinationSceneModel)
+                        }
+                        .toMap(),
+            ) {
+                with(composableScene) {
+                    this@scene.Content(
+                        modifier = Modifier.fillMaxSize(),
+                    )
+                }
+            }
+        }
+    }
+}
+
+// TODO(b/293899074): remove this once we can use the one from SceneTransitionLayout.
+private fun SceneTransitionObservableTransitionState.toModel(): ObservableTransitionState {
+    return when (this) {
+        is SceneTransitionObservableTransitionState.Idle ->
+            ObservableTransitionState.Idle(scene.toModel().key)
+        is SceneTransitionObservableTransitionState.Transition ->
+            ObservableTransitionState.Transition(
+                fromScene = fromScene.toModel().key,
+                toScene = toScene.toModel().key,
+                progress = progress,
+            )
+    }
+}
+
+// TODO(b/293899074): remove this once we can use the one from SceneTransitionLayout.
+private fun toTransitionModels(
+    userAction: UserAction,
+    sceneModel: SceneModel,
+): Pair<SceneTransitionUserAction, SceneTransitionSceneKey> {
+    return userAction.toTransitionUserAction() to sceneModel.key.toTransitionSceneKey()
+}
+
+// TODO(b/293899074): remove this once we can use the one from SceneTransitionLayout.
+private fun SceneKey.toTransitionSceneKey(): SceneTransitionSceneKey {
+    return SceneTransitionSceneKey(
+        name = toString(),
+        identity = this,
+    )
+}
+
+// TODO(b/293899074): remove this once we can use the one from SceneTransitionLayout.
+private fun SceneTransitionSceneKey.toModel(): SceneModel {
+    return SceneModel(key = identity as SceneKey)
+}
+
+// TODO(b/293899074): remove this once we can use the one from SceneTransitionLayout.
+private fun UserAction.toTransitionUserAction(): SceneTransitionUserAction {
+    return when (this) {
+        is UserAction.Swipe ->
+            when (this.direction) {
+                Direction.LEFT -> Swipe.Left
+                Direction.UP -> Swipe.Up
+                Direction.RIGHT -> Swipe.Right
+                Direction.DOWN -> Swipe.Down
+            }
+        is UserAction.Back -> Back
     }
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
index b73e0b2..ff1cb5f 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
@@ -26,6 +26,7 @@
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.unit.dp
+import com.android.compose.animation.scene.SceneScope
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.notifications.ui.composable.Notifications
@@ -63,7 +64,7 @@
             )
 
     @Composable
-    override fun Content(
+    override fun SceneScope.Content(
         modifier: Modifier,
     ) = ShadeScene(viewModel, modifier)
 
@@ -86,11 +87,12 @@
         horizontalAlignment = Alignment.CenterHorizontally,
         verticalArrangement = Arrangement.spacedBy(16.dp),
         modifier =
-            Modifier.fillMaxSize()
+            modifier
+                .fillMaxSize()
                 .clickable(onClick = { viewModel.onContentClicked() })
                 .padding(horizontal = 16.dp, vertical = 48.dp)
     ) {
-        QuickSettings(modifier = modifier.height(160.dp))
-        Notifications(modifier = modifier.weight(1f))
+        QuickSettings(modifier = Modifier.height(160.dp))
+        Notifications(modifier = Modifier.weight(1f))
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/data/repository/SceneContainerRepository.kt b/packages/SystemUI/src/com/android/systemui/scene/data/repository/SceneContainerRepository.kt
index 1fca488..fee3960 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/data/repository/SceneContainerRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/data/repository/SceneContainerRepository.kt
@@ -14,16 +14,23 @@
  * limitations under the License.
  */
 
+@file:OptIn(ExperimentalCoroutinesApi::class)
+
 package com.android.systemui.scene.data.repository
 
+import com.android.systemui.scene.shared.model.ObservableTransitionState
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.SceneKey
 import com.android.systemui.scene.shared.model.SceneModel
 import com.android.systemui.scene.shared.model.SceneTransitionModel
 import javax.inject.Inject
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.flowOf
 
 /** Source of truth for scene framework application state. */
 class SceneContainerRepository
@@ -38,8 +45,17 @@
     private val _currentScene = MutableStateFlow(SceneModel(config.initialSceneKey))
     val currentScene: StateFlow<SceneModel> = _currentScene.asStateFlow()
 
-    private val _transitionProgress = MutableStateFlow(1f)
-    val transitionProgress: StateFlow<Float> = _transitionProgress.asStateFlow()
+    private val transitionState = MutableStateFlow<Flow<ObservableTransitionState>?>(null)
+    val transitionProgress: Flow<Float> =
+        transitionState.flatMapLatest { observableTransitionStateFlow ->
+            observableTransitionStateFlow?.flatMapLatest { observableTransitionState ->
+                when (observableTransitionState) {
+                    is ObservableTransitionState.Idle -> flowOf(1f)
+                    is ObservableTransitionState.Transition -> observableTransitionState.progress
+                }
+            }
+                ?: flowOf(1f)
+        }
 
     private val _transitions = MutableStateFlow<SceneTransitionModel?>(null)
     val transitions: StateFlow<SceneTransitionModel?> = _transitions.asStateFlow()
@@ -92,8 +108,12 @@
         _isVisible.value = isVisible
     }
 
-    /** Sets scene transition progress to the current scene in the container with the given name. */
-    fun setSceneTransitionProgress(progress: Float) {
-        _transitionProgress.value = progress
+    /**
+     * Binds the given flow so the system remembers it.
+     *
+     * Note that you must call is with `null` when the UI is done or risk a memory leak.
+     */
+    fun setTransitionState(transitionState: Flow<ObservableTransitionState>?) {
+        this.transitionState.value = transitionState
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
index 39daad3..b09a5cf 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
@@ -18,11 +18,13 @@
 
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.scene.data.repository.SceneContainerRepository
+import com.android.systemui.scene.shared.model.ObservableTransitionState
 import com.android.systemui.scene.shared.model.RemoteUserInput
 import com.android.systemui.scene.shared.model.SceneKey
 import com.android.systemui.scene.shared.model.SceneModel
 import com.android.systemui.scene.shared.model.SceneTransitionModel
 import javax.inject.Inject
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
@@ -69,13 +71,17 @@
     /** Whether the container with the given name is visible. */
     val isVisible: StateFlow<Boolean> = repository.isVisible
 
-    /** Sets scene transition progress to the current scene in the container with the given name. */
-    fun setSceneTransitionProgress(progress: Float) {
-        repository.setSceneTransitionProgress(progress)
+    /**
+     * Binds the given flow so the system remembers it.
+     *
+     * Note that you must call is with `null` when the UI is done or risk a memory leak.
+     */
+    fun setTransitionState(transitionState: Flow<ObservableTransitionState>?) {
+        repository.setTransitionState(transitionState)
     }
 
     /** Progress of the transition into the current scene in the container with the given name. */
-    val transitionProgress: StateFlow<Float> = repository.transitionProgress
+    val transitionProgress: Flow<Float> = repository.transitionProgress
 
     /**
      * Scene transitions as pairs of keys. A new value is emitted exactly once, each time a scene
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/model/ObservableTransitionState.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/model/ObservableTransitionState.kt
new file mode 100644
index 0000000..9a30aa6
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/model/ObservableTransitionState.kt
@@ -0,0 +1,36 @@
+/*
+ * Copyright 2023 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.scene.shared.model
+
+import kotlinx.coroutines.flow.Flow
+
+/**
+ * This is a fork of a class by the same name in the `com.android.compose.animation.scene` package.
+ *
+ * TODO(b/293899074): remove this fork, once we can compile Compose into System UI.
+ */
+sealed class ObservableTransitionState {
+    /** No transition/animation is currently running. */
+    data class Idle(val scene: SceneKey) : ObservableTransitionState()
+
+    /** There is a transition animating between two scenes. */
+    data class Transition(
+        val fromScene: SceneKey,
+        val toScene: SceneKey,
+        val progress: Flow<Float>,
+    ) : ObservableTransitionState()
+}
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
index f44748a..bd73e36 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
@@ -19,10 +19,12 @@
 import android.view.MotionEvent
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.scene.domain.interactor.SceneInteractor
+import com.android.systemui.scene.shared.model.ObservableTransitionState
 import com.android.systemui.scene.shared.model.RemoteUserInput
 import com.android.systemui.scene.shared.model.SceneKey
 import com.android.systemui.scene.shared.model.SceneModel
 import javax.inject.Inject
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 
 /** Models UI state for the scene container. */
@@ -54,9 +56,13 @@
         interactor.setCurrentScene(scene)
     }
 
-    /** Notifies of the progress of a scene transition. */
-    fun setSceneTransitionProgress(progress: Float) {
-        interactor.setSceneTransitionProgress(progress)
+    /**
+     * Binds the given flow so the system remembers it.
+     *
+     * Note that you must call is with `null` when the UI is done or risk a memory leak.
+     */
+    fun setTransitionState(transitionState: Flow<ObservableTransitionState>?) {
+        interactor.setTransitionState(transitionState)
     }
 
     /** Handles a [MotionEvent] representing remote user input. */
diff --git a/packages/SystemUI/tests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt
index 826a6cc..56e3e96 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt
@@ -22,11 +22,13 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.scene.SceneTestUtils
+import com.android.systemui.scene.shared.model.ObservableTransitionState
 import com.android.systemui.scene.shared.model.SceneKey
 import com.android.systemui.scene.shared.model.SceneModel
 import com.android.systemui.scene.shared.model.SceneTransitionModel
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -91,11 +93,30 @@
         val sceneTransitionProgress by collectLastValue(underTest.transitionProgress)
         assertThat(sceneTransitionProgress).isEqualTo(1f)
 
-        underTest.setSceneTransitionProgress(0.1f)
+        val transitionState =
+            MutableStateFlow<ObservableTransitionState>(
+                ObservableTransitionState.Idle(SceneKey.Lockscreen)
+            )
+        underTest.setTransitionState(transitionState)
+        assertThat(sceneTransitionProgress).isEqualTo(1f)
+
+        val progress = MutableStateFlow(1f)
+        transitionState.value =
+            ObservableTransitionState.Transition(
+                fromScene = SceneKey.Lockscreen,
+                toScene = SceneKey.Shade,
+                progress = progress,
+            )
+        assertThat(sceneTransitionProgress).isEqualTo(1f)
+
+        progress.value = 0.1f
         assertThat(sceneTransitionProgress).isEqualTo(0.1f)
 
-        underTest.setSceneTransitionProgress(0.9f)
+        progress.value = 0.9f
         assertThat(sceneTransitionProgress).isEqualTo(0.9f)
+
+        underTest.setTransitionState(null)
+        assertThat(sceneTransitionProgress).isEqualTo(1f)
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
index 13a602d..c193d83 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
@@ -22,11 +22,13 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.scene.SceneTestUtils
+import com.android.systemui.scene.shared.model.ObservableTransitionState
 import com.android.systemui.scene.shared.model.SceneKey
 import com.android.systemui.scene.shared.model.SceneModel
 import com.android.systemui.scene.shared.model.SceneTransitionModel
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -37,7 +39,8 @@
 class SceneInteractorTest : SysuiTestCase() {
 
     private val utils = SceneTestUtils(this)
-    private val underTest = utils.sceneInteractor()
+    private val repository = utils.fakeSceneContainerRepository()
+    private val underTest = utils.sceneInteractor(repository = repository)
 
     @Test
     fun allSceneKeys() {
@@ -55,11 +58,20 @@
 
     @Test
     fun sceneTransitionProgress() = runTest {
-        val progress by collectLastValue(underTest.transitionProgress)
-        assertThat(progress).isEqualTo(1f)
+        val transitionProgress by collectLastValue(underTest.transitionProgress)
+        assertThat(transitionProgress).isEqualTo(1f)
 
-        underTest.setSceneTransitionProgress(0.55f)
-        assertThat(progress).isEqualTo(0.55f)
+        val progress = MutableStateFlow(0.55f)
+        repository.setTransitionState(
+            MutableStateFlow(
+                ObservableTransitionState.Transition(
+                    fromScene = SceneKey.Lockscreen,
+                    toScene = SceneKey.Shade,
+                    progress = progress,
+                ),
+            )
+        )
+        assertThat(transitionProgress).isEqualTo(0.55f)
     }
 
     @Test
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt
index 26a75d0..70d15a0 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt
@@ -119,9 +119,11 @@
         )
     }
 
-    fun sceneInteractor(): SceneInteractor {
+    fun sceneInteractor(
+        repository: SceneContainerRepository = fakeSceneContainerRepository()
+    ): SceneInteractor {
         return SceneInteractor(
-            repository = fakeSceneContainerRepository(),
+            repository = repository,
         )
     }