Merge "Introduce Kosmos.runTest" into main
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 9db2014..319f1e5 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
@@ -44,6 +44,8 @@
 import com.android.systemui.flags.fakeFeatureFlagsClassic
 import com.android.systemui.keyguard.ui.viewmodel.lockscreenUserActionsViewModel
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.runCurrent
+import com.android.systemui.kosmos.runTest
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.lifecycle.activateIn
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAsleepForTest
@@ -70,7 +72,6 @@
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.runCurrent
-import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -101,7 +102,6 @@
 @EnableSceneContainer
 class SceneFrameworkIntegrationTest : SysuiTestCase() {
     private val kosmos = testKosmos()
-    private val testScope = kosmos.testScope
     private var bouncerSceneJob: Job? = null
 
     @Before
@@ -137,151 +137,149 @@
                 .isTrue()
         }
 
-    @Test
-    fun startsInLockscreenScene() =
-        testScope.runTest { kosmos.assertCurrentScene(Scenes.Lockscreen) }
+    @Test fun startsInLockscreenScene() = kosmos.runTest { assertCurrentScene(Scenes.Lockscreen) }
 
     @Test
     fun clickLockButtonAndEnterCorrectPin_unlocksDevice() =
-        testScope.runTest {
-            kosmos.emulateUserDrivenTransition(Scenes.Bouncer)
+        kosmos.runTest {
+            emulateUserDrivenTransition(Scenes.Bouncer)
 
-            kosmos.fakeSceneDataSource.pause()
-            kosmos.enterPin()
-            kosmos.emulatePendingTransitionProgress(expectedVisible = false)
-            kosmos.assertCurrentScene(Scenes.Gone)
+            fakeSceneDataSource.pause()
+            enterPin()
+            emulatePendingTransitionProgress(expectedVisible = false)
+            assertCurrentScene(Scenes.Gone)
         }
 
     @Test
     fun swipeUpOnLockscreen_enterCorrectPin_unlocksDevice() =
-        testScope.runTest {
-            val actions by collectLastValue(kosmos.lockscreenUserActionsViewModel.actions)
+        kosmos.runTest {
+            val actions by testScope.collectLastValue(kosmos.lockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Bouncer)
-            kosmos.emulateUserDrivenTransition(to = upDestinationSceneKey)
+            emulateUserDrivenTransition(to = upDestinationSceneKey)
 
-            kosmos.fakeSceneDataSource.pause()
-            kosmos.enterPin()
-            kosmos.emulatePendingTransitionProgress(expectedVisible = false)
-            kosmos.assertCurrentScene(Scenes.Gone)
+            fakeSceneDataSource.pause()
+            enterPin()
+            emulatePendingTransitionProgress(expectedVisible = false)
+            assertCurrentScene(Scenes.Gone)
         }
 
     @Test
     fun swipeUpOnLockscreen_withAuthMethodSwipe_dismissesLockscreen() =
-        testScope.runTest {
-            kosmos.setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
+        kosmos.runTest {
+            setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
 
-            val actions by collectLastValue(kosmos.lockscreenUserActionsViewModel.actions)
+            val actions by testScope.collectLastValue(lockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Gone)
-            kosmos.emulateUserDrivenTransition(to = upDestinationSceneKey)
+            emulateUserDrivenTransition(to = upDestinationSceneKey)
         }
 
     @Test
     fun swipeUpOnShadeScene_withAuthMethodSwipe_lockscreenNotDismissed_goesToLockscreen() =
-        testScope.runTest {
-            val actions by collectLastValue(kosmos.shadeUserActionsViewModel.actions)
-            kosmos.setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
-            kosmos.assertCurrentScene(Scenes.Lockscreen)
+        kosmos.runTest {
+            val actions by testScope.collectLastValue(shadeUserActionsViewModel.actions)
+            setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
+            assertCurrentScene(Scenes.Lockscreen)
 
             // Emulate a user swipe to the shade scene.
-            kosmos.emulateUserDrivenTransition(to = Scenes.Shade)
-            kosmos.assertCurrentScene(Scenes.Shade)
+            emulateUserDrivenTransition(to = Scenes.Shade)
+            assertCurrentScene(Scenes.Shade)
 
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Lockscreen)
-            kosmos.emulateUserDrivenTransition(to = Scenes.Lockscreen)
+            emulateUserDrivenTransition(to = Scenes.Lockscreen)
         }
 
     @Test
     fun swipeUpOnShadeScene_withAuthMethodSwipe_lockscreenDismissed_goesToGone() =
-        testScope.runTest {
-            val actions by collectLastValue(kosmos.shadeUserActionsViewModel.actions)
-            val canSwipeToEnter by collectLastValue(kosmos.deviceEntryInteractor.canSwipeToEnter)
+        kosmos.runTest {
+            val actions by testScope.collectLastValue(shadeUserActionsViewModel.actions)
+            val canSwipeToEnter by testScope.collectLastValue(deviceEntryInteractor.canSwipeToEnter)
 
-            kosmos.setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
+            setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
 
             assertThat(canSwipeToEnter).isTrue()
-            kosmos.assertCurrentScene(Scenes.Lockscreen)
+            assertCurrentScene(Scenes.Lockscreen)
 
             // Emulate a user swipe to dismiss the lockscreen.
-            kosmos.emulateUserDrivenTransition(to = Scenes.Gone)
-            kosmos.assertCurrentScene(Scenes.Gone)
+            emulateUserDrivenTransition(to = Scenes.Gone)
+            assertCurrentScene(Scenes.Gone)
 
             // Emulate a user swipe to the shade scene.
-            kosmos.emulateUserDrivenTransition(to = Scenes.Shade)
-            kosmos.assertCurrentScene(Scenes.Shade)
+            emulateUserDrivenTransition(to = Scenes.Shade)
+            assertCurrentScene(Scenes.Shade)
 
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Gone)
-            kosmos.emulateUserDrivenTransition(to = Scenes.Gone)
+            emulateUserDrivenTransition(to = Scenes.Gone)
         }
 
     @Test
     fun withAuthMethodNone_deviceWakeUp_skipsLockscreen() =
-        testScope.runTest {
-            kosmos.setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = false)
-            kosmos.putDeviceToSleep()
-            kosmos.assertCurrentScene(Scenes.Lockscreen)
+        kosmos.runTest {
+            setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = false)
+            putDeviceToSleep()
+            assertCurrentScene(Scenes.Lockscreen)
 
-            kosmos.wakeUpDevice()
-            kosmos.assertCurrentScene(Scenes.Gone)
+            wakeUpDevice()
+            assertCurrentScene(Scenes.Gone)
         }
 
     @Test
     fun withAuthMethodSwipe_deviceWakeUp_doesNotSkipLockscreen() =
-        testScope.runTest {
-            kosmos.setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
-            kosmos.putDeviceToSleep()
-            kosmos.assertCurrentScene(Scenes.Lockscreen)
+        kosmos.runTest {
+            setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
+            putDeviceToSleep()
+            assertCurrentScene(Scenes.Lockscreen)
 
-            kosmos.wakeUpDevice()
-            kosmos.assertCurrentScene(Scenes.Lockscreen)
+            wakeUpDevice()
+            assertCurrentScene(Scenes.Lockscreen)
         }
 
     @Test
     fun lockDeviceLocksDevice() =
-        testScope.runTest {
-            kosmos.unlockDevice()
-            kosmos.assertCurrentScene(Scenes.Gone)
+        kosmos.runTest {
+            unlockDevice()
+            assertCurrentScene(Scenes.Gone)
 
-            kosmos.lockDevice()
-            kosmos.assertCurrentScene(Scenes.Lockscreen)
+            lockDevice()
+            assertCurrentScene(Scenes.Lockscreen)
         }
 
     @Test
     fun deviceGoesToSleep_switchesToLockscreen() =
-        testScope.runTest {
-            kosmos.unlockDevice()
-            kosmos.assertCurrentScene(Scenes.Gone)
+        kosmos.runTest {
+            unlockDevice()
+            assertCurrentScene(Scenes.Gone)
 
-            kosmos.putDeviceToSleep()
-            kosmos.assertCurrentScene(Scenes.Lockscreen)
+            putDeviceToSleep()
+            assertCurrentScene(Scenes.Lockscreen)
         }
 
     @Test
     fun deviceGoesToSleep_wakeUp_unlock() =
-        testScope.runTest {
-            kosmos.unlockDevice()
-            kosmos.assertCurrentScene(Scenes.Gone)
-            kosmos.putDeviceToSleep()
-            kosmos.assertCurrentScene(Scenes.Lockscreen)
-            kosmos.wakeUpDevice()
-            kosmos.assertCurrentScene(Scenes.Lockscreen)
+        kosmos.runTest {
+            unlockDevice()
+            assertCurrentScene(Scenes.Gone)
+            putDeviceToSleep()
+            assertCurrentScene(Scenes.Lockscreen)
+            wakeUpDevice()
+            assertCurrentScene(Scenes.Lockscreen)
 
-            kosmos.unlockDevice()
-            kosmos.assertCurrentScene(Scenes.Gone)
+            unlockDevice()
+            assertCurrentScene(Scenes.Gone)
         }
 
     @Test
     fun swipeUpOnLockscreenWhileUnlocked_dismissesLockscreen() =
-        testScope.runTest {
-            kosmos.unlockDevice()
-            val actions by collectLastValue(kosmos.lockscreenUserActionsViewModel.actions)
+        kosmos.runTest {
+            unlockDevice()
+            val actions by testScope.collectLastValue(lockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Gone)
@@ -289,46 +287,46 @@
 
     @Test
     fun deviceGoesToSleep_withLockTimeout_staysOnLockscreen() =
-        testScope.runTest {
-            kosmos.unlockDevice()
-            kosmos.assertCurrentScene(Scenes.Gone)
-            kosmos.putDeviceToSleep()
-            kosmos.assertCurrentScene(Scenes.Lockscreen)
+        kosmos.runTest {
+            unlockDevice()
+            assertCurrentScene(Scenes.Gone)
+            putDeviceToSleep()
+            assertCurrentScene(Scenes.Lockscreen)
 
             // Pretend like the timeout elapsed and now lock the device.
-            kosmos.lockDevice()
-            kosmos.assertCurrentScene(Scenes.Lockscreen)
+            lockDevice()
+            assertCurrentScene(Scenes.Lockscreen)
         }
 
     @Test
     fun dismissingIme_whileOnPasswordBouncer_navigatesToLockscreen() =
-        testScope.runTest {
-            kosmos.setAuthMethod(AuthenticationMethodModel.Password)
-            val actions by collectLastValue(kosmos.lockscreenUserActionsViewModel.actions)
+        kosmos.runTest {
+            setAuthMethod(AuthenticationMethodModel.Password)
+            val actions by testScope.collectLastValue(lockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Bouncer)
-            kosmos.emulateUserDrivenTransition(to = upDestinationSceneKey)
+            emulateUserDrivenTransition(to = upDestinationSceneKey)
 
-            kosmos.fakeSceneDataSource.pause()
-            kosmos.dismissIme()
+            fakeSceneDataSource.pause()
+            dismissIme()
 
-            kosmos.emulatePendingTransitionProgress()
-            kosmos.assertCurrentScene(Scenes.Lockscreen)
+            emulatePendingTransitionProgress()
+            assertCurrentScene(Scenes.Lockscreen)
         }
 
     @Test
     fun bouncerActionButtonClick_opensEmergencyServicesDialer() =
-        testScope.runTest {
-            kosmos.setAuthMethod(AuthenticationMethodModel.Password)
-            val actions by collectLastValue(kosmos.lockscreenUserActionsViewModel.actions)
+        kosmos.runTest {
+            setAuthMethod(AuthenticationMethodModel.Password)
+            val actions by testScope.collectLastValue(lockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Bouncer)
-            kosmos.emulateUserDrivenTransition(to = upDestinationSceneKey)
+            emulateUserDrivenTransition(to = upDestinationSceneKey)
 
             val bouncerActionButton by
-                collectLastValue(kosmos.bouncerSceneContentViewModel.actionButton)
+                testScope.collectLastValue(bouncerSceneContentViewModel.actionButton)
             assertWithMessage("Bouncer action button not visible")
                 .that(bouncerActionButton)
                 .isNotNull()
@@ -340,56 +338,56 @@
 
     @Test
     fun bouncerActionButtonClick_duringCall_returnsToCall() =
-        testScope.runTest {
-            kosmos.setAuthMethod(AuthenticationMethodModel.Password)
-            kosmos.startPhoneCall()
-            val actions by collectLastValue(kosmos.lockscreenUserActionsViewModel.actions)
+        kosmos.runTest {
+            setAuthMethod(AuthenticationMethodModel.Password)
+            startPhoneCall()
+            val actions by testScope.collectLastValue(lockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Bouncer)
-            kosmos.emulateUserDrivenTransition(to = upDestinationSceneKey)
+            emulateUserDrivenTransition(to = upDestinationSceneKey)
 
             val bouncerActionButton by
-                collectLastValue(kosmos.bouncerSceneContentViewModel.actionButton)
+                testScope.collectLastValue(bouncerSceneContentViewModel.actionButton)
             assertWithMessage("Bouncer action button not visible during call")
                 .that(bouncerActionButton)
                 .isNotNull()
             kosmos.bouncerSceneContentViewModel.onActionButtonClicked(bouncerActionButton!!)
             runCurrent()
 
-            verify(kosmos.mockTelecomManager).showInCallScreen(any())
+            verify(mockTelecomManager).showInCallScreen(any())
         }
 
     @Test
     fun showBouncer_whenLockedSimIntroduced() =
-        testScope.runTest {
-            kosmos.setAuthMethod(AuthenticationMethodModel.None)
-            kosmos.introduceLockedSim()
-            kosmos.assertCurrentScene(Scenes.Bouncer)
+        kosmos.runTest {
+            setAuthMethod(AuthenticationMethodModel.None)
+            introduceLockedSim()
+            assertCurrentScene(Scenes.Bouncer)
         }
 
     @Test
     fun goesToGone_whenSimUnlocked_whileDeviceUnlocked() =
-        testScope.runTest {
-            kosmos.fakeSceneDataSource.pause()
-            kosmos.introduceLockedSim()
-            kosmos.emulatePendingTransitionProgress(expectedVisible = true)
-            kosmos.enterSimPin(
+        kosmos.runTest {
+            fakeSceneDataSource.pause()
+            introduceLockedSim()
+            emulatePendingTransitionProgress(expectedVisible = true)
+            enterSimPin(
                 authMethodAfterSimUnlock = AuthenticationMethodModel.None,
                 enableLockscreen = false,
             )
 
-            kosmos.assertCurrentScene(Scenes.Gone)
+            assertCurrentScene(Scenes.Gone)
         }
 
     @Test
     fun showLockscreen_whenSimUnlocked_whileDeviceLocked() =
-        testScope.runTest {
-            kosmos.fakeSceneDataSource.pause()
-            kosmos.introduceLockedSim()
-            kosmos.emulatePendingTransitionProgress(expectedVisible = true)
-            kosmos.enterSimPin(authMethodAfterSimUnlock = AuthenticationMethodModel.Pin)
-            kosmos.assertCurrentScene(Scenes.Lockscreen)
+        kosmos.runTest {
+            fakeSceneDataSource.pause()
+            introduceLockedSim()
+            emulatePendingTransitionProgress(expectedVisible = true)
+            enterSimPin(authMethodAfterSimUnlock = AuthenticationMethodModel.Pin)
+            assertCurrentScene(Scenes.Lockscreen)
         }
 
     /**
@@ -457,10 +455,10 @@
      */
     private fun Kosmos.emulatePendingTransitionProgress(expectedVisible: Boolean = true) {
         assertWithMessage("The FakeSceneDataSource has to be paused for this to do anything.")
-            .that(kosmos.fakeSceneDataSource.isPaused)
+            .that(fakeSceneDataSource.isPaused)
             .isTrue()
 
-        val to = kosmos.fakeSceneDataSource.pendingScene ?: return
+        val to = fakeSceneDataSource.pendingScene ?: return
         val from = getCurrentSceneInUi()
 
         if (to == from) {
@@ -489,7 +487,7 @@
         // End the transition and report the change.
         transitionState.value = ObservableTransitionState.Idle(to)
 
-        kosmos.fakeSceneDataSource.unpause(force = true)
+        fakeSceneDataSource.unpause(force = true)
         testScope.runCurrent()
 
         assertWithMessage("Visibility mismatch after scene transition from $from to $to!")
@@ -523,7 +521,7 @@
     private fun Kosmos.emulateUserDrivenTransition(to: SceneKey?) {
         checkNotNull(to)
 
-        kosmos.fakeSceneDataSource.pause()
+        fakeSceneDataSource.pause()
         sceneInteractor.changeScene(to, "reason")
 
         emulatePendingTransitionProgress(expectedVisible = to != Scenes.Gone)
@@ -634,7 +632,7 @@
     }
 
     /** Changes device wakefulness state from awake to asleep, going through intermediary states. */
-    private suspend fun Kosmos.putDeviceToSleep() {
+    private fun Kosmos.putDeviceToSleep() {
         val wakefulnessModel = powerInteractor.detailedWakefulness.value
         assertWithMessage("Cannot put device to sleep as it's already asleep!")
             .that(wakefulnessModel.isAwake())
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/GeneralKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/GeneralKosmos.kt
index a9a80b5..ddae581 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/GeneralKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/GeneralKosmos.kt
@@ -6,6 +6,8 @@
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
 
 var Kosmos.testDispatcher by Fixture { StandardTestDispatcher() }
 
@@ -34,3 +36,12 @@
     testScope.backgroundScope.coroutineContext
 }
 var Kosmos.mainCoroutineContext: CoroutineContext by Fixture { testScope.coroutineContext }
+
+/**
+ * Run this test body with a [Kosmos] as receiver, and using the [testScope] currently installed in
+ * that kosmos instance
+ */
+fun Kosmos.runTest(testBody: suspend Kosmos.() -> Unit) =
+    testScope.runTest { this@runTest.testBody() }
+
+fun Kosmos.runCurrent() = testScope.runCurrent()