diff --git a/quickstep/testing/com/android/launcher3/taskbar/bubbles/testing/FakeBubbleViewFactory.kt b/quickstep/testing/com/android/launcher3/taskbar/bubbles/testing/FakeBubbleViewFactory.kt
index 37a07c3..473d8ef 100644
--- a/quickstep/testing/com/android/launcher3/taskbar/bubbles/testing/FakeBubbleViewFactory.kt
+++ b/quickstep/testing/com/android/launcher3/taskbar/bubbles/testing/FakeBubbleViewFactory.kt
@@ -54,7 +54,19 @@
         val flags =
             if (suppressNotification) Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION else 0
         val bubbleInfo =
-            BubbleInfo(key, flags, null, null, 0, context.packageName, null, null, false, true)
+            BubbleInfo(
+                key,
+                flags,
+                null,
+                null,
+                0,
+                context.packageName,
+                null,
+                null,
+                false,
+                true,
+                null,
+            )
         val bubbleView = inflater.inflate(R.layout.bubblebar_item_view, parent, false) as BubbleView
         val dotPath =
             PathParser.createPathFromPathData(
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt
index e736446..d7ce4ed 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt
@@ -35,7 +35,8 @@
 import com.android.launcher3.taskbar.TaskbarStashController.TRANSIENT_TASKBAR_STASH_ALPHA_DURATION
 import com.android.launcher3.taskbar.TaskbarStashController.TRANSIENT_TASKBAR_STASH_DURATION
 import com.android.launcher3.taskbar.TaskbarViewController.ALPHA_INDEX_STASH
-import com.android.launcher3.taskbar.bubbles.BubbleControllers
+import com.android.launcher3.taskbar.bubbles.BubbleBarViewController
+import com.android.launcher3.taskbar.bubbles.stashing.BubbleStashController
 import com.android.launcher3.taskbar.rules.TaskbarModeRule
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.PINNED
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.THREE_BUTTONS
@@ -52,7 +53,6 @@
 import com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING
 import com.android.wm.shell.Flags.FLAG_ENABLE_BUBBLE_BAR
 import com.google.common.truth.Truth.assertThat
-import java.util.Optional
 import org.junit.After
 import org.junit.Before
 import org.junit.Rule
@@ -60,6 +60,7 @@
 import org.junit.runner.RunWith
 
 @RunWith(LauncherMultivalentJUnit::class)
+@EnableFlags(FLAG_ENABLE_BUBBLE_BAR)
 @EmulatedDevices(["pixelTablet2023"])
 class TaskbarStashControllerTest {
     private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
@@ -74,7 +75,8 @@
     @InjectController lateinit var stashedHandleViewController: StashedHandleViewController
     @InjectController lateinit var dragLayerController: TaskbarDragLayerController
     @InjectController lateinit var autohideSuspendController: TaskbarAutohideSuspendController
-    @InjectController lateinit var bubbleControllers: Optional<BubbleControllers>
+    @InjectController lateinit var bubbleBarViewController: BubbleBarViewController
+    @InjectController lateinit var bubbleStashController: BubbleStashController
 
     private val activityContext by taskbarUnitTestRule::activityContext
 
@@ -420,60 +422,55 @@
     }
 
     @Test
-    @EnableFlags(FLAG_ENABLE_BUBBLE_BAR)
     @TaskbarMode(TRANSIENT)
     fun testUpdateAndAnimateTransientTaskbar_unstashTaskbarWithBubbles_bubbleBarUnstashes() {
         getInstrumentation().runOnMainSync {
-            bubbleControllers.get().bubbleBarViewController.setHiddenForBubbles(false)
-            bubbleControllers.get().bubbleStashController.stashBubbleBarImmediate()
+            bubbleBarViewController.setHiddenForBubbles(false)
+            bubbleStashController.stashBubbleBarImmediate()
             stashController.updateAndAnimateTransientTaskbar(false, true)
         }
-        assertThat(bubbleControllers.get().bubbleStashController.isStashed).isFalse()
+        assertThat(bubbleStashController.isStashed).isFalse()
     }
 
     @Test
-    @EnableFlags(FLAG_ENABLE_BUBBLE_BAR)
     @TaskbarMode(TRANSIENT)
     fun testUpdateAndAnimateTransientTaskbar_unstashTaskbarWithoutBubbles_bubbleBarStashed() {
         getInstrumentation().runOnMainSync {
-            bubbleControllers.get().bubbleBarViewController.setHiddenForBubbles(false)
-            bubbleControllers.get().bubbleStashController.stashBubbleBarImmediate()
+            bubbleBarViewController.setHiddenForBubbles(false)
+            bubbleStashController.stashBubbleBarImmediate()
             stashController.updateAndAnimateTransientTaskbar(false, false)
         }
-        assertThat(bubbleControllers.get().bubbleStashController.isStashed).isTrue()
+        assertThat(bubbleStashController.isStashed).isTrue()
     }
 
     @Test
-    @EnableFlags(FLAG_ENABLE_BUBBLE_BAR)
     @TaskbarMode(TRANSIENT)
     fun testUpdateAndAnimateTransientTaskbar_stashTaskbarWithBubbles_bubbleBarStashes() {
         getInstrumentation().runOnMainSync {
-            bubbleControllers.get().bubbleBarViewController.setHiddenForBubbles(false)
-            bubbleControllers.get().bubbleStashController.showBubbleBarImmediate()
+            bubbleBarViewController.setHiddenForBubbles(false)
+            bubbleStashController.showBubbleBarImmediate()
             stashController.updateAndAnimateTransientTaskbar(true, true)
         }
-        assertThat(bubbleControllers.get().bubbleStashController.isStashed).isTrue()
+        assertThat(bubbleStashController.isStashed).isTrue()
     }
 
     @Test
-    @EnableFlags(FLAG_ENABLE_BUBBLE_BAR)
     @TaskbarMode(TRANSIENT)
     fun testUpdateAndAnimateTransientTaskbar_stashTaskbarWithoutBubbles_bubbleBarUnstashed() {
         getInstrumentation().runOnMainSync {
-            bubbleControllers.get().bubbleBarViewController.setHiddenForBubbles(false)
-            bubbleControllers.get().bubbleStashController.showBubbleBarImmediate()
+            bubbleBarViewController.setHiddenForBubbles(false)
+            bubbleStashController.showBubbleBarImmediate()
             stashController.updateAndAnimateTransientTaskbar(true, false)
         }
-        assertThat(bubbleControllers.get().bubbleStashController.isStashed).isFalse()
+        assertThat(bubbleStashController.isStashed).isFalse()
     }
 
     @Test
-    @EnableFlags(FLAG_ENABLE_BUBBLE_BAR)
     @TaskbarMode(TRANSIENT)
     fun testUpdateAndAnimateTransientTaskbar_bubbleBarExpandedBeforeTimeout_expandedAfterwards() {
         getInstrumentation().runOnMainSync {
-            bubbleControllers.get().bubbleBarViewController.setHiddenForBubbles(false)
-            bubbleControllers.get().bubbleBarViewController.isExpanded = true
+            bubbleBarViewController.setHiddenForBubbles(false)
+            bubbleBarViewController.isExpanded = true
             stashController.updateAndAnimateTransientTaskbar(false)
             animatorTestRule.advanceTimeBy(stashController.stashDuration)
         }
@@ -483,7 +480,7 @@
             stashController.timeoutAlarm.finishAlarm()
             animatorTestRule.advanceTimeBy(stashController.stashDuration)
         }
-        assertThat(bubbleControllers.get().bubbleBarViewController.isExpanded).isTrue()
+        assertThat(bubbleBarViewController.isExpanded).isTrue()
     }
 
     @Test
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewTest.kt
index 94f9cf5..2caff01 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewTest.kt
@@ -65,7 +65,19 @@
             overflowView.setOverflow(BubbleBarOverflow(overflowView), bitmap)
 
             val bubbleInfo =
-                BubbleInfo("key", 0, null, null, 0, context.packageName, null, null, false, true)
+                BubbleInfo(
+                    "key",
+                    0,
+                    null,
+                    null,
+                    0,
+                    context.packageName,
+                    null,
+                    null,
+                    false,
+                    true,
+                    null,
+                )
             bubbleView = inflater.inflate(R.layout.bubblebar_item_view, null, false) as BubbleView
             bubble =
                 BubbleBarBubble(bubbleInfo, bubbleView, bitmap, bitmap, Color.WHITE, Path(), "")
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt
index 84e872d..3913165 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt
@@ -86,7 +86,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpandedNoOp,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -135,7 +135,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpandedNoOp,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -183,7 +183,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpandedNoOp,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -228,7 +228,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpandedNoOp,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -274,7 +274,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpandedNoOp,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -311,7 +311,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpanded,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -355,7 +355,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpanded,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -405,7 +405,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpanded,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -454,7 +454,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpandedNoOp,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -504,7 +504,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpanded,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -538,7 +538,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpandedNoOp,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -577,7 +577,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpanded,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -625,7 +625,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpanded,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -666,7 +666,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpandedNoOp,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -713,7 +713,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpanded,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -760,7 +760,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpanded,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -818,7 +818,7 @@
                 bubbleBarView,
                 bubbleStashController,
                 onExpanded,
-                animatorScheduler
+                animatorScheduler,
             )
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
@@ -870,7 +870,19 @@
             bubbleBarView.addView(overflowView)
 
             val bubbleInfo =
-                BubbleInfo("key", 0, null, null, 0, context.packageName, null, null, false, true)
+                BubbleInfo(
+                    "key",
+                    0,
+                    null,
+                    null,
+                    0,
+                    context.packageName,
+                    null,
+                    null,
+                    false,
+                    true,
+                    null,
+                )
             bubbleView =
                 inflater.inflate(R.layout.bubblebar_item_view, bubbleBarView, false) as BubbleView
             bubble =
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
index cb5e464..b0d706f 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
@@ -28,9 +28,11 @@
 import com.android.launcher3.LauncherAppState
 import com.android.launcher3.statehandlers.DesktopVisibilityController
 import com.android.launcher3.taskbar.TaskbarActivityContext
+import com.android.launcher3.taskbar.TaskbarControllers
 import com.android.launcher3.taskbar.TaskbarManager
 import com.android.launcher3.taskbar.TaskbarNavButtonController.TaskbarNavButtonCallbacks
 import com.android.launcher3.taskbar.TaskbarViewController
+import com.android.launcher3.taskbar.bubbles.BubbleControllers
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR
 import com.android.launcher3.util.LauncherMultivalentJUnit.Companion.isRunningInRobolectric
@@ -38,6 +40,9 @@
 import com.android.quickstep.AllAppsActionManager
 import com.android.quickstep.TouchInteractionService
 import com.android.quickstep.TouchInteractionService.TISBinder
+import java.lang.reflect.Field
+import java.lang.reflect.ParameterizedType
+import java.util.Optional
 import org.junit.Assume.assumeTrue
 import org.junit.rules.RuleChain
 import org.junit.rules.TestRule
@@ -180,19 +185,38 @@
     fun recreateTaskbar() = instrumentation.runOnMainSync { taskbarManager.recreateTaskbar() }
 
     private fun injectControllers() {
-        val controllers = activityContext.controllers
-        val controllerFieldsByType = controllers.javaClass.fields.associateBy { it.type }
+        val bubbleControllerTypes =
+            BubbleControllers::class.java.fields.map { f ->
+                if (f.type == Optional::class.java) {
+                    (f.genericType as ParameterizedType).actualTypeArguments[0] as Class<*>
+                } else {
+                    f.type
+                }
+            }
         testInstance.javaClass.fields
             .filter { it.isAnnotationPresent(InjectController::class.java) }
             .forEach {
-                it.set(
-                    testInstance,
-                    controllerFieldsByType[it.type]?.get(controllers)
-                        ?: throw NoSuchElementException("Failed to find controller for ${it.type}"),
-                )
+                val controllers: Any =
+                    if (it.type in bubbleControllerTypes) {
+                        activityContext.controllers.bubbleControllers.orElseThrow {
+                            NoSuchElementException("Bubble controllers are not initialized")
+                        }
+                    } else {
+                        activityContext.controllers
+                    }
+                injectController(it, testInstance, controllers)
             }
     }
 
+    private fun injectController(field: Field, testInstance: Any, controllers: Any) {
+        val controllerFieldsByType = controllers.javaClass.fields.associateBy { it.type }
+        field.set(
+            testInstance,
+            controllerFieldsByType[field.type]?.get(controllers)
+                ?: throw NoSuchElementException("Failed to find controller for ${field.type}"),
+        )
+    }
+
     /**
      * Annotates test controller fields to inject the corresponding controllers from the current
      * [TaskbarControllers] instance.
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
index 5d4fdc5..52ca78d 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
@@ -16,18 +16,24 @@
 
 package com.android.launcher3.taskbar.rules
 
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.taskbar.TaskbarActivityContext
 import com.android.launcher3.taskbar.TaskbarKeyguardController
 import com.android.launcher3.taskbar.TaskbarManager
 import com.android.launcher3.taskbar.TaskbarStashController
+import com.android.launcher3.taskbar.bubbles.BubbleBarController
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.NavBarKidsMode
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.UserSetupMode
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
+import com.android.wm.shell.Flags
 import com.google.common.truth.Truth.assertThat
 import org.junit.Assert.assertThrows
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.Description
 import org.junit.runner.RunWith
@@ -39,6 +45,8 @@
 
     private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
 
+    @get:Rule(order = 0) val setFlagsRule = SetFlagsRule()
+
     @Test
     fun testSetup_taskbarInitialized() {
         onSetup { assertThat(activityContext).isInstanceOf(TaskbarActivityContext::class.java) }
@@ -127,6 +135,44 @@
         }
     }
 
+    @EnableFlags(Flags.FLAG_ENABLE_BUBBLE_BAR)
+    @Test
+    fun testInjectBubbleController_bubbleFlagOn_isInjected() {
+        val testClass =
+            object {
+                @InjectController lateinit var controller: BubbleBarController
+                val isInjected: Boolean
+                    get() = ::controller.isInitialized
+            }
+
+        TaskbarUnitTestRule(testClass, context).apply(EMPTY_STATEMENT, DESCRIPTION).evaluate()
+
+        onSetup(TaskbarUnitTestRule(testClass, context)) {
+            assertThat(testClass.isInjected).isTrue()
+        }
+    }
+
+    @DisableFlags(Flags.FLAG_ENABLE_BUBBLE_BAR)
+    @Test
+    fun testInjectBubbleController_bubbleFlagOff_exceptionThrown() {
+        val testClass =
+            object {
+                @InjectController lateinit var controller: BubbleBarController
+            }
+
+        // We cannot use #assertThrows because we also catch an assumption violated exception when
+        // running #evaluate on devices that do not support Taskbar.
+        val result =
+            try {
+                TaskbarUnitTestRule(testClass, context)
+                    .apply(EMPTY_STATEMENT, DESCRIPTION)
+                    .evaluate()
+            } catch (e: NoSuchElementException) {
+                e
+            }
+        assertThat(result).isInstanceOf(NoSuchElementException::class.java)
+    }
+
     @Test
     fun testUserSetupMode_default_isComplete() {
         onSetup { assertThat(activityContext.isUserSetupComplete).isTrue() }