diff --git a/packages/SystemUI/Android.bp b/packages/SystemUI/Android.bp
index 2529157..ed62c5f 100644
--- a/packages/SystemUI/Android.bp
+++ b/packages/SystemUI/Android.bp
@@ -204,6 +204,45 @@
     path: "tests/utils/src",
 }
 
+filegroup {
+    name: "SystemUI-tests-robolectric-pilots",
+    srcs: [
+        // data
+        "tests/src/com/android/systemui/keyguard/data/quickaffordance/CameraQuickAffordanceConfigTest.kt",
+        "tests/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfigTest.kt",
+        "tests/src/com/android/systemui/keyguard/data/quickaffordance/FakeKeyguardQuickAffordanceConfig.kt",
+        "tests/src/com/android/systemui/keyguard/data/quickaffordance/FlashlightQuickAffordanceConfigTest.kt",
+        "tests/src/com/android/systemui/keyguard/data/quickaffordance/HomeControlsKeyguardQuickAffordanceConfigTest.kt",
+        "tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceLegacySettingSyncerTest.kt",
+        "tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceLocalUserSelectionManagerTest.kt",
+        "tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceRemoteUserSelectionManagerTest.kt",
+        "tests/src/com/android/systemui/keyguard/data/quickaffordance/QrCodeScannerKeyguardQuickAffordanceConfigTest.kt",
+        "tests/src/com/android/systemui/keyguard/data/quickaffordance/QuickAccessWalletKeyguardQuickAffordanceConfigTest.kt",
+        "tests/src/com/android/systemui/keyguard/data/repository/KeyguardQuickAffordanceRepositoryTest.kt",
+        "tests/src/com/android/systemui/keyguard/data/repository/KeyguardRepositoryImplTest.kt",
+        // domain
+        "tests/src/com/android/systemui/keyguard/domain/interactor/AlternateBouncerInteractorTest.kt",
+        "tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt",
+        "tests/src/com/android/systemui/keyguard/domain/interactor/AlternateBouncerInteractorTest.kt",
+        "tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt",
+        "tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt",
+        "tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt",
+        "tests/src/com/android/systemui/keyguard/domain/interactor/LightRevealScrimInteractorTest.kt",
+        "tests/src/com/android/systemui/keyguard/domain/quickaffordance/FakeKeyguardQuickAffordanceRegistry.kt",
+        "tests/src/com/android/systemui/keyguard/domain/interactor/PrimaryBouncerInteractorWithCoroutinesTest.kt",
+        "tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt",
+        "tests/src/com/android/systemui/keyguard/domain/interactor/PrimaryBouncerInteractorWithCoroutinesTest.kt",
+        // ui
+        "tests/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToLockscreenTransitionViewModelTest.kt",
+        "tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBouncerViewModelTest.kt",
+        "tests/src/com/android/systemui/keyguard/ui/viewmodel/GoneToDreamingTransitionViewModelTest.kt",
+        "tests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToDreamingTransitionViewModelTest.kt",
+        "tests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToOccludedTransitionViewModelTest.kt",
+        "tests/src/com/android/systemui/keyguard/ui/viewmodel/OccludedToLockscreenTransitionViewModelTest.kt",
+    ],
+    path: "tests/src",
+}
+
 java_library {
     name: "SystemUI-tests-concurrency",
     srcs: [
@@ -315,8 +354,16 @@
     defaults: [
         "platform_app_defaults",
         "SystemUI_app_defaults",
+        "SystemUI_compose_defaults",
     ],
     manifest: "tests/AndroidManifest-base.xml",
+
+    srcs: [
+        "src/**/*.kt",
+        "src/**/*.java",
+        "src/**/I*.aidl",
+        ":ReleaseJavaFiles",
+    ],
     static_libs: [
         "SystemUI-tests-base",
     ],
@@ -330,6 +377,9 @@
     certificate: "platform",
     privileged: true,
     resource_dirs: [],
+    kotlincflags: ["-Xjvm-default=all"],
+
+    plugins: ["dagger2-compiler"],
 }
 
 android_robolectric_test {
@@ -337,6 +387,13 @@
     srcs: [
         "tests/robolectric/src/**/*.kt",
         "tests/robolectric/src/**/*.java",
+        ":SystemUI-tests-utils",
+        ":SystemUI-tests-robolectric-pilots",
+    ],
+    static_libs: [
+        "androidx.test.uiautomator_uiautomator",
+        "androidx.test.ext.junit",
+        "inline-mockito-robolectric-prebuilt",
     ],
     libs: [
         "android.test.runner",
@@ -344,7 +401,9 @@
         "android.test.mock",
         "truth-prebuilt",
     ],
-    kotlincflags: ["-Xjvm-default=enable"],
+
+    upstream: true,
+
     instrumentation_for: "SystemUIRobo-stub",
     java_resource_dirs: ["tests/robolectric/config"],
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/CameraQuickAffordanceConfigTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/CameraQuickAffordanceConfigTest.kt
index 58cdec4..db18ba6 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/CameraQuickAffordanceConfigTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/CameraQuickAffordanceConfigTest.kt
@@ -20,6 +20,7 @@
 import android.app.StatusBarManager
 import android.content.Context
 import android.content.pm.PackageManager
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.camera.CameraGestureHelper
@@ -31,14 +32,13 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.Mock
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class CameraQuickAffordanceConfigTest : SysuiTestCase() {
 
     @Mock private lateinit var cameraGestureHelper: CameraGestureHelper
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfigTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfigTest.kt
index 15b85de..64839e2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfigTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfigTest.kt
@@ -22,6 +22,7 @@
 import android.provider.Settings.Global.ZEN_MODE_OFF
 import android.provider.Settings.Secure.ZEN_DURATION_FOREVER
 import android.provider.Settings.Secure.ZEN_DURATION_PROMPT
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.settingslib.notification.EnableZenModeDialog
 import com.android.systemui.R
@@ -51,7 +52,6 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.ArgumentCaptor
 import org.mockito.Captor
 import org.mockito.Mock
@@ -60,7 +60,7 @@
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class DoNotDisturbQuickAffordanceConfigTest : SysuiTestCase() {
 
     @Mock private lateinit var zenModeController: ZenModeController
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/FlashlightQuickAffordanceConfigTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/FlashlightQuickAffordanceConfigTest.kt
index 9fa7db1..31391ee 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/FlashlightQuickAffordanceConfigTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/FlashlightQuickAffordanceConfigTest.kt
@@ -18,6 +18,7 @@
 package com.android.systemui.keyguard.data.quickaffordance
 
 import android.content.Context
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.R
 import com.android.systemui.common.shared.model.Icon
@@ -35,13 +36,12 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.Mock
 import org.mockito.MockitoAnnotations
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class FlashlightQuickAffordanceConfigTest : LeakCheckedTest() {
 
     @Mock private lateinit var context: Context
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/HomeControlsKeyguardQuickAffordanceConfigTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/HomeControlsKeyguardQuickAffordanceConfigTest.kt
index 659c1e5..2c1c04c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/HomeControlsKeyguardQuickAffordanceConfigTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/HomeControlsKeyguardQuickAffordanceConfigTest.kt
@@ -17,6 +17,7 @@
 
 package com.android.systemui.keyguard.data.quickaffordance
 
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
@@ -34,13 +35,12 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.Mock
 import org.mockito.Mockito.`when` as whenever
 import org.mockito.MockitoAnnotations
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class HomeControlsKeyguardQuickAffordanceConfigTest : SysuiTestCase() {
 
     @Mock private lateinit var component: ControlsComponent
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceLegacySettingSyncerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceLegacySettingSyncerTest.kt
index 3b0169d..3bae7f7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceLegacySettingSyncerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceLegacySettingSyncerTest.kt
@@ -20,6 +20,7 @@
 import android.content.Context
 import android.content.res.Resources
 import android.provider.Settings
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
@@ -40,7 +41,6 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.ArgumentMatchers.anyString
 import org.mockito.Mock
@@ -48,7 +48,7 @@
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class KeyguardQuickAffordanceLegacySettingSyncerTest : SysuiTestCase() {
 
     @Mock private lateinit var sharedPrefs: FakeSharedPreferences
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceLocalUserSelectionManagerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceLocalUserSelectionManagerTest.kt
index 3d65713..1259b47 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceLocalUserSelectionManagerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceLocalUserSelectionManagerTest.kt
@@ -20,6 +20,7 @@
 import android.content.Intent
 import android.content.SharedPreferences
 import android.content.pm.UserInfo
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
@@ -40,7 +41,6 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.ArgumentMatchers.anyString
 import org.mockito.Mock
@@ -51,7 +51,7 @@
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class KeyguardQuickAffordanceLocalUserSelectionManagerTest : SysuiTestCase() {
 
     @Mock private lateinit var userFileManager: UserFileManager
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceRemoteUserSelectionManagerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceRemoteUserSelectionManagerTest.kt
index b21cec9..c08ef42 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceRemoteUserSelectionManagerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceRemoteUserSelectionManagerTest.kt
@@ -19,6 +19,7 @@
 
 import android.content.pm.UserInfo
 import android.os.UserHandle
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.settings.FakeUserTracker
@@ -37,13 +38,12 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.Mock
 import org.mockito.MockitoAnnotations
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class KeyguardQuickAffordanceRemoteUserSelectionManagerTest : SysuiTestCase() {
 
     @Mock private lateinit var userHandle: UserHandle
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/MuteQuickAffordanceCoreStartableTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/MuteQuickAffordanceCoreStartableTest.kt
index 34f3ed8..facc747 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/MuteQuickAffordanceCoreStartableTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/MuteQuickAffordanceCoreStartableTest.kt
@@ -21,6 +21,7 @@
 import android.media.AudioManager
 import androidx.lifecycle.MutableLiveData
 import androidx.lifecycle.Observer
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.flags.FeatureFlags
@@ -47,7 +48,6 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.Mock
 import org.mockito.Mockito.verify
 import org.mockito.Mockito.verifyZeroInteractions
@@ -55,8 +55,8 @@
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
-class MuteQuickAffordanceCoreStartableTest : SysuiTestCase()  {
+@RunWith(AndroidJUnit4::class)
+class MuteQuickAffordanceCoreStartableTest : SysuiTestCase() {
 
     @Mock
     private lateinit var featureFlags: FeatureFlags
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/QrCodeScannerKeyguardQuickAffordanceConfigTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/QrCodeScannerKeyguardQuickAffordanceConfigTest.kt
index 9d2ddff..1adf808 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/QrCodeScannerKeyguardQuickAffordanceConfigTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/QrCodeScannerKeyguardQuickAffordanceConfigTest.kt
@@ -18,6 +18,7 @@
 package com.android.systemui.keyguard.data.quickaffordance
 
 import android.content.Intent
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceConfig.OnTriggeredResult
@@ -33,13 +34,12 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.Mock
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class QrCodeScannerKeyguardQuickAffordanceConfigTest : SysuiTestCase() {
 
     @Mock private lateinit var controller: QRCodeScannerController
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/QuickAccessWalletKeyguardQuickAffordanceConfigTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/QuickAccessWalletKeyguardQuickAffordanceConfigTest.kt
index 8f56b95..752963f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/QuickAccessWalletKeyguardQuickAffordanceConfigTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/QuickAccessWalletKeyguardQuickAffordanceConfigTest.kt
@@ -20,6 +20,7 @@
 import android.graphics.drawable.Drawable
 import android.service.quickaccesswallet.GetWalletCardsResponse
 import android.service.quickaccesswallet.QuickAccessWalletClient
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
@@ -41,14 +42,13 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.Mock
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class QuickAccessWalletKeyguardQuickAffordanceConfigTest : SysuiTestCase() {
 
     @Mock private lateinit var walletController: QuickAccessWalletController
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/VideoCameraQuickAffordanceConfigTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/VideoCameraQuickAffordanceConfigTest.kt
index 805dcec..5bd86bd 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/VideoCameraQuickAffordanceConfigTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/quickaffordance/VideoCameraQuickAffordanceConfigTest.kt
@@ -17,6 +17,7 @@
 
 package com.android.systemui.keyguard.data.quickaffordance
 
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.ActivityIntentHelper
 import com.android.systemui.SysuiTestCase
@@ -32,7 +33,6 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.ArgumentMatchers.anyBoolean
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.Mock
@@ -40,7 +40,7 @@
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class VideoCameraQuickAffordanceConfigTest : SysuiTestCase() {
 
     @Mock private lateinit var activityIntentHelper: ActivityIntentHelper
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardBouncerRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardBouncerRepositoryTest.kt
index 444a2a7..ff22f1e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardBouncerRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardBouncerRepositoryTest.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard.data.repository
 
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.keyguard.ViewMediatorCallback
 import com.android.systemui.SysuiTestCase
@@ -26,13 +27,12 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.Mock
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class KeyguardBouncerRepositoryTest : SysuiTestCase() {
 
     @Mock private lateinit var systemClock: SystemClock
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardQuickAffordanceRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardQuickAffordanceRepositoryTest.kt
index 6099f01..86e8c9a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardQuickAffordanceRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardQuickAffordanceRepositoryTest.kt
@@ -19,9 +19,11 @@
 
 import android.content.pm.UserInfo
 import android.os.UserHandle
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyguard.data.quickaffordance.FakeKeyguardQuickAffordanceConfig
 import com.android.systemui.keyguard.data.quickaffordance.FakeKeyguardQuickAffordanceProviderClientFactory
 import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceConfig
@@ -39,23 +41,19 @@
 import com.android.systemui.util.mockito.whenever
 import com.android.systemui.util.settings.FakeSettings
 import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.launchIn
-import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.yield
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.ArgumentMatchers.anyString
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class KeyguardQuickAffordanceRepositoryTest : SysuiTestCase() {
 
     private lateinit var underTest: KeyguardQuickAffordanceRepository
@@ -65,12 +63,14 @@
     private lateinit var userTracker: FakeUserTracker
     private lateinit var client1: FakeCustomizationProviderClient
     private lateinit var client2: FakeCustomizationProviderClient
+    private lateinit var testScope: TestScope
 
     @Before
     fun setUp() {
         config1 = FakeKeyguardQuickAffordanceConfig(FakeCustomizationProviderClient.AFFORDANCE_1)
         config2 = FakeKeyguardQuickAffordanceConfig(FakeCustomizationProviderClient.AFFORDANCE_2)
-        val scope = CoroutineScope(IMMEDIATE)
+        val testDispatcher = StandardTestDispatcher()
+        testScope = TestScope(testDispatcher)
         userTracker = FakeUserTracker()
         val localUserSelectionManager =
             KeyguardQuickAffordanceLocalUserSelectionManager(
@@ -93,7 +93,7 @@
         client2 = FakeCustomizationProviderClient()
         val remoteUserSelectionManager =
             KeyguardQuickAffordanceRemoteUserSelectionManager(
-                scope = scope,
+                scope = testScope.backgroundScope,
                 userTracker = userTracker,
                 clientFactory =
                     FakeKeyguardQuickAffordanceProviderClientFactory(
@@ -116,14 +116,14 @@
         underTest =
             KeyguardQuickAffordanceRepository(
                 appContext = context,
-                scope = scope,
+                scope = testScope.backgroundScope,
                 localUserSelectionManager = localUserSelectionManager,
                 remoteUserSelectionManager = remoteUserSelectionManager,
                 userTracker = userTracker,
                 legacySettingSyncer =
                     KeyguardQuickAffordanceLegacySettingSyncer(
-                        scope = scope,
-                        backgroundDispatcher = IMMEDIATE,
+                        scope = testScope.backgroundScope,
+                        backgroundDispatcher = testDispatcher,
                         secureSettings = FakeSettings(),
                         selectionsManager = localUserSelectionManager,
                     ),
@@ -135,15 +135,14 @@
 
     @Test
     fun setSelections() =
-        runBlocking(IMMEDIATE) {
-            var configsBySlotId: Map<String, List<KeyguardQuickAffordanceConfig>>? = null
-            val job = underTest.selections.onEach { configsBySlotId = it }.launchIn(this)
+        testScope.runTest {
+            val configsBySlotId = collectLastValue(underTest.selections)
             val slotId1 = "slot1"
             val slotId2 = "slot2"
 
             underTest.setSelections(slotId1, listOf(config1.key))
             assertSelections(
-                configsBySlotId,
+                configsBySlotId(),
                 mapOf(
                     slotId1 to listOf(config1),
                 ),
@@ -151,7 +150,7 @@
 
             underTest.setSelections(slotId2, listOf(config2.key))
             assertSelections(
-                configsBySlotId,
+                configsBySlotId(),
                 mapOf(
                     slotId1 to listOf(config1),
                     slotId2 to listOf(config2),
@@ -161,19 +160,17 @@
             underTest.setSelections(slotId1, emptyList())
             underTest.setSelections(slotId2, listOf(config1.key))
             assertSelections(
-                configsBySlotId,
+                configsBySlotId(),
                 mapOf(
                     slotId1 to emptyList(),
                     slotId2 to listOf(config1),
                 ),
             )
-
-            job.cancel()
         }
 
     @Test
     fun getAffordancePickerRepresentations() =
-        runBlocking(IMMEDIATE) {
+        testScope.runTest {
             assertThat(underTest.getAffordancePickerRepresentations())
                 .isEqualTo(
                     listOf(
@@ -226,7 +223,7 @@
 
     @Test
     fun `selections for secondary user`() =
-        runBlocking(IMMEDIATE) {
+        testScope.runTest {
             userTracker.set(
                 userInfos =
                     listOf(
@@ -252,12 +249,10 @@
                 slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
                 affordanceId = FakeCustomizationProviderClient.AFFORDANCE_2,
             )
-            val observed = mutableListOf<Map<String, List<KeyguardQuickAffordanceConfig>>>()
-            val job = underTest.selections.onEach { observed.add(it) }.launchIn(this)
-            yield()
+            val observed = collectLastValue(underTest.selections)
 
             assertSelections(
-                observed = observed.last(),
+                observed = observed(),
                 expected =
                     mapOf(
                         KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
@@ -266,8 +261,6 @@
                             ),
                     )
             )
-
-            job.cancel()
         }
 
     private fun assertSelections(
@@ -283,7 +276,6 @@
     }
 
     companion object {
-        private val IMMEDIATE = Dispatchers.Main.immediate
         private const val SECONDARY_USER_1 = UserHandle.MIN_SECONDARY_USER_ID + 1
         private const val SECONDARY_USER_2 = UserHandle.MIN_SECONDARY_USER_ID + 2
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardRepositoryImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardRepositoryImplTest.kt
index f997d18..8bb6a85 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardRepositoryImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardRepositoryImplTest.kt
@@ -18,6 +18,7 @@
 
 import android.graphics.Point
 import android.hardware.biometrics.BiometricSourceType
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.keyguard.KeyguardUpdateMonitor
 import com.android.keyguard.KeyguardUpdateMonitorCallback
@@ -54,13 +55,12 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.Mock
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class KeyguardRepositoryImplTest : SysuiTestCase() {
 
     @Mock private lateinit var statusBarStateController: StatusBarStateController
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt
index 32cec09..ae227b4 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt
@@ -20,6 +20,7 @@
 import android.util.Log
 import android.util.Log.TerribleFailure
 import android.util.Log.TerribleFailureHandler
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.Interpolators
@@ -43,10 +44,9 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class KeyguardTransitionRepositoryTest : SysuiTestCase() {
 
     private lateinit var underTest: KeyguardTransitionRepository
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/TrustRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/TrustRepositoryTest.kt
index 4b06905..a181137 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/TrustRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/TrustRepositoryTest.kt
@@ -18,6 +18,7 @@
 
 import android.app.trust.TrustManager
 import android.content.pm.UserInfo
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.keyguard.logging.TrustRepositoryLogger
 import com.android.systemui.SysuiTestCase
@@ -33,7 +34,6 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.ArgumentCaptor
 import org.mockito.Captor
 import org.mockito.Mock
@@ -43,7 +43,7 @@
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class TrustRepositoryTest : SysuiTestCase() {
     @Mock private lateinit var trustManager: TrustManager
     @Captor private lateinit var listenerCaptor: ArgumentCaptor<TrustManager.TrustListener>
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/AlternateBouncerInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/AlternateBouncerInteractorTest.kt
index 8caf60f..7ded354 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/AlternateBouncerInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/AlternateBouncerInteractorTest.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard.domain.interactor
 
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.keyguard.KeyguardUpdateMonitor
 import com.android.keyguard.ViewMediatorCallback
@@ -36,14 +37,13 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.Mock
 import org.mockito.Mockito.mock
 import org.mockito.MockitoAnnotations
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class AlternateBouncerInteractorTest : SysuiTestCase() {
     private lateinit var underTest: AlternateBouncerInteractor
     private lateinit var bouncerRepository: KeyguardBouncerRepository
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
index d938243..7d4861b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
@@ -19,6 +19,7 @@
 
 import android.app.StatusBarManager
 import android.content.Context
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
@@ -36,12 +37,12 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.Mockito.mock
+import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class KeyguardInteractorTest : SysuiTestCase() {
     private lateinit var commandQueue: FakeCommandQueue
     private lateinit var featureFlags: FakeFeatureFlags
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardLongPressInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardLongPressInteractorTest.kt
index 9d60b16..51988ef 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardLongPressInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardLongPressInteractorTest.kt
@@ -18,6 +18,7 @@
 package com.android.systemui.keyguard.domain.interactor
 
 import android.content.Intent
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.internal.logging.UiEventLogger
 import com.android.systemui.SysuiTestCase
@@ -39,7 +40,6 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.ArgumentMatchers.anyBoolean
 import org.mockito.Mock
 import org.mockito.Mockito.never
@@ -48,7 +48,7 @@
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class KeyguardLongPressInteractorTest : SysuiTestCase() {
 
     @Mock private lateinit var activityStarter: ActivityStarter
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt
index 8cff0ae..ec70857 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt
@@ -18,6 +18,7 @@
 package com.android.systemui.keyguard.domain.interactor
 
 import android.os.UserHandle
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.internal.widget.LockPatternUtils
 import com.android.systemui.R
@@ -60,7 +61,6 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.ArgumentMatchers.anyString
 import org.mockito.Mock
@@ -68,7 +68,7 @@
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class KeyguardQuickAffordanceInteractorTest : SysuiTestCase() {
 
     @Mock private lateinit var lockPatternUtils: LockPatternUtils
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt
index 6333b24..3d13d80 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt
@@ -17,6 +17,7 @@
 
 package com.android.systemui.keyguard.domain.interactor
 
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
@@ -29,17 +30,17 @@
 import com.android.systemui.keyguard.shared.model.TransitionState.STARTED
 import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
+@kotlinx.coroutines.ExperimentalCoroutinesApi
 class KeyguardTransitionInteractorTest : SysuiTestCase() {
 
     private lateinit var underTest: KeyguardTransitionInteractor
@@ -53,7 +54,7 @@
 
     @Test
     fun `transition collectors receives only appropriate events`() =
-        runBlocking(IMMEDIATE) {
+        runTest(UnconfinedTestDispatcher()) {
             var lockscreenToAodSteps = mutableListOf<TransitionStep>()
             val job1 =
                 underTest.lockscreenToAodTransition
@@ -87,10 +88,9 @@
 
     @Test
     fun dozeAmountTransitionTest() =
-        runBlocking(IMMEDIATE) {
+        runTest(UnconfinedTestDispatcher()) {
             var dozeAmountSteps = mutableListOf<TransitionStep>()
-            val job =
-                underTest.dozeAmountTransition.onEach { dozeAmountSteps.add(it) }.launchIn(this)
+            val job = underTest.dozeAmountTransition.onEach { dozeAmountSteps.add(it) }.launchIn(this)
 
             val steps = mutableListOf<TransitionStep>()
 
@@ -119,10 +119,9 @@
 
     @Test
     fun keyguardStateTests() =
-        runBlocking(IMMEDIATE) {
+        runTest(UnconfinedTestDispatcher()) {
             var finishedSteps = mutableListOf<KeyguardState>()
-            val job =
-                underTest.finishedKeyguardState.onEach { finishedSteps.add(it) }.launchIn(this)
+            val job = underTest.finishedKeyguardState.onEach { finishedSteps.add(it) }.launchIn(this)
 
             val steps = mutableListOf<TransitionStep>()
 
@@ -143,12 +142,10 @@
 
     @Test
     fun finishedKeyguardTransitionStepTests() =
-        runBlocking(IMMEDIATE) {
+        runTest(UnconfinedTestDispatcher()) {
             var finishedSteps = mutableListOf<TransitionStep>()
             val job =
-                underTest.finishedKeyguardTransitionStep
-                    .onEach { finishedSteps.add(it) }
-                    .launchIn(this)
+                underTest.finishedKeyguardTransitionStep.onEach { finishedSteps.add(it) }.launchIn(this)
 
             val steps = mutableListOf<TransitionStep>()
 
@@ -169,12 +166,10 @@
 
     @Test
     fun startedKeyguardTransitionStepTests() =
-        runBlocking(IMMEDIATE) {
+        runTest(UnconfinedTestDispatcher()) {
             var startedSteps = mutableListOf<TransitionStep>()
             val job =
-                underTest.startedKeyguardTransitionStep
-                    .onEach { startedSteps.add(it) }
-                    .launchIn(this)
+                underTest.startedKeyguardTransitionStep.onEach { startedSteps.add(it) }.launchIn(this)
 
             val steps = mutableListOf<TransitionStep>()
 
@@ -192,8 +187,4 @@
 
             job.cancel()
         }
-
-    companion object {
-        private val IMMEDIATE = Dispatchers.Main.immediate
-    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/LightRevealScrimInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/LightRevealScrimInteractorTest.kt
index 3166214..6236616 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/LightRevealScrimInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/LightRevealScrimInteractorTest.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard.domain.interactor
 
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
@@ -33,11 +34,10 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.MockitoAnnotations
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class LightRevealScrimInteractorTest : SysuiTestCase() {
     private val fakeKeyguardTransitionRepository = FakeKeyguardTransitionRepository()
     private val fakeLightRevealScrimRepository = FakeLightRevealScrimRepository()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/PrimaryBouncerCallbackInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/PrimaryBouncerCallbackInteractorTest.kt
index fbfeca9..f86ac79 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/PrimaryBouncerCallbackInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/PrimaryBouncerCallbackInteractorTest.kt
@@ -17,18 +17,18 @@
 package com.android.systemui.keyguard.domain.interactor
 
 import android.view.View
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.Mock
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class PrimaryBouncerCallbackInteractorTest : SysuiTestCase() {
     private val mPrimaryBouncerCallbackInteractor = PrimaryBouncerCallbackInteractor()
     @Mock
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/PrimaryBouncerInteractorWithCoroutinesTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/PrimaryBouncerInteractorWithCoroutinesTest.kt
index ea7bc91..75b74b0 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/PrimaryBouncerInteractorWithCoroutinesTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/PrimaryBouncerInteractorWithCoroutinesTest.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.keyguard.domain.interactor
 
 import android.os.Looper
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.keyguard.KeyguardSecurityModel
 import com.android.keyguard.KeyguardUpdateMonitor
@@ -35,12 +36,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 import org.mockito.Mock
 import org.mockito.MockitoAnnotations
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class PrimaryBouncerInteractorWithCoroutinesTest : SysuiTestCase() {
     private lateinit var repository: FakeKeyguardBouncerRepository
     @Mock private lateinit var bouncerView: BouncerView
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToLockscreenTransitionViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToLockscreenTransitionViewModelTest.kt
index 06e397d..706154e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToLockscreenTransitionViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToLockscreenTransitionViewModelTest.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
@@ -33,10 +34,9 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class DreamingToLockscreenTransitionViewModelTest : SysuiTestCase() {
     private lateinit var underTest: DreamingToLockscreenTransitionViewModel
     private lateinit var repository: FakeKeyguardTransitionRepository
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/GoneToDreamingTransitionViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/GoneToDreamingTransitionViewModelTest.kt
index 14c3b50..b15ce10 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/GoneToDreamingTransitionViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/GoneToDreamingTransitionViewModelTest.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
@@ -32,10 +33,9 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class GoneToDreamingTransitionViewModelTest : SysuiTestCase() {
     private lateinit var underTest: GoneToDreamingTransitionViewModel
     private lateinit var repository: FakeKeyguardTransitionRepository
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBouncerViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBouncerViewModelTest.kt
index 3727134..586af62 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBouncerViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBouncerViewModelTest.kt
@@ -16,28 +16,29 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.keyguard.data.BouncerView
 import com.android.systemui.keyguard.domain.interactor.PrimaryBouncerInteractor
 import com.android.systemui.keyguard.shared.model.BouncerShowMessageModel
 import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.runBlocking
+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.junit.runners.JUnit4
 import org.mockito.Mock
 import org.mockito.Mockito
 import org.mockito.MockitoAnnotations
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
+@kotlinx.coroutines.ExperimentalCoroutinesApi
 class KeyguardBouncerViewModelTest : SysuiTestCase() {
     lateinit var underTest: KeyguardBouncerViewModel
     @Mock lateinit var bouncerView: BouncerView
@@ -51,7 +52,7 @@
 
     @Test
     fun setMessage() =
-        runBlocking(Dispatchers.Main.immediate) {
+        runTest {
             val flow = MutableStateFlow<BouncerShowMessageModel?>(null)
             var message: BouncerShowMessageModel? = null
             Mockito.`when`(bouncerInteractor.showMessage)
@@ -62,6 +63,8 @@
             flow.value = BouncerShowMessageModel(message = "abc", colorStateList = null)
 
             val job = underTest.bouncerShowMessage.onEach { message = it }.launchIn(this)
+            // Run the tasks that are pending at this point of virtual time.
+            runCurrent()
             assertThat(message?.message).isEqualTo("abc")
             job.cancel()
         }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToDreamingTransitionViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToDreamingTransitionViewModelTest.kt
index ed31dc3..d94c108 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToDreamingTransitionViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToDreamingTransitionViewModelTest.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
@@ -32,10 +33,9 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class LockscreenToDreamingTransitionViewModelTest : SysuiTestCase() {
     private lateinit var underTest: LockscreenToDreamingTransitionViewModel
     private lateinit var repository: FakeKeyguardTransitionRepository
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToOccludedTransitionViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToOccludedTransitionViewModelTest.kt
index 458b315..12ec24d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToOccludedTransitionViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToOccludedTransitionViewModelTest.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
@@ -32,10 +33,9 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class LockscreenToOccludedTransitionViewModelTest : SysuiTestCase() {
     private lateinit var underTest: LockscreenToOccludedTransitionViewModel
     private lateinit var repository: FakeKeyguardTransitionRepository
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/OccludedToLockscreenTransitionViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/OccludedToLockscreenTransitionViewModelTest.kt
index a36214e..0c4e845 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/OccludedToLockscreenTransitionViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/OccludedToLockscreenTransitionViewModelTest.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
@@ -32,10 +33,9 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
 
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(AndroidJUnit4::class)
 class OccludedToLockscreenTransitionViewModelTest : SysuiTestCase() {
     private lateinit var underTest: OccludedToLockscreenTransitionViewModel
     private lateinit var repository: FakeKeyguardTransitionRepository