Move most SceneFrameworkIntegrationTest state to Kosmos

This simplifies the logic by removing the need for lazy
in most places, and also should make many of the state methods
more reusable, since they only depend on a single Kosmos for
state

Bug: 353687680
Test: local deviceless
Flag: TEST_ONLY
Change-Id: Ibec38154b0eb593acd0067d9e22cd9a882849043

4 files changed, 225 insertions, 249 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
index 4f7c01358a7f..2d42c4247ab7 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
@@ -18,7 +18,6 @@
 
 package com.android.systemui.scene
 
-import android.telecom.TelecomManager
 import android.telephony.TelephonyManager
 import android.testing.TestableLooper.RunWithLooper
 import androidx.test.ext.junit.runners.AndroidJUnit4
@@ -28,54 +27,43 @@ import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.Swipe
 import com.android.compose.animation.scene.UserActionResult
 import com.android.internal.R
-import com.android.internal.util.EmergencyAffordanceManager
 import com.android.internal.util.emergencyAffordanceManager
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.authentication.data.repository.FakeAuthenticationRepository
 import com.android.systemui.authentication.data.repository.fakeAuthenticationRepository
 import com.android.systemui.authentication.domain.interactor.authenticationInteractor
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
-import com.android.systemui.bouncer.domain.interactor.BouncerActionButtonInteractor
-import com.android.systemui.bouncer.domain.interactor.bouncerActionButtonInteractor
-import com.android.systemui.bouncer.ui.viewmodel.BouncerSceneContentViewModel
 import com.android.systemui.bouncer.ui.viewmodel.PasswordBouncerViewModel
 import com.android.systemui.bouncer.ui.viewmodel.PinBouncerViewModel
 import com.android.systemui.bouncer.ui.viewmodel.bouncerSceneContentViewModel
-import com.android.systemui.classifier.domain.interactor.falsingInteractor
-import com.android.systemui.communal.domain.interactor.communalInteractor
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.deviceentry.data.repository.fakeDeviceEntryRepository
 import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
 import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.flags.Flags
 import com.android.systemui.flags.fakeFeatureFlagsClassic
-import com.android.systemui.keyguard.ui.viewmodel.LockscreenUserActionsViewModel
+import com.android.systemui.keyguard.ui.viewmodel.lockscreenUserActionsViewModel
+import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.lifecycle.activateIn
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAsleepForTest
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAwakeForTest
 import com.android.systemui.power.domain.interactor.powerInteractor
-import com.android.systemui.qs.ui.adapter.fakeQSSceneAdapter
 import com.android.systemui.scene.domain.interactor.sceneInteractor
 import com.android.systemui.scene.domain.resolver.homeSceneFamilyResolver
 import com.android.systemui.scene.domain.startable.sceneContainerStartable
-import com.android.systemui.scene.shared.logger.sceneLogger
 import com.android.systemui.scene.shared.model.SceneFamilies
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.scene.shared.model.fakeSceneDataSource
 import com.android.systemui.scene.ui.viewmodel.SceneContainerViewModel
-import com.android.systemui.shade.domain.interactor.shadeInteractor
-import com.android.systemui.shade.ui.viewmodel.ShadeSceneContentViewModel
-import com.android.systemui.shade.ui.viewmodel.ShadeUserActionsViewModel
 import com.android.systemui.shade.ui.viewmodel.shadeSceneContentViewModel
 import com.android.systemui.shade.ui.viewmodel.shadeUserActionsViewModel
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.telephony.data.repository.fakeTelephonyRepository
 import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.whenever
-import com.android.telecom.telecomManager
+import com.android.telecom.mockTelecomManager
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -83,14 +71,12 @@ import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mockito.verify
-import org.mockito.MockitoAnnotations
 
 /**
  * Integration test cases for the Scene Framework.
@@ -116,195 +102,131 @@ import org.mockito.MockitoAnnotations
 @RunWithLooper
 @EnableSceneContainer
 class SceneFrameworkIntegrationTest : SysuiTestCase() {
-
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
-    private val sceneContainerConfig by lazy { kosmos.sceneContainerConfig }
-    private val sceneInteractor by lazy { kosmos.sceneInteractor }
-    private val authenticationInteractor by lazy { kosmos.authenticationInteractor }
-    private val deviceEntryInteractor by lazy { kosmos.deviceEntryInteractor }
-    private val communalInteractor by lazy { kosmos.communalInteractor }
-
-    private val transitionState by lazy {
-        MutableStateFlow<ObservableTransitionState>(
-            ObservableTransitionState.Idle(sceneContainerConfig.initialSceneKey)
-        )
-    }
-    private val sceneContainerViewModel by lazy {
-        SceneContainerViewModel(
-                sceneInteractor = sceneInteractor,
-                falsingInteractor = kosmos.falsingInteractor,
-                powerInteractor = kosmos.powerInteractor,
-                logger = kosmos.sceneLogger,
-                motionEventHandlerReceiver = {},
-            )
-            .apply { setTransitionState(transitionState) }
-    }
-
-    private lateinit var mobileConnectionsRepository: FakeMobileConnectionsRepository
-    private lateinit var bouncerActionButtonInteractor: BouncerActionButtonInteractor
-    private lateinit var bouncerSceneContentViewModel: BouncerSceneContentViewModel
-
-    private val mLockscreenUserActionsViewModel by lazy {
-        LockscreenUserActionsViewModel(
-            deviceEntryInteractor = deviceEntryInteractor,
-            communalInteractor = communalInteractor,
-            shadeInteractor = kosmos.shadeInteractor,
-        )
-    }
-
-    private lateinit var shadeSceneContentViewModel: ShadeSceneContentViewModel
-    private lateinit var mShadeUserActionsViewModel: ShadeUserActionsViewModel
-
-    private val powerInteractor by lazy { kosmos.powerInteractor }
-
     private var bouncerSceneJob: Job? = null
 
-    private val qsFlexiglassAdapter = kosmos.fakeQSSceneAdapter
-
-    private lateinit var emergencyAffordanceManager: EmergencyAffordanceManager
-    private lateinit var telecomManager: TelecomManager
-    private val fakeSceneDataSource = kosmos.fakeSceneDataSource
-
     @Before
-    fun setUp() {
-        MockitoAnnotations.initMocks(this)
-
-        overrideResource(R.bool.config_enable_emergency_call_while_sim_locked, true)
-        telecomManager = checkNotNull(kosmos.telecomManager)
-        whenever(telecomManager.isInCall).thenReturn(false)
-        emergencyAffordanceManager = kosmos.emergencyAffordanceManager
-        whenever(emergencyAffordanceManager.needsEmergencyAffordance()).thenReturn(true)
+    fun setUp() =
+        kosmos.run {
+            overrideResource(R.bool.config_enable_emergency_call_while_sim_locked, true)
+            whenever(mockTelecomManager.isInCall).thenReturn(false)
+            whenever(emergencyAffordanceManager.needsEmergencyAffordance()).thenReturn(true)
 
-        kosmos.fakeFeatureFlagsClassic.apply { set(Flags.NEW_NETWORK_SLICE_UI, false) }
-
-        mobileConnectionsRepository = kosmos.fakeMobileConnectionsRepository
-        mobileConnectionsRepository.isAnySimSecure.value = false
-
-        kosmos.fakeTelephonyRepository.apply {
-            setHasTelephonyRadio(true)
-            setCallState(TelephonyManager.CALL_STATE_IDLE)
-            setIsInCall(false)
-        }
+            fakeFeatureFlagsClassic.apply { set(Flags.NEW_NETWORK_SLICE_UI, false) }
 
-        bouncerActionButtonInteractor = kosmos.bouncerActionButtonInteractor
-        bouncerSceneContentViewModel = kosmos.bouncerSceneContentViewModel
+            fakeMobileConnectionsRepository.isAnySimSecure.value = false
 
-        shadeSceneContentViewModel = kosmos.shadeSceneContentViewModel
-        mShadeUserActionsViewModel = kosmos.shadeUserActionsViewModel
+            fakeTelephonyRepository.apply {
+                setHasTelephonyRadio(true)
+                setCallState(TelephonyManager.CALL_STATE_IDLE)
+                setIsInCall(false)
+            }
 
-        val startable = kosmos.sceneContainerStartable
-        startable.start()
+            sceneContainerStartable.start()
 
-        mLockscreenUserActionsViewModel.activateIn(testScope)
-        shadeSceneContentViewModel.activateIn(testScope)
-        mShadeUserActionsViewModel.activateIn(testScope)
-        bouncerSceneContentViewModel.activateIn(testScope)
-        sceneContainerViewModel.activateIn(testScope)
+            lockscreenUserActionsViewModel.activateIn(testScope)
+            shadeSceneContentViewModel.activateIn(testScope)
+            shadeUserActionsViewModel.activateIn(testScope)
+            bouncerSceneContentViewModel.activateIn(testScope)
+            sceneContainerViewModel.activateIn(testScope)
 
-        assertWithMessage("Initial scene key mismatch!")
-            .that(sceneContainerViewModel.currentScene.value)
-            .isEqualTo(sceneContainerConfig.initialSceneKey)
-        assertWithMessage("Initial scene container visibility mismatch!")
-            .that(sceneContainerViewModel.isVisible)
-            .isTrue()
-    }
+            assertWithMessage("Initial scene key mismatch!")
+                .that(sceneContainerViewModel.currentScene.value)
+                .isEqualTo(sceneContainerConfig.initialSceneKey)
+            assertWithMessage("Initial scene container visibility mismatch!")
+                .that(sceneContainerViewModel.isVisible)
+                .isTrue()
+        }
 
     @Test
-    fun startsInLockscreenScene() = testScope.runTest { assertCurrentScene(Scenes.Lockscreen) }
+    fun startsInLockscreenScene() =
+        testScope.runTest { kosmos.assertCurrentScene(Scenes.Lockscreen) }
 
     @Test
     fun clickLockButtonAndEnterCorrectPin_unlocksDevice() =
         testScope.runTest {
-            emulateUserDrivenTransition(Scenes.Bouncer)
+            kosmos.emulateUserDrivenTransition(Scenes.Bouncer)
 
-            fakeSceneDataSource.pause()
-            enterPin()
-            emulatePendingTransitionProgress(
-                expectedVisible = false,
-            )
-            assertCurrentScene(Scenes.Gone)
+            kosmos.fakeSceneDataSource.pause()
+            kosmos.enterPin()
+            kosmos.emulatePendingTransitionProgress(expectedVisible = false)
+            kosmos.assertCurrentScene(Scenes.Gone)
         }
 
     @Test
     fun swipeUpOnLockscreen_enterCorrectPin_unlocksDevice() =
         testScope.runTest {
-            val actions by collectLastValue(mLockscreenUserActionsViewModel.actions)
+            val actions by collectLastValue(kosmos.lockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Bouncer)
-            emulateUserDrivenTransition(
+            kosmos.emulateUserDrivenTransition(
                 to = upDestinationSceneKey,
             )
 
-            fakeSceneDataSource.pause()
-            enterPin()
-            emulatePendingTransitionProgress(
-                expectedVisible = false,
-            )
-            assertCurrentScene(Scenes.Gone)
+            kosmos.fakeSceneDataSource.pause()
+            kosmos.enterPin()
+            kosmos.emulatePendingTransitionProgress(expectedVisible = false)
+            kosmos.assertCurrentScene(Scenes.Gone)
         }
 
     @Test
     fun swipeUpOnLockscreen_withAuthMethodSwipe_dismissesLockscreen() =
         testScope.runTest {
-            setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
+            kosmos.setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
 
-            val actions by collectLastValue(mLockscreenUserActionsViewModel.actions)
+            val actions by collectLastValue(kosmos.lockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Gone)
-            emulateUserDrivenTransition(
-                to = upDestinationSceneKey,
-            )
+            kosmos.emulateUserDrivenTransition(to = upDestinationSceneKey)
         }
 
     @Test
     fun swipeUpOnShadeScene_withAuthMethodSwipe_lockscreenNotDismissed_goesToLockscreen() =
         testScope.runTest {
-            val actions by collectLastValue(mShadeUserActionsViewModel.actions)
+            val actions by collectLastValue(kosmos.shadeUserActionsViewModel.actions)
             val homeScene by collectLastValue(kosmos.homeSceneFamilyResolver.resolvedScene)
-            setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
-            assertCurrentScene(Scenes.Lockscreen)
+            kosmos.setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
+            kosmos.assertCurrentScene(Scenes.Lockscreen)
 
             // Emulate a user swipe to the shade scene.
-            emulateUserDrivenTransition(to = Scenes.Shade)
-            assertCurrentScene(Scenes.Shade)
+            kosmos.emulateUserDrivenTransition(to = Scenes.Shade)
+            kosmos.assertCurrentScene(Scenes.Shade)
 
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(SceneFamilies.Home)
             assertThat(homeScene).isEqualTo(Scenes.Lockscreen)
-            emulateUserDrivenTransition(
-                to = homeScene,
-            )
+            kosmos.emulateUserDrivenTransition(to = homeScene)
         }
 
     @Test
     fun swipeUpOnShadeScene_withAuthMethodSwipe_lockscreenDismissed_goesToGone() =
         testScope.runTest {
-            val actions by collectLastValue(mShadeUserActionsViewModel.actions)
-            val canSwipeToEnter by collectLastValue(deviceEntryInteractor.canSwipeToEnter)
+            val actions by collectLastValue(kosmos.shadeUserActionsViewModel.actions)
+            val canSwipeToEnter by collectLastValue(kosmos.deviceEntryInteractor.canSwipeToEnter)
             val homeScene by collectLastValue(kosmos.homeSceneFamilyResolver.resolvedScene)
 
-            setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
+            kosmos.setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
 
             assertThat(canSwipeToEnter).isTrue()
-            assertCurrentScene(Scenes.Lockscreen)
+            kosmos.assertCurrentScene(Scenes.Lockscreen)
 
             // Emulate a user swipe to dismiss the lockscreen.
-            emulateUserDrivenTransition(to = Scenes.Gone)
-            assertCurrentScene(Scenes.Gone)
+            kosmos.emulateUserDrivenTransition(to = Scenes.Gone)
+            kosmos.assertCurrentScene(Scenes.Gone)
 
             // Emulate a user swipe to the shade scene.
-            emulateUserDrivenTransition(to = Scenes.Shade)
-            assertCurrentScene(Scenes.Shade)
+            kosmos.emulateUserDrivenTransition(to = Scenes.Shade)
+            kosmos.assertCurrentScene(Scenes.Shade)
 
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(SceneFamilies.Home)
             assertThat(homeScene).isEqualTo(Scenes.Gone)
-            emulateUserDrivenTransition(
+            kosmos.emulateUserDrivenTransition(
                 to = homeScene,
             )
         }
@@ -312,64 +234,64 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
     @Test
     fun withAuthMethodNone_deviceWakeUp_skipsLockscreen() =
         testScope.runTest {
-            setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = false)
-            putDeviceToSleep(instantlyLockDevice = false)
-            assertCurrentScene(Scenes.Lockscreen)
+            kosmos.setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = false)
+            kosmos.putDeviceToSleep(instantlyLockDevice = false)
+            kosmos.assertCurrentScene(Scenes.Lockscreen)
 
-            wakeUpDevice()
-            assertCurrentScene(Scenes.Gone)
+            kosmos.wakeUpDevice()
+            kosmos.assertCurrentScene(Scenes.Gone)
         }
 
     @Test
     fun withAuthMethodSwipe_deviceWakeUp_doesNotSkipLockscreen() =
         testScope.runTest {
-            setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
-            putDeviceToSleep(instantlyLockDevice = false)
-            assertCurrentScene(Scenes.Lockscreen)
+            kosmos.setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
+            kosmos.putDeviceToSleep(instantlyLockDevice = false)
+            kosmos.assertCurrentScene(Scenes.Lockscreen)
 
-            wakeUpDevice()
-            assertCurrentScene(Scenes.Lockscreen)
+            kosmos.wakeUpDevice()
+            kosmos.assertCurrentScene(Scenes.Lockscreen)
         }
 
     @Test
     fun lockDeviceLocksDevice() =
         testScope.runTest {
-            unlockDevice()
-            assertCurrentScene(Scenes.Gone)
+            kosmos.unlockDevice()
+            kosmos.assertCurrentScene(Scenes.Gone)
 
-            lockDevice()
-            assertCurrentScene(Scenes.Lockscreen)
+            kosmos.lockDevice()
+            kosmos.assertCurrentScene(Scenes.Lockscreen)
         }
 
     @Test
     fun deviceGoesToSleep_switchesToLockscreen() =
         testScope.runTest {
-            unlockDevice()
-            assertCurrentScene(Scenes.Gone)
+            kosmos.unlockDevice()
+            kosmos.assertCurrentScene(Scenes.Gone)
 
-            putDeviceToSleep()
-            assertCurrentScene(Scenes.Lockscreen)
+            kosmos.putDeviceToSleep()
+            kosmos.assertCurrentScene(Scenes.Lockscreen)
         }
 
     @Test
     fun deviceGoesToSleep_wakeUp_unlock() =
         testScope.runTest {
-            unlockDevice()
-            assertCurrentScene(Scenes.Gone)
-            putDeviceToSleep()
-            assertCurrentScene(Scenes.Lockscreen)
-            wakeUpDevice()
-            assertCurrentScene(Scenes.Lockscreen)
-
-            unlockDevice()
-            assertCurrentScene(Scenes.Gone)
+            kosmos.unlockDevice()
+            kosmos.assertCurrentScene(Scenes.Gone)
+            kosmos.putDeviceToSleep()
+            kosmos.assertCurrentScene(Scenes.Lockscreen)
+            kosmos.wakeUpDevice()
+            kosmos.assertCurrentScene(Scenes.Lockscreen)
+
+            kosmos.unlockDevice()
+            kosmos.assertCurrentScene(Scenes.Gone)
         }
 
     @Test
     fun swipeUpOnLockscreenWhileUnlocked_dismissesLockscreen() =
         testScope.runTest {
-            unlockDevice()
-            val actions by collectLastValue(mLockscreenUserActionsViewModel.actions)
+            kosmos.unlockDevice()
+            val actions by collectLastValue(kosmos.lockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Gone)
@@ -378,46 +300,47 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
     @Test
     fun deviceGoesToSleep_withLockTimeout_staysOnLockscreen() =
         testScope.runTest {
-            unlockDevice()
-            assertCurrentScene(Scenes.Gone)
-            putDeviceToSleep(instantlyLockDevice = false)
-            assertCurrentScene(Scenes.Lockscreen)
+            kosmos.unlockDevice()
+            kosmos.assertCurrentScene(Scenes.Gone)
+            kosmos.putDeviceToSleep(instantlyLockDevice = false)
+            kosmos.assertCurrentScene(Scenes.Lockscreen)
 
             // Pretend like the timeout elapsed and now lock the device.
-            lockDevice()
-            assertCurrentScene(Scenes.Lockscreen)
+            kosmos.lockDevice()
+            kosmos.assertCurrentScene(Scenes.Lockscreen)
         }
 
     @Test
     fun dismissingIme_whileOnPasswordBouncer_navigatesToLockscreen() =
         testScope.runTest {
-            setAuthMethod(AuthenticationMethodModel.Password)
-            val actions by collectLastValue(mLockscreenUserActionsViewModel.actions)
+            kosmos.setAuthMethod(AuthenticationMethodModel.Password)
+            val actions by collectLastValue(kosmos.lockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Bouncer)
-            emulateUserDrivenTransition(
+            kosmos.emulateUserDrivenTransition(
                 to = upDestinationSceneKey,
             )
 
-            fakeSceneDataSource.pause()
-            dismissIme()
+            kosmos.fakeSceneDataSource.pause()
+            kosmos.dismissIme()
 
-            emulatePendingTransitionProgress()
-            assertCurrentScene(Scenes.Lockscreen)
+            kosmos.emulatePendingTransitionProgress()
+            kosmos.assertCurrentScene(Scenes.Lockscreen)
         }
 
     @Test
     fun bouncerActionButtonClick_opensEmergencyServicesDialer() =
         testScope.runTest {
-            setAuthMethod(AuthenticationMethodModel.Password)
-            val actions by collectLastValue(mLockscreenUserActionsViewModel.actions)
+            kosmos.setAuthMethod(AuthenticationMethodModel.Password)
+            val actions by collectLastValue(kosmos.lockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Bouncer)
-            emulateUserDrivenTransition(to = upDestinationSceneKey)
+            kosmos.emulateUserDrivenTransition(to = upDestinationSceneKey)
 
-            val bouncerActionButton by collectLastValue(bouncerSceneContentViewModel.actionButton)
+            val bouncerActionButton by
+                collectLastValue(kosmos.bouncerSceneContentViewModel.actionButton)
             assertWithMessage("Bouncer action button not visible")
                 .that(bouncerActionButton)
                 .isNotNull()
@@ -430,54 +353,55 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
     @Test
     fun bouncerActionButtonClick_duringCall_returnsToCall() =
         testScope.runTest {
-            setAuthMethod(AuthenticationMethodModel.Password)
-            startPhoneCall()
-            val actions by collectLastValue(mLockscreenUserActionsViewModel.actions)
+            kosmos.setAuthMethod(AuthenticationMethodModel.Password)
+            kosmos.startPhoneCall()
+            val actions by collectLastValue(kosmos.lockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Bouncer)
-            emulateUserDrivenTransition(to = upDestinationSceneKey)
+            kosmos.emulateUserDrivenTransition(to = upDestinationSceneKey)
 
-            val bouncerActionButton by collectLastValue(bouncerSceneContentViewModel.actionButton)
+            val bouncerActionButton by
+                collectLastValue(kosmos.bouncerSceneContentViewModel.actionButton)
             assertWithMessage("Bouncer action button not visible during call")
                 .that(bouncerActionButton)
                 .isNotNull()
             bouncerActionButton?.onClick?.invoke()
             runCurrent()
 
-            verify(telecomManager).showInCallScreen(any())
+            verify(kosmos.mockTelecomManager).showInCallScreen(any())
         }
 
     @Test
     fun showBouncer_whenLockedSimIntroduced() =
         testScope.runTest {
-            setAuthMethod(AuthenticationMethodModel.None)
-            introduceLockedSim()
-            assertCurrentScene(Scenes.Bouncer)
+            kosmos.setAuthMethod(AuthenticationMethodModel.None)
+            kosmos.introduceLockedSim()
+            kosmos.assertCurrentScene(Scenes.Bouncer)
         }
 
     @Test
     fun goesToGone_whenSimUnlocked_whileDeviceUnlocked() =
         testScope.runTest {
-            fakeSceneDataSource.pause()
-            introduceLockedSim()
-            emulatePendingTransitionProgress(expectedVisible = true)
-            enterSimPin(
+            kosmos.fakeSceneDataSource.pause()
+            kosmos.introduceLockedSim()
+            kosmos.emulatePendingTransitionProgress(expectedVisible = true)
+            kosmos.enterSimPin(
                 authMethodAfterSimUnlock = AuthenticationMethodModel.None,
                 enableLockscreen = false
             )
 
-            assertCurrentScene(Scenes.Gone)
+            kosmos.assertCurrentScene(Scenes.Gone)
         }
 
     @Test
     fun showLockscreen_whenSimUnlocked_whileDeviceLocked() =
         testScope.runTest {
-            fakeSceneDataSource.pause()
-            introduceLockedSim()
-            emulatePendingTransitionProgress(expectedVisible = true)
-            enterSimPin(authMethodAfterSimUnlock = AuthenticationMethodModel.Pin)
-            assertCurrentScene(Scenes.Lockscreen)
+            kosmos.fakeSceneDataSource.pause()
+            kosmos.introduceLockedSim()
+            kosmos.emulatePendingTransitionProgress(expectedVisible = true)
+            kosmos.enterSimPin(authMethodAfterSimUnlock = AuthenticationMethodModel.Pin)
+            kosmos.assertCurrentScene(Scenes.Lockscreen)
         }
 
     /**
@@ -485,8 +409,8 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
      *
      * Note that this doesn't assert what the current scene is in the UI.
      */
-    private fun TestScope.assertCurrentScene(expected: SceneKey) {
-        runCurrent()
+    private fun Kosmos.assertCurrentScene(expected: SceneKey) {
+        testScope.runCurrent()
         assertWithMessage("Current scene mismatch!")
             .that(sceneContainerViewModel.currentScene.value)
             .isEqualTo(expected)
@@ -498,7 +422,7 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
      * This can be different than the value in [SceneContainerViewModel.currentScene], by design, as
      * the UI must gradually transition between scenes.
      */
-    private fun getCurrentSceneInUi(): SceneKey {
+    private fun Kosmos.getCurrentSceneInUi(): SceneKey {
         return when (val state = transitionState.value) {
             is ObservableTransitionState.Idle -> state.currentScene
             is ObservableTransitionState.Transition.ChangeScene -> state.fromScene
@@ -508,7 +432,7 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
     }
 
     /** Updates the current authentication method and related states in the data layer. */
-    private fun TestScope.setAuthMethod(
+    private fun Kosmos.setAuthMethod(
         authMethod: AuthenticationMethodModel,
         enableLockscreen: Boolean = true
     ) {
@@ -520,20 +444,20 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
         // Set the lockscreen enabled bit _before_ set the auth method as the code picks up on the
         // lockscreen enabled bit _after_ the auth method is changed and the lockscreen enabled bit
         // is not an observable that can trigger a new evaluation.
-        kosmos.fakeDeviceEntryRepository.setLockscreenEnabled(enableLockscreen)
-        kosmos.fakeAuthenticationRepository.setAuthenticationMethod(authMethod)
-        runCurrent()
+        fakeDeviceEntryRepository.setLockscreenEnabled(enableLockscreen)
+        fakeAuthenticationRepository.setAuthenticationMethod(authMethod)
+        testScope.runCurrent()
     }
 
     /** Emulates a phone call in progress. */
-    private fun TestScope.startPhoneCall() {
-        whenever(telecomManager.isInCall).thenReturn(true)
-        kosmos.fakeTelephonyRepository.apply {
+    private fun Kosmos.startPhoneCall() {
+        whenever(mockTelecomManager.isInCall).thenReturn(true)
+        fakeTelephonyRepository.apply {
             setHasTelephonyRadio(true)
             setIsInCall(true)
             setCallState(TelephonyManager.CALL_STATE_OFFHOOK)
         }
-        runCurrent()
+        testScope.runCurrent()
     }
 
     /**
@@ -543,14 +467,12 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
      *
      * In order to use this, the [fakeSceneDataSource] must be paused before this method is called.
      */
-    private fun TestScope.emulatePendingTransitionProgress(
-        expectedVisible: Boolean = true,
-    ) {
+    private fun Kosmos.emulatePendingTransitionProgress(expectedVisible: Boolean = true) {
         assertWithMessage("The FakeSceneDataSource has to be paused for this to do anything.")
-            .that(fakeSceneDataSource.isPaused)
+            .that(kosmos.fakeSceneDataSource.isPaused)
             .isTrue()
 
-        val to = fakeSceneDataSource.pendingScene ?: return
+        val to = kosmos.fakeSceneDataSource.pendingScene ?: return
         val from = getCurrentSceneInUi()
 
         if (to == from) {
@@ -568,19 +490,19 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
                 isInitiatedByUserInput = false,
                 isUserInputOngoing = flowOf(false),
             )
-        runCurrent()
+        testScope.runCurrent()
 
         // Report progress of transition.
         while (progressFlow.value < 1f) {
             progressFlow.value += 0.2f
-            runCurrent()
+            testScope.runCurrent()
         }
 
         // End the transition and report the change.
         transitionState.value = ObservableTransitionState.Idle(to)
 
-        fakeSceneDataSource.unpause(force = true)
-        runCurrent()
+        kosmos.fakeSceneDataSource.unpause(force = true)
+        testScope.runCurrent()
 
         assertWithMessage("Visibility mismatch after scene transition from $from to $to!")
             .that(sceneContainerViewModel.isVisible)
@@ -598,7 +520,7 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
                 bouncerSceneJob?.cancel()
                 null
             }
-        runCurrent()
+        testScope.runCurrent()
     }
 
     /**
@@ -610,12 +532,10 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
      *
      * @param to The scene to transition to.
      */
-    private fun TestScope.emulateUserDrivenTransition(
-        to: SceneKey?,
-    ) {
+    private fun Kosmos.emulateUserDrivenTransition(to: SceneKey?) {
         checkNotNull(to)
 
-        fakeSceneDataSource.pause()
+        kosmos.fakeSceneDataSource.pause()
         sceneInteractor.changeScene(to, "reason")
 
         emulatePendingTransitionProgress(
@@ -630,17 +550,17 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
      *
      * Not to be confused with [putDeviceToSleep], which may also instantly lock the device.
      */
-    private suspend fun TestScope.lockDevice() {
+    private suspend fun Kosmos.lockDevice() {
         val authMethod = authenticationInteractor.getAuthenticationMethod()
         assertWithMessage("The authentication method of $authMethod is not secure, cannot lock!")
             .that(authMethod.isSecure)
             .isTrue()
-        kosmos.sceneInteractor.changeScene(Scenes.Lockscreen, "")
-        runCurrent()
+        sceneInteractor.changeScene(Scenes.Lockscreen, "")
+        testScope.runCurrent()
     }
 
     /** Unlocks the device by entering the correct PIN. Ends up in the Gone scene. */
-    private fun TestScope.unlockDevice() {
+    private fun Kosmos.unlockDevice() {
         assertWithMessage("Cannot unlock a device that's already unlocked!")
             .that(deviceEntryInteractor.isUnlocked.value)
             .isFalse()
@@ -662,12 +582,12 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
      *
      * Does not assert that the device is locked or unlocked.
      */
-    private fun TestScope.enterPin() {
+    private fun Kosmos.enterPin() {
         assertWithMessage("Cannot enter PIN when not on the Bouncer scene!")
             .that(getCurrentSceneInUi())
             .isEqualTo(Scenes.Bouncer)
         val authMethodViewModel by
-            collectLastValue(bouncerSceneContentViewModel.authMethodViewModel)
+            testScope.collectLastValue(bouncerSceneContentViewModel.authMethodViewModel)
         assertWithMessage("Cannot enter PIN when not using a PIN authentication method!")
             .that(authMethodViewModel)
             .isInstanceOf(PinBouncerViewModel::class.java)
@@ -677,7 +597,7 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
             pinBouncerViewModel.onPinButtonClicked(digit)
         }
         pinBouncerViewModel.onAuthenticateButtonClicked()
-        runCurrent()
+        testScope.runCurrent()
     }
 
     /**
@@ -688,7 +608,7 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
      *
      * Does not assert that the device is locked or unlocked.
      */
-    private fun TestScope.enterSimPin(
+    private fun Kosmos.enterSimPin(
         authMethodAfterSimUnlock: AuthenticationMethodModel = AuthenticationMethodModel.None,
         enableLockscreen: Boolean = true,
     ) {
@@ -696,7 +616,7 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
             .that(getCurrentSceneInUi())
             .isEqualTo(Scenes.Bouncer)
         val authMethodViewModel by
-            collectLastValue(bouncerSceneContentViewModel.authMethodViewModel)
+            testScope.collectLastValue(bouncerSceneContentViewModel.authMethodViewModel)
         assertWithMessage("Cannot enter PIN when not using a PIN authentication method!")
             .that(authMethodViewModel)
             .isInstanceOf(PinBouncerViewModel::class.java)
@@ -706,26 +626,26 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
             pinBouncerViewModel.onPinButtonClicked(digit)
         }
         pinBouncerViewModel.onAuthenticateButtonClicked()
-        kosmos.fakeMobileConnectionsRepository.isAnySimSecure.value = false
-        runCurrent()
+        fakeMobileConnectionsRepository.isAnySimSecure.value = false
+        testScope.runCurrent()
 
         setAuthMethod(authMethodAfterSimUnlock, enableLockscreen)
-        runCurrent()
+        testScope.runCurrent()
     }
 
     /** Changes device wakefulness state from asleep to awake, going through intermediary states. */
-    private fun TestScope.wakeUpDevice() {
+    private fun Kosmos.wakeUpDevice() {
         val wakefulnessModel = powerInteractor.detailedWakefulness.value
         assertWithMessage("Cannot wake up device as it's already awake!")
             .that(wakefulnessModel.isAwake())
             .isFalse()
 
         powerInteractor.setAwakeForTest()
-        runCurrent()
+        testScope.runCurrent()
     }
 
     /** Changes device wakefulness state from awake to asleep, going through intermediary states. */
-    private suspend fun TestScope.putDeviceToSleep(
+    private suspend fun Kosmos.putDeviceToSleep(
         instantlyLockDevice: Boolean = true,
     ) {
         val wakefulnessModel = powerInteractor.detailedWakefulness.value
@@ -734,7 +654,7 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
             .isTrue()
 
         powerInteractor.setAsleepForTest()
-        runCurrent()
+        testScope.runCurrent()
 
         if (instantlyLockDevice) {
             lockDevice()
@@ -742,16 +662,16 @@ class SceneFrameworkIntegrationTest : SysuiTestCase() {
     }
 
     /** Emulates the dismissal of the IME (soft keyboard). */
-    private fun TestScope.dismissIme() {
+    private fun Kosmos.dismissIme() {
         (bouncerSceneContentViewModel.authMethodViewModel.value as? PasswordBouncerViewModel)?.let {
             it.onImeDismissed()
-            runCurrent()
+            testScope.runCurrent()
         }
     }
 
-    private fun TestScope.introduceLockedSim() {
+    private fun Kosmos.introduceLockedSim() {
         setAuthMethod(AuthenticationMethodModel.Sim)
-        kosmos.fakeMobileConnectionsRepository.isAnySimSecure.value = true
-        runCurrent()
+        fakeMobileConnectionsRepository.isAnySimSecure.value = true
+        testScope.runCurrent()
     }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenUserActionsViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenUserActionsViewModelKosmos.kt
new file mode 100644
index 000000000000..a25b29fd18cb
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenUserActionsViewModelKosmos.kt
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 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.keyguard.ui.viewmodel
+
+import com.android.systemui.communal.domain.interactor.communalInteractor
+import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.shade.domain.interactor.shadeInteractor
+
+val Kosmos.lockscreenUserActionsViewModel by Fixture {
+    LockscreenUserActionsViewModel(
+        deviceEntryInteractor = deviceEntryInteractor,
+        communalInteractor = communalInteractor,
+        shadeInteractor = shadeInteractor,
+    )
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneKosmos.kt
index b3664e107d72..db43fcf080ec 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneKosmos.kt
@@ -1,11 +1,18 @@
 package com.android.systemui.scene
 
+import com.android.compose.animation.scene.ObservableTransitionState
+import com.android.systemui.classifier.domain.interactor.falsingInteractor
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.power.domain.interactor.powerInteractor
+import com.android.systemui.scene.domain.interactor.sceneInteractor
+import com.android.systemui.scene.shared.logger.sceneLogger
 import com.android.systemui.scene.shared.model.Overlays
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.scene.ui.FakeOverlay
+import com.android.systemui.scene.ui.viewmodel.SceneContainerViewModel
+import kotlinx.coroutines.flow.MutableStateFlow
 
 var Kosmos.sceneKeys by Fixture {
     listOf(
@@ -49,3 +56,20 @@ var Kosmos.sceneContainerConfig by Fixture {
         navigationDistances = navigationDistances,
     )
 }
+
+val Kosmos.transitionState by Fixture {
+    MutableStateFlow<ObservableTransitionState>(
+        ObservableTransitionState.Idle(sceneContainerConfig.initialSceneKey)
+    )
+}
+
+val Kosmos.sceneContainerViewModel by Fixture {
+    SceneContainerViewModel(
+            sceneInteractor = sceneInteractor,
+            falsingInteractor = falsingInteractor,
+            powerInteractor = powerInteractor,
+            motionEventHandlerReceiver = {},
+            logger = sceneLogger
+        )
+        .apply { setTransitionState(transitionState) }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/telecom/TelecomManagerKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/telecom/TelecomManagerKosmos.kt
index 4e0c0883eb02..d1bee6149cb0 100644
--- a/packages/SystemUI/tests/utils/src/com/android/telecom/TelecomManagerKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/telecom/TelecomManagerKosmos.kt
@@ -21,4 +21,5 @@ import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
 import com.android.systemui.util.mockito.mock
 
-var Kosmos.telecomManager by Fixture<TelecomManager?> { mock() }
+val Kosmos.mockTelecomManager by Fixture<TelecomManager> { mock() }
+var Kosmos.telecomManager by Fixture<TelecomManager?> { mockTelecomManager }