Expose previewProgress and isInPreviewStage in ObservableTransitionState

Bug: 350705972
Test: ObservableTransitionStateTest
Flag: com.android.systemui.scene_container
Change-Id: I708eb96d81ea75e2bb264b054d3a215a5f30c7e9

3 files changed, 105 insertions, 3 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 ae5344fd7922..8f1a4141176a 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
@@ -74,6 +74,12 @@ sealed interface ObservableTransitionState {
          * the transition completes/settles.
          */
         val isUserInputOngoing: Flow<Boolean>,
+
+        /** Current progress of the preview part of the transition */
+        val previewProgress: Flow<Float> = flowOf(0f),
+
+        /** Whether the transition is currently in the preview stage or not */
+        val isInPreviewStage: Flow<Boolean> = flowOf(false),
     ) : ObservableTransitionState {
         override fun toString(): String =
             """Transition
@@ -113,6 +119,8 @@ fun SceneTransitionLayoutState.observableTransitionState(): Flow<ObservableTrans
                         progress = snapshotFlow { state.progress },
                         isInitiatedByUserInput = state.isInitiatedByUserInput,
                         isUserInputOngoing = snapshotFlow { state.isUserInputOngoing },
+                        previewProgress = snapshotFlow { state.previewProgress },
+                        isInPreviewStage = snapshotFlow { state.isInPreviewStage }
                     )
                 }
             }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt
index f717301dba38..0543e7f09e5d 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt
@@ -16,11 +16,16 @@
 
 package com.android.compose.animation.scene
 
+import androidx.activity.BackEventCompat
+import androidx.activity.ComponentActivity
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 import androidx.compose.runtime.snapshots.Snapshot
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.test.junit4.createAndroidComposeRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.compose.animation.scene.TestScenes.SceneA
 import com.android.compose.animation.scene.TestScenes.SceneB
@@ -36,7 +41,7 @@ import org.junit.runner.RunWith
 
 @RunWith(AndroidJUnit4::class)
 class ObservableTransitionStateTest {
-    @get:Rule val rule = createComposeRule()
+    @get:Rule val rule = createAndroidComposeRule<ComponentActivity>()
 
     @Test
     fun testObservableTransitionState() = runTest {
@@ -145,6 +150,82 @@ class ObservableTransitionStateTest {
         assertThat(currentScene.value).isEqualTo(SceneA)
     }
 
+    @Test
+    fun testObservablePreviewTransitionState() = runTest {
+        val layoutState =
+            rule.runOnUiThread {
+                MutableSceneTransitionLayoutState(
+                    SceneA,
+                    transitions = transitions { from(SceneA, to = SceneB, preview = {}) }
+                )
+            }
+        rule.setContent {
+            SceneTransitionLayout(layoutState) {
+                scene(SceneA, mapOf(Back to SceneB)) { Box(Modifier.fillMaxSize()) }
+                scene(SceneB) { Box(Modifier.fillMaxSize()) }
+            }
+        }
+
+        var observableState: ObservableTransitionState? = null
+        backgroundScope.launch {
+            layoutState.observableTransitionState().collect { observableState = it }
+        }
+
+        fun observableState(): ObservableTransitionState {
+            runCurrent()
+            return observableState!!
+        }
+
+        fun ObservableTransitionState.Transition.previewProgress(): Float {
+            var lastProgress = -1f
+            backgroundScope.launch { previewProgress.collect { lastProgress = it } }
+            runCurrent()
+            return lastProgress
+        }
+
+        fun ObservableTransitionState.Transition.isInPreviewStage(): Boolean {
+            var lastIsInPreviewStage = false
+            backgroundScope.launch { isInPreviewStage.collect { lastIsInPreviewStage = it } }
+            runCurrent()
+            return lastIsInPreviewStage
+        }
+
+        // Start back.
+        val dispatcher = rule.activity.onBackPressedDispatcher
+        rule.runOnUiThread {
+            dispatcher.dispatchOnBackStarted(backEvent())
+            dispatcher.dispatchOnBackProgressed(backEvent(progress = 0.4f))
+        }
+
+        var state = observableState()
+        assertThat(state).isInstanceOf(ObservableTransitionState.Transition::class.java)
+        assertThat((state as ObservableTransitionState.Transition).fromScene).isEqualTo(SceneA)
+        assertThat(state.previewProgress()).isEqualTo(0.4f)
+        assertThat(state.isInPreviewStage()).isEqualTo(true)
+
+        // Cancel it.
+        rule.runOnUiThread { dispatcher.dispatchOnBackCancelled() }
+        rule.waitForIdle()
+        state = observableState()
+        assertThat(state).isInstanceOf(ObservableTransitionState.Idle::class.java)
+
+        // Start again and commit it.
+        rule.runOnUiThread {
+            dispatcher.dispatchOnBackStarted(backEvent())
+            dispatcher.dispatchOnBackProgressed(backEvent(progress = 0.4f))
+            dispatcher.onBackPressed()
+        }
+        state = observableState()
+        assertThat(state).isInstanceOf(ObservableTransitionState.Transition::class.java)
+        assertThat((state as ObservableTransitionState.Transition).fromScene).isEqualTo(SceneA)
+        assertThat(state.previewProgress()).isEqualTo(0.4f)
+        assertThat(state.isInPreviewStage()).isEqualTo(false)
+
+        rule.waitForIdle()
+        state = observableState()
+        assertThat(state).isInstanceOf(ObservableTransitionState.Idle::class.java)
+    }
+
     // See http://shortn/_hj4Mhikmos for inspiration.
     private fun runTestWithSnapshots(testBody: suspend TestScope.() -> Unit) {
         val globalWriteObserverHandle =
@@ -159,4 +240,13 @@ class ObservableTransitionStateTest {
             globalWriteObserverHandle.dispose()
         }
     }
+
+    private fun backEvent(progress: Float = 0f): BackEventCompat {
+        return BackEventCompat(
+            touchX = 0f,
+            touchY = 0f,
+            progress = progress,
+            swipeEdge = BackEventCompat.EDGE_LEFT,
+        )
+    }
 }
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 2f17ca853d16..53d3c0121b1f 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
@@ -69,6 +69,8 @@ fun Transition(
     progress: Flow<Float> = flowOf(0f),
     isInitiatedByUserInput: Boolean = false,
     isUserInputOngoing: Flow<Boolean> = flowOf(false),
+    previewProgress: Flow<Float> = flowOf(0f),
+    isInPreviewStage: Flow<Boolean> = flowOf(false)
 ): ObservableTransitionState.Transition {
     return ObservableTransitionState.Transition(
         fromScene = from,
@@ -76,7 +78,9 @@ fun Transition(
         currentScene = currentScene,
         progress = progress,
         isInitiatedByUserInput = isInitiatedByUserInput,
-        isUserInputOngoing = isUserInputOngoing
+        isUserInputOngoing = isUserInputOngoing,
+        previewProgress = previewProgress,
+        isInPreviewStage = isInPreviewStage
     )
 }