diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarUnitTestRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarUnitTestRule.kt
index 1c900b8..bbf738e 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarUnitTestRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarUnitTestRule.kt
@@ -31,14 +31,15 @@
 import com.android.quickstep.TouchInteractionService
 import com.android.quickstep.TouchInteractionService.TISBinder
 import org.junit.Assume.assumeTrue
-import org.junit.rules.MethodRule
-import org.junit.runners.model.FrameworkMethod
+import org.junit.rules.TestRule
+import org.junit.runner.Description
 import org.junit.runners.model.Statement
 
 /**
  * Manages the Taskbar lifecycle for unit tests.
  *
- * Tests need to provide their target [context] through the constructor.
+ * Tests should pass in themselves as [testInstance]. They also need to provide their target
+ * [context] through the constructor.
  *
  * See [InjectController] for grabbing controller(s) under test with minimal boilerplate.
  *
@@ -61,12 +62,11 @@
  * }
  * ```
  */
-class TaskbarUnitTestRule(private val context: Context) : MethodRule {
+class TaskbarUnitTestRule(private val testInstance: Any, private val context: Context) : TestRule {
     private val instrumentation = InstrumentationRegistry.getInstrumentation()
     private val serviceTestRule = ServiceTestRule()
 
     private lateinit var taskbarManager: TaskbarManager
-    private lateinit var target: Any
 
     val activityContext: TaskbarActivityContext
         get() {
@@ -74,10 +74,9 @@
                 ?: throw RuntimeException("Failed to obtain TaskbarActivityContext.")
         }
 
-    override fun apply(base: Statement, method: FrameworkMethod, target: Any): Statement {
+    override fun apply(base: Statement, description: Description): Statement {
         return object : Statement() {
             override fun evaluate() {
-                this@TaskbarUnitTestRule.target = target
 
                 instrumentation.runOnMainSync {
                     assumeTrue(
@@ -141,11 +140,11 @@
     private fun injectControllers() {
         val controllers = activityContext.controllers
         val controllerFieldsByType = controllers.javaClass.fields.associateBy { it.type }
-        target.javaClass.fields
+        testInstance.javaClass.fields
             .filter { it.isAnnotationPresent(InjectController::class.java) }
             .forEach {
                 it.set(
-                    target,
+                    testInstance,
                     controllerFieldsByType[it.type]?.get(controllers)
                         ?: throw NoSuchElementException("Failed to find controller for ${it.type}"),
                 )
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
index 9a514bf..bfad697 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
@@ -42,7 +42,8 @@
 @EmulatedDevices(["pixelFoldable2023", "pixelTablet2023"])
 class TaskbarAllAppsControllerTest {
 
-    @get:Rule val taskbarUnitTestRule = TaskbarUnitTestRule(getInstrumentation().targetContext)
+    @get:Rule
+    val taskbarUnitTestRule = TaskbarUnitTestRule(this, getInstrumentation().targetContext)
     @get:Rule val animatorTestRule = AnimatorTestRule(this)
 
     @InjectController lateinit var allAppsController: TaskbarAllAppsController
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
index 918ec7d..72bdc16 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
@@ -41,7 +41,8 @@
 @EmulatedDevices(["pixelFoldable2023"])
 class TaskbarOverlayControllerTest {
 
-    @get:Rule val taskbarUnitTestRule = TaskbarUnitTestRule(getInstrumentation().targetContext)
+    @get:Rule
+    val taskbarUnitTestRule = TaskbarUnitTestRule(this, getInstrumentation().targetContext)
     @InjectController lateinit var overlayController: TaskbarOverlayController
 
     private val taskbarContext: TaskbarActivityContext