Adding a special runner to run simulate multiple devices when running robo tests

Bug: 338128923
Test: Verrified by modifying a test
Flag: None
Change-Id: I133842d68e437b971da39e2081c8859515463cd3
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/EmulatedDeviceAndroidJUnit.kt b/tests/multivalentTests/src/com/android/launcher3/util/EmulatedDeviceAndroidJUnit.kt
new file mode 100644
index 0000000..694f257
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/util/EmulatedDeviceAndroidJUnit.kt
@@ -0,0 +1,76 @@
+/*
+ * 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.launcher3.util
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.collect.ImmutableList
+import java.util.Locale
+import kotlin.annotation.AnnotationRetention.RUNTIME
+import kotlin.annotation.AnnotationTarget.CLASS
+import org.junit.runner.Runner
+import org.junit.runners.Suite
+
+/**
+ * A custom runner which emulates multiple devices when running in robolectric framework. Runs
+ * normally when running on device
+ */
+class EmulatedDeviceAndroidJUnit(klass: Class<*>?) : Suite(klass, ImmutableList.of()) {
+
+    val runners: List<Runner> =
+        testClass.getAnnotation(Devices::class.java)?.value?.let { devices ->
+            if (devices.isEmpty() || !isRunningInRobolectric) {
+                return@let null
+            }
+            try {
+                (testClass.javaClass.classLoader.loadClass(ROBOLECTRIC_RUNNER) as Class<Runner>)
+                    .getConstructor(Class::class.java, String::class.java)
+                    .let { ctor ->
+                        devices.map { deviceName ->
+                            ctor.newInstance(testClass.javaClass, deviceName)
+                        }
+                    }
+            } catch (e: Exception) {
+                null
+            }
+        }
+            ?: listOf(AndroidJUnit4(testClass.javaClass))
+
+    override fun getChildren() = runners
+
+    @Retention(RUNTIME) @Target(CLASS) annotation class Devices(val value: Array<String>)
+
+    companion object {
+        private const val ROBOLECTRIC_RUNNER =
+            "com.android.launcher3.util.RobolectricEmulatedDeviceRunner"
+
+        val isRunningInRobolectric: Boolean
+            get() =
+                if (
+                    System.getProperty("java.runtime.name")
+                        .lowercase(Locale.getDefault())
+                        .contains("android")
+                ) {
+                    false
+                } else {
+                    try {
+                        // Check if robolectric runner exists
+                        Class.forName("org.robolectric.RobolectricTestRunner") != null
+                    } catch (e: ClassNotFoundException) {
+                        false
+                    }
+                }
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/rule/RobolectricUiThreadRule.kt b/tests/multivalentTests/src/com/android/launcher3/util/rule/RobolectricUiThreadRule.kt
index 18cd1e4..b65c443 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/rule/RobolectricUiThreadRule.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/util/rule/RobolectricUiThreadRule.kt
@@ -18,7 +18,7 @@
 
 import androidx.test.annotation.UiThreadTest
 import androidx.test.platform.app.InstrumentationRegistry
-import java.util.Locale
+import com.android.launcher3.util.EmulatedDeviceAndroidJUnit.Companion.isRunningInRobolectric
 import java.util.concurrent.atomic.AtomicReference
 import org.junit.rules.TestRule
 import org.junit.runner.Description
@@ -35,7 +35,7 @@
         if (!shouldRunOnUiThread(description)) base else UiThreadStatement(base)
 
     private fun shouldRunOnUiThread(description: Description): Boolean {
-        if (!isRunningInRobolectric()) {
+        if (!isRunningInRobolectric) {
             // If not running in robolectric, let the default runner handle this
             return false
         }
@@ -58,21 +58,6 @@
         return true
     }
 
-    private fun isRunningInRobolectric(): Boolean {
-        if (
-            System.getProperty("java.runtime.name")
-                .lowercase(Locale.getDefault())
-                .contains("android")
-        )
-            return false
-        return try {
-            // Check if robolectric runner exists
-            Class.forName("org.robolectric.RobolectricTestRunner") != null
-        } catch (e: ClassNotFoundException) {
-            false
-        }
-    }
-
     private class UiThreadStatement(val base: Statement) : Statement() {
 
         override fun evaluate() {