diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/pipeline/MediaDataManager.kt b/packages/SystemUI/src/com/android/systemui/media/controls/pipeline/MediaDataManager.kt
index 14dd99023b923..3012bb41445e6 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/pipeline/MediaDataManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/pipeline/MediaDataManager.kt
@@ -82,7 +82,6 @@ import com.android.systemui.util.traceSection
 import java.io.IOException
 import java.io.PrintWriter
 import java.util.concurrent.Executor
-import java.util.concurrent.Executors
 import javax.inject.Inject
 
 // URI fields to try loading album art from
@@ -154,6 +153,7 @@ private fun allowMediaRecommendations(context: Context): Boolean {
 class MediaDataManager(
     private val context: Context,
     @Background private val backgroundExecutor: Executor,
+    @Main private val uiExecutor: Executor,
     @Main private val foregroundExecutor: DelayableExecutor,
     private val mediaControllerFactory: MediaControllerFactory,
     private val broadcastDispatcher: BroadcastDispatcher,
@@ -171,7 +171,8 @@ class MediaDataManager(
     private val systemClock: SystemClock,
     private val tunerService: TunerService,
     private val mediaFlags: MediaFlags,
-    private val logger: MediaUiEventLogger
+    private val logger: MediaUiEventLogger,
+    private val smartspaceManager: SmartspaceManager,
 ) : Dumpable, BcSmartspaceDataPlugin.SmartspaceTargetListener {
 
     companion object {
@@ -218,6 +219,7 @@ class MediaDataManager(
     constructor(
         context: Context,
         @Background backgroundExecutor: Executor,
+        @Main uiExecutor: Executor,
         @Main foregroundExecutor: DelayableExecutor,
         mediaControllerFactory: MediaControllerFactory,
         dumpManager: DumpManager,
@@ -233,10 +235,12 @@ class MediaDataManager(
         clock: SystemClock,
         tunerService: TunerService,
         mediaFlags: MediaFlags,
-        logger: MediaUiEventLogger
+        logger: MediaUiEventLogger,
+        smartspaceManager: SmartspaceManager,
     ) : this(
         context,
         backgroundExecutor,
+        uiExecutor,
         foregroundExecutor,
         mediaControllerFactory,
         broadcastDispatcher,
@@ -254,7 +258,8 @@ class MediaDataManager(
         clock,
         tunerService,
         mediaFlags,
-        logger
+        logger,
+        smartspaceManager,
     )
 
     private val appChangeReceiver =
@@ -314,21 +319,18 @@ class MediaDataManager(
 
         // Register for Smartspace data updates.
         smartspaceMediaDataProvider.registerListener(this)
-        val smartspaceManager: SmartspaceManager =
-            context.getSystemService(SmartspaceManager::class.java)
         smartspaceSession =
             smartspaceManager.createSmartspaceSession(
                 SmartspaceConfig.Builder(context, SMARTSPACE_UI_SURFACE_LABEL).build()
             )
         smartspaceSession?.let {
             it.addOnTargetsAvailableListener(
-                // Use a new thread listening to Smartspace updates instead of using the existing
-                // backgroundExecutor. SmartspaceSession has scheduled routine updates which can be
-                // unpredictable on test simulators, using the backgroundExecutor makes it's hard to
-                // test the threads numbers.
-                // Switch to use backgroundExecutor when SmartspaceSession has a good way to be
-                // mocked.
-                Executors.newCachedThreadPool(),
+                // Use a main uiExecutor thread listening to Smartspace updates instead of using
+                // the existing background executor.
+                // SmartspaceSession has scheduled routine updates which can be unpredictable on
+                // test simulators, using the backgroundExecutor makes it's hard to test the threads
+                // numbers.
+                uiExecutor,
                 SmartspaceSession.OnTargetsAvailableListener { targets ->
                     smartspaceMediaDataProvider.onTargetsAvailable(targets)
                 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/controls/pipeline/MediaDataManagerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/controls/pipeline/MediaDataManagerTest.kt
index 11eb26b1da02f..8bd71632638d4 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/controls/pipeline/MediaDataManagerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/controls/pipeline/MediaDataManagerTest.kt
@@ -20,6 +20,8 @@ import android.app.Notification
 import android.app.Notification.MediaStyle
 import android.app.PendingIntent
 import android.app.smartspace.SmartspaceAction
+import android.app.smartspace.SmartspaceConfig
+import android.app.smartspace.SmartspaceManager
 import android.app.smartspace.SmartspaceTarget
 import android.content.Intent
 import android.graphics.Bitmap
@@ -106,6 +108,7 @@ class MediaDataManagerTest : SysuiTestCase() {
     lateinit var metadataBuilder: MediaMetadata.Builder
     lateinit var backgroundExecutor: FakeExecutor
     lateinit var foregroundExecutor: FakeExecutor
+    lateinit var uiExecutor: FakeExecutor
     @Mock lateinit var dumpManager: DumpManager
     @Mock lateinit var broadcastDispatcher: BroadcastDispatcher
     @Mock lateinit var mediaTimeoutListener: MediaTimeoutListener
@@ -117,6 +120,7 @@ class MediaDataManagerTest : SysuiTestCase() {
     @Mock lateinit var listener: MediaDataManager.Listener
     @Mock lateinit var pendingIntent: PendingIntent
     @Mock lateinit var activityStarter: ActivityStarter
+    @Mock lateinit var smartspaceManager: SmartspaceManager
     lateinit var smartspaceMediaDataProvider: SmartspaceMediaDataProvider
     @Mock lateinit var mediaSmartspaceTarget: SmartspaceTarget
     @Mock private lateinit var mediaRecommendationItem: SmartspaceAction
@@ -131,6 +135,7 @@ class MediaDataManagerTest : SysuiTestCase() {
     @Mock private lateinit var tunerService: TunerService
     @Captor lateinit var tunableCaptor: ArgumentCaptor<TunerService.Tunable>
     @Captor lateinit var callbackCaptor: ArgumentCaptor<(String, PlaybackState) -> Unit>
+    @Captor lateinit var smartSpaceConfigBuilderCaptor: ArgumentCaptor<SmartspaceConfig>
 
     private val instanceIdSequence = InstanceIdSequenceFake(1 shl 20)
 
@@ -145,6 +150,7 @@ class MediaDataManagerTest : SysuiTestCase() {
     fun setup() {
         foregroundExecutor = FakeExecutor(clock)
         backgroundExecutor = FakeExecutor(clock)
+        uiExecutor = FakeExecutor(clock)
         smartspaceMediaDataProvider = SmartspaceMediaDataProvider()
         Settings.Secure.putInt(
             context.contentResolver,
@@ -155,6 +161,7 @@ class MediaDataManagerTest : SysuiTestCase() {
             MediaDataManager(
                 context = context,
                 backgroundExecutor = backgroundExecutor,
+                uiExecutor = uiExecutor,
                 foregroundExecutor = foregroundExecutor,
                 mediaControllerFactory = mediaControllerFactory,
                 broadcastDispatcher = broadcastDispatcher,
@@ -172,7 +179,8 @@ class MediaDataManagerTest : SysuiTestCase() {
                 systemClock = clock,
                 tunerService = tunerService,
                 mediaFlags = mediaFlags,
-                logger = logger
+                logger = logger,
+                smartspaceManager = smartspaceManager,
             )
         verify(tunerService)
             .addTunable(capture(tunableCaptor), eq(Settings.Secure.MEDIA_CONTROLS_RECOMMENDATION))
@@ -191,6 +199,7 @@ class MediaDataManagerTest : SysuiTestCase() {
                 putString(MediaMetadata.METADATA_KEY_ARTIST, SESSION_ARTIST)
                 putString(MediaMetadata.METADATA_KEY_TITLE, SESSION_TITLE)
             }
+        verify(smartspaceManager).createSmartspaceSession(capture(smartSpaceConfigBuilderCaptor))
         whenever(mediaControllerFactory.create(eq(session.sessionToken))).thenReturn(controller)
         whenever(controller.transportControls).thenReturn(transportControls)
         whenever(controller.playbackInfo).thenReturn(playbackInfo)
@@ -767,15 +776,14 @@ class MediaDataManagerTest : SysuiTestCase() {
             .onSmartspaceMediaDataLoaded(anyObject(), anyObject(), anyBoolean())
     }
 
-    @Ignore("b/233283726")
     @Test
     fun testOnSmartspaceMediaDataLoaded_hasNoneMediaTarget_callsRemoveListener() {
         smartspaceMediaDataProvider.onTargetsAvailable(listOf(mediaSmartspaceTarget))
         verify(logger).getNewInstanceId()
 
         smartspaceMediaDataProvider.onTargetsAvailable(listOf())
-        foregroundExecutor.advanceClockToLast()
-        foregroundExecutor.runAllReady()
+        uiExecutor.advanceClockToLast()
+        uiExecutor.runAllReady()
 
         verify(listener).onSmartspaceMediaDataRemoved(eq(KEY_MEDIA_SMARTSPACE), eq(false))
         verifyNoMoreInteractions(logger)