Moves runTestWithSnapshots to receive a TestScope instead of a Kosmos.

- Placed in a new compose package under tests/utils so it's not
  perceived as part of the Kosmos package
- Kosmos.runTest automatically uses this as well so nobody has to worry
  about it
- Converts existing usages to Kosmos.runTest

Bug: 387321113
Flag: TEST_ONLY
Test: irrelevant
Change-Id: I8d8f621d303ab422258fef62dfee94c32f95c3cd

3 files changed, 52 insertions, 28 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/view/SceneJankMonitorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/view/SceneJankMonitorTest.kt
index 19369e7b343d..984f8fd13cde 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/view/SceneJankMonitorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/view/SceneJankMonitorTest.kt
@@ -28,7 +28,7 @@ import com.android.systemui.keyguard.data.repository.fakeDeviceEntryFingerprintA
 import com.android.systemui.keyguard.shared.model.SuccessFingerprintAuthenticationStatus
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.runCurrent
-import com.android.systemui.kosmos.runTestWithSnapshots
+import com.android.systemui.kosmos.runTest
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.lifecycle.activateIn
 import com.android.systemui.scene.shared.model.Scenes
@@ -59,7 +59,7 @@ class SceneJankMonitorTest : SysuiTestCase() {
 
     @Test
     fun onTransitionStart_withProvidedCuj_beginsThatCuj() =
-        kosmos.runTestWithSnapshots {
+        kosmos.runTest {
             val cuj = 1337
             underTest.onTransitionStart(
                 view = mock(),
@@ -73,7 +73,7 @@ class SceneJankMonitorTest : SysuiTestCase() {
 
     @Test
     fun onTransitionEnd_withProvidedCuj_endsThatCuj() =
-        kosmos.runTestWithSnapshots {
+        kosmos.runTest {
             val cuj = 1337
             underTest.onTransitionEnd(from = Scenes.Communal, to = Scenes.Dream, cuj = cuj)
             verify(interactionJankMonitor, never()).begin(any(), anyInt())
@@ -82,7 +82,7 @@ class SceneJankMonitorTest : SysuiTestCase() {
 
     @Test
     fun bouncer_authMethodPin() =
-        kosmos.runTestWithSnapshots {
+        kosmos.runTest {
             bouncer(
                 authenticationMethod = AuthenticationMethodModel.Pin,
                 appearCuj = Cuj.CUJ_LOCKSCREEN_PIN_APPEAR,
@@ -92,7 +92,7 @@ class SceneJankMonitorTest : SysuiTestCase() {
 
     @Test
     fun bouncer_authMethodSim() =
-        kosmos.runTestWithSnapshots {
+        kosmos.runTest {
             bouncer(
                 authenticationMethod = AuthenticationMethodModel.Sim,
                 appearCuj = Cuj.CUJ_LOCKSCREEN_PIN_APPEAR,
@@ -109,7 +109,7 @@ class SceneJankMonitorTest : SysuiTestCase() {
 
     @Test
     fun bouncer_authMethodPattern() =
-        kosmos.runTestWithSnapshots {
+        kosmos.runTest {
             bouncer(
                 authenticationMethod = AuthenticationMethodModel.Pattern,
                 appearCuj = Cuj.CUJ_LOCKSCREEN_PATTERN_APPEAR,
@@ -119,7 +119,7 @@ class SceneJankMonitorTest : SysuiTestCase() {
 
     @Test
     fun bouncer_authMethodPassword() =
-        kosmos.runTestWithSnapshots {
+        kosmos.runTest {
             bouncer(
                 authenticationMethod = AuthenticationMethodModel.Password,
                 appearCuj = Cuj.CUJ_LOCKSCREEN_PASSWORD_APPEAR,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/compose/Snapshot.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/compose/Snapshot.kt
new file mode 100644
index 000000000000..fb6699c44d62
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/compose/Snapshot.kt
@@ -0,0 +1,41 @@
+/*
+ * Copyright (C) 2025 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.compose
+
+import androidx.compose.runtime.snapshots.Snapshot
+import com.android.systemui.kosmos.runCurrent
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+
+/**
+ * Runs the given test [block] in a [TestScope] that's set up such that the Compose snapshot state
+ * is settled eagerly. This is the Compose equivalent to using an [UnconfinedTestDispatcher] or
+ * using [runCurrent] a lot.
+ *
+ * Note that this shouldn't be needed or used in a Compose test environment.
+ */
+fun TestScope.runTestWithSnapshots(block: suspend TestScope.() -> Unit) {
+    val handle = Snapshot.registerGlobalWriteObserver { Snapshot.sendApplyNotifications() }
+
+    try {
+        runTest { block() }
+    } finally {
+        handle.dispose()
+    }
+}
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 1881a94c8984..afe48214832f 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
@@ -1,7 +1,7 @@
 package com.android.systemui.kosmos
 
-import androidx.compose.runtime.snapshots.Snapshot
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.compose.runTestWithSnapshots
 import com.android.systemui.coroutines.FlowValue
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
@@ -17,7 +17,6 @@ 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
 import org.mockito.kotlin.verify
 
 var Kosmos.testDispatcher by Fixture { StandardTestDispatcher() }
@@ -53,26 +52,10 @@ var Kosmos.brightnessWarningToast: BrightnessWarningToast by
 
 /**
  * Run this test body with a [Kosmos] as receiver, and using the [testScope] currently installed in
- * that kosmos instance
+ * that Kosmos instance
  */
-fun Kosmos.runTest(testBody: suspend Kosmos.() -> Unit) =
-    testScope.runTest testBody@{ this@runTest.testBody() }
-
-/**
- * Runs the given [Kosmos]-scoped test [block] in an environment where compose snapshot state is
- * settled eagerly. This is the compose equivalent to using an [UnconfinedTestDispatcher] or using
- * [runCurrent] a lot.
- *
- * Note that this shouldn't be needed or used in a compose test environment.
- */
-fun Kosmos.runTestWithSnapshots(block: suspend Kosmos.() -> Unit) {
-    val handle = Snapshot.registerGlobalWriteObserver { Snapshot.sendApplyNotifications() }
-
-    try {
-        testScope.runTest { block() }
-    } finally {
-        handle.dispose()
-    }
+fun Kosmos.runTest(testBody: suspend Kosmos.() -> Unit) {
+    testScope.runTestWithSnapshots testBody@{ this@runTest.testBody() }
 }
 
 fun Kosmos.runCurrent() = testScope.runCurrent()