Merge "Re-enable SharedNotificationContainerViewModelTest tests" into main

2 files changed, 116 insertions, 72 deletions

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 a023033e3aa3..3408e06fe713 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
@@ -25,13 +25,12 @@ import android.platform.test.flag.junit.FlagsParameterization
 import androidx.test.filters.SmallTest
 import com.android.systemui.Flags.FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX
 import com.android.systemui.Flags.FLAG_MIGRATE_CLOCKS_TO_BLUEPRINT
-import com.android.systemui.Flags.FLAG_SCENE_CONTAINER
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.common.shared.model.NotificationContainerBounds
 import com.android.systemui.common.ui.data.repository.fakeConfigurationRepository
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
-import com.android.systemui.flags.BrokenWithSceneContainer
+import com.android.systemui.flags.DisableSceneContainer
 import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.flags.Flags
 import com.android.systemui.flags.andSceneContainer
@@ -51,9 +50,8 @@ 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.shared.flag.SceneContainerFlag
-import com.android.systemui.shade.data.repository.shadeRepository
 import com.android.systemui.shade.mockLargeScreenHeaderHelper
+import com.android.systemui.shade.shadeTestUtil
 import com.android.systemui.statusbar.notification.stack.domain.interactor.sharedNotificationContainerInteractor
 import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.any
@@ -116,8 +114,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         get() = kosmos.keyguardRootViewModel
     val keyguardTransitionRepository
         get() = kosmos.fakeKeyguardTransitionRepository
-    val shadeRepository
-        get() = kosmos.shadeRepository
+    val shadeTestUtil
+        get() = kosmos.shadeTestUtil
     val sharedNotificationContainerInteractor
         get() = kosmos.sharedNotificationContainerInteractor
     val largeScreenHeaderHelper
@@ -127,7 +125,6 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
 
     @Before
     fun setUp() {
-        assertThat(SceneContainerFlag.isEnabled).isEqualTo(SceneContainerFlag.isEnabled)
         overrideResource(R.bool.config_use_split_notification_shade, false)
         movementFlow = MutableStateFlow(BurnInModel())
         whenever(aodBurnInViewModel.movement(any())).thenReturn(movementFlow)
@@ -234,7 +231,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     @Test
-    @DisableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX, FLAG_SCENE_CONTAINER)
+    @DisableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
+    @DisableSceneContainer
     fun validateMarginTopWithLargeScreenHeader_refactorFlagOff_usesResource() =
         testScope.runTest {
             val headerResourceHeight = 50
@@ -273,7 +271,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     @EnableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
     fun validateMarginTopWithLargeScreenHeader_refactorFlagOn_usesHelper() =
         testScope.runTest {
@@ -313,7 +311,6 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     @Test
-    @BrokenWithSceneContainer(bugId = 333132830)
     fun glanceableHubAlpha_lockscreenToHub() =
         testScope.runTest {
             val alpha by collectLastValue(underTest.glanceableHubAlpha)
@@ -357,7 +354,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
 
             // While state is GLANCEABLE_HUB, verify alpha is restored to full if glanceable hub is
             // not fully visible.
-            shadeRepository.setLockscreenShadeExpansion(0.1f)
+            shadeTestUtil.setLockscreenShadeExpansion(0.1f)
             assertThat(alpha).isEqualTo(1f)
         }
 
@@ -463,14 +460,13 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     @Test
-    @BrokenWithSceneContainer(bugId = 333132830)
     fun isOnLockscreenWithoutShade() =
         testScope.runTest {
             val isOnLockscreenWithoutShade by collectLastValue(underTest.isOnLockscreenWithoutShade)
 
             // First on AOD
-            shadeRepository.setLockscreenShadeExpansion(0f)
-            shadeRepository.setQsExpansion(0f)
+            shadeTestUtil.setLockscreenShadeExpansion(0f)
+            shadeTestUtil.setQsExpansion(0f)
             keyguardTransitionRepository.sendTransitionSteps(
                 from = KeyguardState.LOCKSCREEN,
                 to = KeyguardState.OCCLUDED,
@@ -482,25 +478,24 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
             showLockscreen()
 
             // While state is LOCKSCREEN, validate variations of both shade and qs expansion
-            shadeRepository.setLockscreenShadeExpansion(0.1f)
-            shadeRepository.setQsExpansion(0f)
+            shadeTestUtil.setQsExpansion(0f)
+            shadeTestUtil.setLockscreenShadeExpansion(0.1f)
             assertThat(isOnLockscreenWithoutShade).isFalse()
 
-            shadeRepository.setLockscreenShadeExpansion(0.1f)
-            shadeRepository.setQsExpansion(0.1f)
+            shadeTestUtil.setLockscreenShadeExpansion(0.1f)
+            shadeTestUtil.setShadeAndQsExpansion(0.1f, .9f)
             assertThat(isOnLockscreenWithoutShade).isFalse()
 
-            shadeRepository.setLockscreenShadeExpansion(0f)
-            shadeRepository.setQsExpansion(0.1f)
+            shadeTestUtil.setLockscreenShadeExpansion(0f)
+            shadeTestUtil.setQsExpansion(0.1f)
             assertThat(isOnLockscreenWithoutShade).isFalse()
 
-            shadeRepository.setLockscreenShadeExpansion(0f)
-            shadeRepository.setQsExpansion(0f)
+            shadeTestUtil.setQsExpansion(0f)
+            shadeTestUtil.setLockscreenShadeExpansion(0f)
             assertThat(isOnLockscreenWithoutShade).isTrue()
         }
 
     @Test
-    @BrokenWithSceneContainer(bugId = 333132830)
     fun isOnGlanceableHubWithoutShade() =
         testScope.runTest {
             val isOnGlanceableHubWithoutShade by
@@ -519,25 +514,25 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
             assertThat(isOnGlanceableHubWithoutShade).isTrue()
 
             // While state is GLANCEABLE_HUB, validate variations of both shade and qs expansion
-            shadeRepository.setLockscreenShadeExpansion(0.1f)
-            shadeRepository.setQsExpansion(0f)
+            shadeTestUtil.setQsExpansion(0f)
+            shadeTestUtil.setLockscreenShadeExpansion(0.1f)
             assertThat(isOnGlanceableHubWithoutShade).isFalse()
 
-            shadeRepository.setLockscreenShadeExpansion(0.1f)
-            shadeRepository.setQsExpansion(0.1f)
+            shadeTestUtil.setLockscreenShadeExpansion(0.1f)
+            shadeTestUtil.setShadeAndQsExpansion(0.1f, .9f)
             assertThat(isOnGlanceableHubWithoutShade).isFalse()
 
-            shadeRepository.setLockscreenShadeExpansion(0f)
-            shadeRepository.setQsExpansion(0.1f)
+            shadeTestUtil.setLockscreenShadeExpansion(0f)
+            shadeTestUtil.setQsExpansion(0.1f)
             assertThat(isOnGlanceableHubWithoutShade).isFalse()
 
-            shadeRepository.setLockscreenShadeExpansion(0f)
-            shadeRepository.setQsExpansion(0f)
+            shadeTestUtil.setQsExpansion(0f)
+            shadeTestUtil.setLockscreenShadeExpansion(0f)
             assertThat(isOnGlanceableHubWithoutShade).isTrue()
         }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun boundsOnLockscreenNotInSplitShade() =
         testScope.runTest {
             val bounds by collectLastValue(underTest.bounds)
@@ -558,7 +553,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     @Test
-    @DisableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX, FLAG_SCENE_CONTAINER)
+    @DisableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
+    @DisableSceneContainer
     fun boundsOnLockscreenInSplitShade_refactorFlagOff_usesLargeHeaderResource() =
         testScope.runTest {
             val bounds by collectLastValue(underTest.bounds)
@@ -593,7 +589,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
 
     @Test
     @EnableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun boundsOnLockscreenInSplitShade_refactorFlagOn_usesLargeHeaderHelper() =
         testScope.runTest {
             val bounds by collectLastValue(underTest.bounds)
@@ -621,7 +617,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun boundsOnShade() =
         testScope.runTest {
             val bounds by collectLastValue(underTest.bounds)
@@ -637,7 +633,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun boundsOnQS() =
         testScope.runTest {
             val bounds by collectLastValue(underTest.bounds)
@@ -682,7 +678,6 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     @Test
-    @BrokenWithSceneContainer(bugId = 333132830)
     fun maxNotificationsOnLockscreen_DoesNotUpdateWhenUserInteracting() =
         testScope.runTest {
             var notificationCount = 10
@@ -700,26 +695,25 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
             assertThat(maxNotifications).isEqualTo(10)
 
             // Shade expanding... still 10
-            shadeRepository.setLockscreenShadeExpansion(0.5f)
+            shadeTestUtil.setLockscreenShadeExpansion(0.5f)
             assertThat(maxNotifications).isEqualTo(10)
 
             notificationCount = 25
 
             // When shade is expanding by user interaction
-            shadeRepository.setLegacyLockscreenShadeTracking(true)
+            shadeTestUtil.setLockscreenShadeTracking(true)
 
             // Should still be 10, since the user is interacting
             assertThat(maxNotifications).isEqualTo(10)
 
-            shadeRepository.setLegacyLockscreenShadeTracking(false)
-            shadeRepository.setLockscreenShadeExpansion(0f)
+            shadeTestUtil.setLockscreenShadeTracking(false)
+            shadeTestUtil.setLockscreenShadeExpansion(0f)
 
             // Stopped tracking, show 25
             assertThat(maxNotifications).isEqualTo(25)
         }
 
     @Test
-    @BrokenWithSceneContainer(bugId = 333132830)
     fun maxNotificationsOnShade() =
         testScope.runTest {
             val calculateSpace = { space: Float, useExtraShelfSpace: Boolean -> 10 }
@@ -739,7 +733,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun translationYUpdatesOnKeyguardForBurnIn() =
         testScope.runTest {
             val translationY by collectLastValue(underTest.translationY(BurnInParameters()))
@@ -752,7 +746,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun translationYUpdatesOnKeyguard() =
         testScope.runTest {
             val translationY by collectLastValue(underTest.translationY(BurnInParameters()))
@@ -764,7 +758,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
             configurationRepository.onAnyConfigurationChange()
 
             // legacy expansion means the user is swiping up, usually for the bouncer
-            shadeRepository.setLegacyShadeExpansion(0.5f)
+            shadeTestUtil.setShadeExpansion(0.5f)
 
             showLockscreen()
 
@@ -773,7 +767,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun translationYDoesNotUpdateWhenShadeIsExpanded() =
         testScope.runTest {
             val translationY by collectLastValue(underTest.translationY(BurnInParameters()))
@@ -786,7 +780,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
 
             // legacy expansion means the user is swiping up, usually for the bouncer but also for
             // shade collapsing
-            shadeRepository.setLegacyShadeExpansion(0.5f)
+            shadeTestUtil.setShadeExpansion(0.5f)
 
             showLockscreenWithShadeExpanded()
 
@@ -794,7 +788,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     @Test
-    @DisableFlags(FLAG_SCENE_CONTAINER)
+    @DisableSceneContainer
     fun updateBounds_fromKeyguardRoot() =
         testScope.runTest {
             val bounds by collectLastValue(underTest.bounds)
@@ -806,7 +800,6 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     @Test
-    @BrokenWithSceneContainer(bugId = 333132830)
     fun alphaOnFullQsExpansion() =
         testScope.runTest {
             val viewState = ViewStateAccessor()
@@ -815,13 +808,13 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
             showLockscreenWithQSExpanded()
 
             // Alpha fades out as QS expands
-            shadeRepository.setQsExpansion(0.5f)
+            shadeTestUtil.setQsExpansion(0.5f)
             assertThat(alpha).isWithin(0.01f).of(0.5f)
-            shadeRepository.setQsExpansion(0.9f)
+            shadeTestUtil.setQsExpansion(0.9f)
             assertThat(alpha).isWithin(0.01f).of(0.1f)
 
             // Ensure that alpha is set back to 1f when QS is fully expanded
-            shadeRepository.setQsExpansion(1f)
+            shadeTestUtil.setQsExpansion(1f)
             assertThat(alpha).isEqualTo(1f)
         }
 
@@ -856,7 +849,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
             assertThat(alpha).isEqualTo(0f)
 
             // An attempt to override by the shade should be ignored
-            shadeRepository.setQsExpansion(0.5f)
+            shadeTestUtil.setQsExpansion(0.5f)
             assertThat(alpha).isEqualTo(0f)
         }
 
@@ -891,7 +884,7 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
             assertThat(alpha).isEqualTo(0f)
 
             // An attempt to override by the shade should be ignored
-            shadeRepository.setQsExpansion(0.5f)
+            shadeTestUtil.setQsExpansion(0.5f)
             assertThat(alpha).isEqualTo(0f)
         }
 
@@ -914,7 +907,6 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     @Test
-    @BrokenWithSceneContainer(bugId = 333132830)
     fun shadeCollapseFadeIn() =
         testScope.runTest {
             val fadeIn by collectValues(underTest.shadeCollapseFadeIn)
@@ -961,8 +953,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
         }
 
     private suspend fun TestScope.showLockscreen() {
-        shadeRepository.setLockscreenShadeExpansion(0f)
-        shadeRepository.setQsExpansion(0f)
+        shadeTestUtil.setQsExpansion(0f)
+        shadeTestUtil.setLockscreenShadeExpansion(0f)
         runCurrent()
         keyguardRepository.setStatusBarState(StatusBarState.KEYGUARD)
         runCurrent()
@@ -974,8 +966,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
     }
 
     private suspend fun TestScope.showDream() {
-        shadeRepository.setLockscreenShadeExpansion(0f)
-        shadeRepository.setQsExpansion(0f)
+        shadeTestUtil.setQsExpansion(0f)
+        shadeTestUtil.setLockscreenShadeExpansion(0f)
         runCurrent()
         keyguardRepository.setDreaming(true)
         runCurrent()
@@ -987,8 +979,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
     }
 
     private suspend fun TestScope.showLockscreenWithShadeExpanded() {
-        shadeRepository.setLockscreenShadeExpansion(1f)
-        shadeRepository.setQsExpansion(0f)
+        shadeTestUtil.setQsExpansion(0f)
+        shadeTestUtil.setLockscreenShadeExpansion(1f)
         runCurrent()
         keyguardRepository.setStatusBarState(StatusBarState.SHADE_LOCKED)
         runCurrent()
@@ -1000,8 +992,8 @@ class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) :
     }
 
     private suspend fun TestScope.showLockscreenWithQSExpanded() {
-        shadeRepository.setLockscreenShadeExpansion(0f)
-        shadeRepository.setQsExpansion(1f)
+        shadeTestUtil.setLockscreenShadeExpansion(0f)
+        shadeTestUtil.setQsExpansion(1f)
         runCurrent()
         keyguardRepository.setStatusBarState(StatusBarState.SHADE_LOCKED)
         runCurrent()
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ShadeTestUtil.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ShadeTestUtil.kt
index d17dd6c518a4..0101961c9b72 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ShadeTestUtil.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ShadeTestUtil.kt
@@ -23,7 +23,6 @@ import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shade.data.repository.FakeShadeRepository
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import org.junit.Assert
@@ -33,12 +32,14 @@ class ShadeTestUtil constructor(val delegate: ShadeTestUtilDelegate) {
 
     /** Sets shade expansion to a value between 0-1. */
     fun setShadeExpansion(shadeExpansion: Float) {
-        setShadeAndQsExpansion(shadeExpansion, 0f)
+        delegate.assertFlagValid()
+        delegate.setShadeExpansion(shadeExpansion)
     }
 
     /** Sets QS expansion to a value between 0-1. */
     fun setQsExpansion(qsExpansion: Float) {
-        setShadeAndQsExpansion(0f, qsExpansion)
+        delegate.assertFlagValid()
+        delegate.setQsExpansion(qsExpansion)
     }
 
     /** Sets both shade and QS expansion. One value must be zero or values must add up to 1f. */
@@ -48,7 +49,7 @@ class ShadeTestUtil constructor(val delegate: ShadeTestUtilDelegate) {
             shadeExpansion == 0f || qsExpansion == 0f || shadeExpansion + qsExpansion == 1f,
         )
         delegate.assertFlagValid()
-        delegate.setShadeAndQsExpansionInternal(shadeExpansion, qsExpansion)
+        delegate.setShadeAndQsExpansion(shadeExpansion, qsExpansion)
     }
 
     /** Sets the shade expansion on the lockscreen to the given amount from 0-1. */
@@ -56,6 +57,12 @@ class ShadeTestUtil constructor(val delegate: ShadeTestUtilDelegate) {
         delegate.assertFlagValid()
         delegate.setLockscreenShadeExpansion(lockscreenShadeExpansion)
     }
+
+    /** Sets whether the user is moving the shade with touch input. */
+    fun setLockscreenShadeTracking(lockscreenShadeTracking: Boolean) {
+        delegate.assertFlagValid()
+        delegate.setLockscreenShadeTracking(lockscreenShadeTracking)
+    }
 }
 
 /** Sets up shade state for tests for a specific value of the scene container flag. */
@@ -64,16 +71,25 @@ interface ShadeTestUtilDelegate {
     fun assertFlagValid()
 
     /** Sets both shade and QS expansion. One value must be zero or values must add up to 1f. */
-    fun setShadeAndQsExpansionInternal(shadeExpansion: Float, qsExpansion: Float)
+    fun setShadeAndQsExpansion(shadeExpansion: Float, qsExpansion: Float)
 
     /** Sets the shade expansion on the lockscreen to the given amount from 0-1. */
     fun setLockscreenShadeExpansion(lockscreenShadeExpansion: Float)
+
+    /** Sets whether the user is moving the shade with touch input. */
+    fun setLockscreenShadeTracking(lockscreenShadeTracking: Boolean)
+
+    /** Sets shade expansion to a value between 0-1. */
+    fun setShadeExpansion(shadeExpansion: Float)
+
+    /** Sets QS expansion to a value between 0-1. */
+    fun setQsExpansion(qsExpansion: Float)
 }
 
 /** Sets up shade state for tests when the scene container flag is disabled. */
 class ShadeTestUtilLegacyImpl(val testScope: TestScope, val shadeRepository: FakeShadeRepository) :
     ShadeTestUtilDelegate {
-    override fun setShadeAndQsExpansionInternal(shadeExpansion: Float, qsExpansion: Float) {
+    override fun setShadeAndQsExpansion(shadeExpansion: Float, qsExpansion: Float) {
         shadeRepository.setLegacyShadeExpansion(shadeExpansion)
         shadeRepository.setQsExpansion(qsExpansion)
         testScope.runCurrent()
@@ -83,24 +99,56 @@ class ShadeTestUtilLegacyImpl(val testScope: TestScope, val shadeRepository: Fak
         shadeRepository.setLockscreenShadeExpansion(lockscreenShadeExpansion)
     }
 
+    override fun setLockscreenShadeTracking(lockscreenShadeTracking: Boolean) {
+        shadeRepository.setLegacyLockscreenShadeTracking(lockscreenShadeTracking)
+    }
+
     override fun assertFlagValid() {
         Assert.assertFalse(SceneContainerFlag.isEnabled)
     }
+
+    /** Sets shade expansion to a value between 0-1. */
+    override fun setShadeExpansion(shadeExpansion: Float) {
+        shadeRepository.setLegacyShadeExpansion(shadeExpansion)
+        testScope.runCurrent()
+    }
+
+    /** Sets QS expansion to a value between 0-1. */
+    override fun setQsExpansion(qsExpansion: Float) {
+        shadeRepository.setQsExpansion(qsExpansion)
+        testScope.runCurrent()
+    }
 }
 
 /** Sets up shade state for tests when the scene container flag is enabled. */
 class ShadeTestUtilSceneImpl(val testScope: TestScope, val sceneInteractor: SceneInteractor) :
     ShadeTestUtilDelegate {
-    override fun setShadeAndQsExpansionInternal(shadeExpansion: Float, qsExpansion: Float) {
+    val isUserInputOngoing = MutableStateFlow(true)
+
+    override fun setShadeAndQsExpansion(shadeExpansion: Float, qsExpansion: Float) {
         if (shadeExpansion == 0f) {
             setTransitionProgress(Scenes.Lockscreen, Scenes.QuickSettings, qsExpansion)
         } else if (qsExpansion == 0f) {
             setTransitionProgress(Scenes.Lockscreen, Scenes.Shade, shadeExpansion)
+        } else if (shadeExpansion == 1f) {
+            setIdleScene(Scenes.Shade)
+        } else if (qsExpansion == 1f) {
+            setIdleScene(Scenes.QuickSettings)
         } else {
             setTransitionProgress(Scenes.Shade, Scenes.QuickSettings, qsExpansion)
         }
     }
 
+    /** Sets shade expansion to a value between 0-1. */
+    override fun setShadeExpansion(shadeExpansion: Float) {
+        setShadeAndQsExpansion(shadeExpansion, 0f)
+    }
+
+    /** Sets QS expansion to a value between 0-1. */
+    override fun setQsExpansion(qsExpansion: Float) {
+        setShadeAndQsExpansion(0f, qsExpansion)
+    }
+
     override fun setLockscreenShadeExpansion(lockscreenShadeExpansion: Float) {
         if (lockscreenShadeExpansion == 0f) {
             setIdleScene(Scenes.Lockscreen)
@@ -111,6 +159,10 @@ class ShadeTestUtilSceneImpl(val testScope: TestScope, val sceneInteractor: Scen
         }
     }
 
+    override fun setLockscreenShadeTracking(lockscreenShadeTracking: Boolean) {
+        isUserInputOngoing.value = lockscreenShadeTracking
+    }
+
     private fun setIdleScene(scene: SceneKey) {
         sceneInteractor.changeScene(scene, "test")
         val transitionState =
@@ -127,8 +179,8 @@ class ShadeTestUtilSceneImpl(val testScope: TestScope, val sceneInteractor: Scen
                     fromScene = from,
                     toScene = to,
                     progress = MutableStateFlow(progress),
-                    isInitiatedByUserInput = false,
-                    isUserInputOngoing = flowOf(false),
+                    isInitiatedByUserInput = true,
+                    isUserInputOngoing = isUserInputOngoing,
                 )
             )
         sceneInteractor.setTransitionState(transitionState)