diff --git a/tests/Android.bp b/tests/Android.bp
index 3822ff8..5ec2263 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -23,7 +23,8 @@
     srcs: [
         "src/**/*.java",
         "src/**/*.kt",
-        ":launcher3-robo-src",
+        "multivalentTests/src/**/*.java",
+        "multivalentTests/src/**/*.kt",
     ],
     exclude_srcs: [
         ":launcher-non-quickstep-tests-src",
@@ -37,6 +38,8 @@
     srcs: [
         "multivalentTests/src/**/*.java",
         "multivalentTests/src/**/*.kt",
+        "src_deviceless/**/*.java",
+        "src_deviceless/**/*.kt",
     ],
 }
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/icons/FastBitmapDrawableTest.java b/tests/multivalentTests/src/com/android/launcher3/icons/FastBitmapDrawableTest.java
index 7e9b68d..58dce0b 100644
--- a/tests/multivalentTests/src/com/android/launcher3/icons/FastBitmapDrawableTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/icons/FastBitmapDrawableTest.java
@@ -37,15 +37,12 @@
 import android.view.animation.PathInterpolator;
 
 import androidx.test.annotation.UiThreadTest;
-import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
-import com.android.launcher3.util.rule.RobolectricUiThreadRule;
+import com.android.launcher3.util.LauncherMultivalentJUnit;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
-import org.junit.rules.TestRule;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
@@ -55,14 +52,11 @@
  * Tests for FastBitmapDrawable.
  */
 @SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(LauncherMultivalentJUnit.class)
 @UiThreadTest
 public class FastBitmapDrawableTest {
     private static final float EPSILON = 0.00001f;
 
-    @Rule
-    public final TestRule roboUiThreadRule = new RobolectricUiThreadRule();
-
     @Spy
     FastBitmapDrawable mFastBitmapDrawable =
             spy(new FastBitmapDrawable(Bitmap.createBitmap(100, 100, Bitmap.Config.ARGB_8888)));
diff --git a/tests/multivalentTests/src/com/android/launcher3/logging/StartupLatencyLoggerTest.kt b/tests/multivalentTests/src/com/android/launcher3/logging/StartupLatencyLoggerTest.kt
index 12f6c8c..713d4d5 100644
--- a/tests/multivalentTests/src/com/android/launcher3/logging/StartupLatencyLoggerTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/logging/StartupLatencyLoggerTest.kt
@@ -2,22 +2,18 @@
 
 import androidx.core.util.isEmpty
 import androidx.test.annotation.UiThreadTest
-import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
-import com.android.launcher3.util.rule.RobolectricUiThreadRule
+import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.google.common.truth.Truth.assertThat
 import org.junit.Before
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 
 /** Unit test for [ColdRebootStartupLatencyLogger]. */
 @SmallTest
-@RunWith(AndroidJUnit4::class)
+@RunWith(LauncherMultivalentJUnit::class)
 class StartupLatencyLoggerTest {
 
-    @get:Rule val roboUiThreadRule = RobolectricUiThreadRule()
-
     private val underTest = ColdRebootStartupLatencyLogger()
 
     @Before
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/EmulatedDeviceAndroidJUnit.kt b/tests/multivalentTests/src/com/android/launcher3/util/LauncherMultivalentJUnit.kt
similarity index 68%
rename from tests/multivalentTests/src/com/android/launcher3/util/EmulatedDeviceAndroidJUnit.kt
rename to tests/multivalentTests/src/com/android/launcher3/util/LauncherMultivalentJUnit.kt
index 694f257..e8560af 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/EmulatedDeviceAndroidJUnit.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/util/LauncherMultivalentJUnit.kt
@@ -24,23 +24,24 @@
 import org.junit.runners.Suite
 
 /**
- * A custom runner which emulates multiple devices when running in robolectric framework. Runs
- * normally when running on device
+ * A custom runner for multivalent tests with launcher specific features
+ * 1) Adds support for @UiThread annotations in deviceless tests
+ * 2) Allows emulating multiple devices when running in deviceless mode
  */
-class EmulatedDeviceAndroidJUnit(klass: Class<*>?) : Suite(klass, ImmutableList.of()) {
+class LauncherMultivalentJUnit(klass: Class<*>?) : Suite(klass, ImmutableList.of()) {
 
     val runners: List<Runner> =
-        testClass.getAnnotation(Devices::class.java)?.value?.let { devices ->
-            if (devices.isEmpty() || !isRunningInRobolectric) {
+        (testClass.getAnnotation(EmulatedDevices::class.java)?.value ?: emptyArray()).let { devices
+            ->
+            if (!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)
-                        }
+                        if (devices.isEmpty()) listOf(ctor.newInstance(testClass.javaClass, null))
+                        else devices.map { ctor.newInstance(testClass.javaClass, it) }
                     }
             } catch (e: Exception) {
                 null
@@ -50,11 +51,13 @@
 
     override fun getChildren() = runners
 
-    @Retention(RUNTIME) @Target(CLASS) annotation class Devices(val value: Array<String>)
+    /**
+     * Annotation to be added to a test so run it on a list of emulated devices for deviceless test
+     */
+    @Retention(RUNTIME) @Target(CLASS) annotation class EmulatedDevices(val value: Array<String>)
 
     companion object {
-        private const val ROBOLECTRIC_RUNNER =
-            "com.android.launcher3.util.RobolectricEmulatedDeviceRunner"
+        private const val ROBOLECTRIC_RUNNER = "com.android.launcher3.util.RobolectricDeviceRunner"
 
         val isRunningInRobolectric: Boolean
             get() =
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/rule/RobolectricUiThreadRule.kt b/tests/multivalentTests/src/com/android/launcher3/util/rule/RobolectricUiThreadRule.kt
deleted file mode 100644
index b65c443..0000000
--- a/tests/multivalentTests/src/com/android/launcher3/util/rule/RobolectricUiThreadRule.kt
+++ /dev/null
@@ -1,75 +0,0 @@
-/*
- * 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.rule
-
-import androidx.test.annotation.UiThreadTest
-import androidx.test.platform.app.InstrumentationRegistry
-import com.android.launcher3.util.EmulatedDeviceAndroidJUnit.Companion.isRunningInRobolectric
-import java.util.concurrent.atomic.AtomicReference
-import org.junit.rules.TestRule
-import org.junit.runner.Description
-import org.junit.runners.model.Statement
-
-/**
- * A test rule to add support for @UiThreadTest annotations when running in robolectric until is it
- * natively supported by the robolectric runner:
- * https://github.com/robolectric/robolectric/issues/9026
- */
-class RobolectricUiThreadRule : TestRule {
-
-    override fun apply(base: Statement, description: Description): Statement =
-        if (!shouldRunOnUiThread(description)) base else UiThreadStatement(base)
-
-    private fun shouldRunOnUiThread(description: Description): Boolean {
-        if (!isRunningInRobolectric) {
-            // If not running in robolectric, let the default runner handle this
-            return false
-        }
-        var clazz = description.testClass
-        try {
-            if (
-                clazz
-                    .getDeclaredMethod(description.methodName)
-                    .getAnnotation(UiThreadTest::class.java) != null
-            ) {
-                return true
-            }
-        } catch (_: Exception) {
-            // Ignore
-        }
-
-        while (!clazz.isAnnotationPresent(UiThreadTest::class.java)) {
-            clazz = clazz.superclass ?: return false
-        }
-        return true
-    }
-
-    private class UiThreadStatement(val base: Statement) : Statement() {
-
-        override fun evaluate() {
-            val exceptionRef = AtomicReference<Throwable>()
-            InstrumentationRegistry.getInstrumentation().runOnMainSync {
-                try {
-                    base.evaluate()
-                } catch (throwable: Throwable) {
-                    exceptionRef.set(throwable)
-                }
-            }
-            exceptionRef.get()?.let { throw it }
-        }
-    }
-}
diff --git a/tests/src_deviceless/com/android/launcher3/util/RobolectricDeviceRunner.kt b/tests/src_deviceless/com/android/launcher3/util/RobolectricDeviceRunner.kt
new file mode 100644
index 0000000..dc6d716
--- /dev/null
+++ b/tests/src_deviceless/com/android/launcher3/util/RobolectricDeviceRunner.kt
@@ -0,0 +1,107 @@
+/*
+ * 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.annotation.UiThreadTest
+import androidx.test.platform.app.InstrumentationRegistry
+import java.lang.reflect.Method
+import java.util.concurrent.atomic.AtomicReference
+import org.junit.runners.model.FrameworkMethod
+import org.junit.runners.model.Statement
+import org.robolectric.RobolectricTestRunner
+import org.robolectric.internal.bytecode.Sandbox
+import org.robolectric.util.ReflectionHelpers
+import org.robolectric.util.ReflectionHelpers.ClassParameter
+
+/** Runner which emulates the provided display before running the actual test */
+class RobolectricDeviceRunner(testClass: Class<*>?, private val deviceName: String?) :
+    RobolectricTestRunner(testClass) {
+
+    private val nameSuffix = deviceName?.let { "-$it" } ?: ""
+
+    override fun getName() = super.getName() + nameSuffix
+
+    override fun testName(method: FrameworkMethod) = super.testName(method) + nameSuffix
+
+    @Throws(Throwable::class)
+    override fun beforeTest(sandbox: Sandbox, method: FrameworkMethod, bootstrappedMethod: Method) {
+        super.beforeTest(sandbox, method, bootstrappedMethod)
+
+        deviceName ?: return
+
+        val emulator =
+            try {
+                ReflectionHelpers.loadClass(
+                    bootstrappedMethod.declaringClass.classLoader,
+                    DEVICE_EMULATOR
+                )
+            } catch (e: Exception) {
+                // Ignore, if the device emulator is not present
+                return
+            }
+        ReflectionHelpers.callStaticMethod<Any>(
+            emulator,
+            "updateDevice",
+            ClassParameter.from(String::class.java, deviceName)
+        )
+    }
+
+    override fun getHelperTestRunner(clazz: Class<*>) = MyHelperTestRunner(clazz)
+
+    class MyHelperTestRunner(clazz: Class<*>) : HelperTestRunner(clazz) {
+
+        override fun methodBlock(method: FrameworkMethod): Statement =
+            // this needs to be run in the test classLoader
+            ReflectionHelpers.callStaticMethod(
+                method.declaringClass.classLoader,
+                RobolectricDeviceRunner::class.qualifiedName,
+                "wrapUiThreadMethod",
+                ClassParameter.from(FrameworkMethod::class.java, method),
+                ClassParameter.from(Statement::class.java, super.methodBlock(method))
+            )
+    }
+
+    private class UiThreadStatement(val base: Statement) : Statement() {
+
+        override fun evaluate() {
+            val exceptionRef = AtomicReference<Throwable>()
+            InstrumentationRegistry.getInstrumentation().runOnMainSync {
+                try {
+                    base.evaluate()
+                } catch (throwable: Throwable) {
+                    exceptionRef.set(throwable)
+                }
+            }
+            exceptionRef.get()?.let { throw it }
+        }
+    }
+
+    companion object {
+
+        private const val DEVICE_EMULATOR = "com.android.launcher3.util.RoboDeviceEmulator"
+
+        @JvmStatic
+        fun wrapUiThreadMethod(method: FrameworkMethod, base: Statement): Statement =
+            if (
+                method.method.isAnnotationPresent(UiThreadTest::class.java) ||
+                    method.declaringClass.isAnnotationPresent(UiThreadTest::class.java)
+            ) {
+                UiThreadStatement(base)
+            } else {
+                base
+            }
+    }
+}