Add loading states for media changes.

Flag: ACONFIG media_controls_refactor DISABLED
Bug: 326281896
Test: atest SystemUiRoboTests:MediaCarouselInteractorTest
Change-Id: I09477546282bc297b317f695e81871404b7dd417
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/data/repository/MediaFilterRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/data/repository/MediaFilterRepositoryTest.kt
index 956ef66..33eb90a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/data/repository/MediaFilterRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/data/repository/MediaFilterRepositoryTest.kt
@@ -26,7 +26,9 @@
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.media.controls.MediaTestHelper
 import com.android.systemui.media.controls.shared.model.MediaData
+import com.android.systemui.media.controls.shared.model.MediaDataLoadingModel
 import com.android.systemui.media.controls.shared.model.SmartspaceMediaData
+import com.android.systemui.media.controls.shared.model.SmartspaceMediaLoadingModel
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.test.runTest
@@ -144,6 +146,37 @@
             assertThat(smartspaceMediaData?.isActive).isFalse()
         }
 
+    @Test
+    fun addMediaDataLoadingState() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(underTest.mediaDataLoadedStates)
+            val instanceId = InstanceId.fakeInstanceId(123)
+            val mediaLoadedStates = mutableListOf(MediaDataLoadingModel.Loaded(instanceId))
+
+            underTest.addMediaDataLoadingState(MediaDataLoadingModel.Loaded(instanceId))
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStates)
+
+            mediaLoadedStates.remove(MediaDataLoadingModel.Loaded(instanceId))
+
+            underTest.addMediaDataLoadingState(MediaDataLoadingModel.Removed(instanceId))
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStates)
+        }
+
+    @Test
+    fun setRecommendationsLoadingState() =
+        testScope.runTest {
+            val recommendationsLoadingState by
+                collectLastValue(underTest.recommendationsLoadingState)
+            val recommendationsLoadingModel =
+                SmartspaceMediaLoadingModel.Loaded(KEY_MEDIA_SMARTSPACE)
+
+            underTest.setRecommedationsLoadingState(recommendationsLoadingModel)
+
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
+        }
+
     companion object {
         private const val KEY = "KEY"
         private const val KEY_MEDIA_SMARTSPACE = "MEDIA_SMARTSPACE_ID"
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/interactor/MediaCarouselInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/interactor/MediaCarouselInteractorTest.kt
index d9d84f2..a0a1eb3 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/interactor/MediaCarouselInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/interactor/MediaCarouselInteractorTest.kt
@@ -20,6 +20,7 @@
 import android.graphics.drawable.Icon
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.internal.logging.InstanceId
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.flags.Flags
@@ -28,13 +29,20 @@
 import com.android.systemui.media.controls.MediaTestHelper
 import com.android.systemui.media.controls.data.repository.MediaFilterRepository
 import com.android.systemui.media.controls.data.repository.mediaFilterRepository
+import com.android.systemui.media.controls.domain.pipeline.MediaDataFilterImpl
 import com.android.systemui.media.controls.domain.pipeline.interactor.MediaCarouselInteractor
 import com.android.systemui.media.controls.domain.pipeline.interactor.mediaCarouselInteractor
+import com.android.systemui.media.controls.domain.pipeline.mediaDataFilter
 import com.android.systemui.media.controls.shared.model.MediaData
+import com.android.systemui.media.controls.shared.model.MediaDataLoadingModel
 import com.android.systemui.media.controls.shared.model.SmartspaceMediaData
+import com.android.systemui.media.controls.shared.model.SmartspaceMediaLoadingModel
+import com.android.systemui.statusbar.notificationLockscreenUserManager
 import com.android.systemui.testKosmos
+import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.test.runTest
+import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -45,9 +53,17 @@
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
 
+    private val mediaDataFilter: MediaDataFilterImpl = kosmos.mediaDataFilter
+    private val notificationLockscreenUserManager = kosmos.notificationLockscreenUserManager
     private val mediaFilterRepository: MediaFilterRepository = kosmos.mediaFilterRepository
+
     private val underTest: MediaCarouselInteractor = kosmos.mediaCarouselInteractor
 
+    @Before
+    fun setUp() {
+        underTest.start()
+    }
+
     @Test
     fun addUserMediaEntry_activeThenInactivate() =
         testScope.runTest {
@@ -56,7 +72,7 @@
             val hasActiveMedia by collectLastValue(underTest.hasActiveMedia)
             val hasAnyMedia by collectLastValue(underTest.hasAnyMedia)
 
-            val userMedia = MediaData().copy(active = true)
+            val userMedia = MediaData(active = true)
 
             mediaFilterRepository.addSelectedUserMediaEntry(userMedia)
 
@@ -79,7 +95,7 @@
             val hasActiveMedia by collectLastValue(underTest.hasActiveMedia)
             val hasAnyMedia by collectLastValue(underTest.hasAnyMedia)
 
-            val userMedia = MediaData().copy(active = false)
+            val userMedia = MediaData(active = false)
             val instanceId = userMedia.instanceId
 
             mediaFilterRepository.addSelectedUserMediaEntry(userMedia)
@@ -112,7 +128,7 @@
                     isActive = true,
                     recommendations = MediaTestHelper.getValidRecommendationList(icon),
                 )
-            val userMedia = MediaData().copy(active = false)
+            val userMedia = MediaData(active = false)
 
             mediaFilterRepository.setRecommendation(userMediaRecommendation)
 
@@ -199,7 +215,80 @@
     fun hasActiveMediaOrRecommendation_nothingSet_returnsFalse() =
         testScope.runTest { assertThat(underTest.hasActiveMediaOrRecommendation.value).isFalse() }
 
+    @Test
+    fun onMediaDataUpdated_updatesLoadingState() =
+        testScope.runTest {
+            whenever(notificationLockscreenUserManager.isCurrentProfile(USER_ID)).thenReturn(true)
+            whenever(notificationLockscreenUserManager.isProfileAvailable(USER_ID)).thenReturn(true)
+            val mediaDataLoadedStates by collectLastValue(underTest.mediaDataLoadedStates)
+            val instanceId = InstanceId.fakeInstanceId(123)
+            val mediaLoadedStates: MutableList<MediaDataLoadingModel> = mutableListOf()
+
+            mediaLoadedStates.add(MediaDataLoadingModel.Loaded(instanceId))
+            mediaDataFilter.onMediaDataLoaded(
+                KEY,
+                KEY,
+                MediaData(userId = USER_ID, instanceId = instanceId)
+            )
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStates)
+
+            val newInstanceId = InstanceId.fakeInstanceId(321)
+
+            mediaLoadedStates.add(MediaDataLoadingModel.Loaded(newInstanceId))
+            mediaDataFilter.onMediaDataLoaded(
+                KEY_2,
+                KEY_2,
+                MediaData(userId = USER_ID, instanceId = newInstanceId)
+            )
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStates)
+
+            mediaLoadedStates.remove(MediaDataLoadingModel.Loaded(instanceId))
+
+            mediaDataFilter.onMediaDataRemoved(KEY)
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStates)
+
+            mediaLoadedStates.remove(MediaDataLoadingModel.Loaded(newInstanceId))
+
+            mediaDataFilter.onMediaDataRemoved(KEY_2)
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStates)
+        }
+
+    @Test
+    fun onMediaRecommendationsUpdated_updatesLoadingState() =
+        testScope.runTest {
+            whenever(notificationLockscreenUserManager.isCurrentProfile(USER_ID)).thenReturn(true)
+            whenever(notificationLockscreenUserManager.isProfileAvailable(USER_ID)).thenReturn(true)
+            val recommendationsLoadingState by
+                collectLastValue(underTest.recommendationsLoadingState)
+            val icon = Icon.createWithResource(context, R.drawable.ic_media_play)
+            val mediaRecommendations =
+                SmartspaceMediaData(
+                    targetId = KEY_MEDIA_SMARTSPACE,
+                    isActive = true,
+                    recommendations = MediaTestHelper.getValidRecommendationList(icon),
+                )
+            var recommendationsLoadingModel: SmartspaceMediaLoadingModel =
+                SmartspaceMediaLoadingModel.Loaded(KEY_MEDIA_SMARTSPACE, isPrioritized = true)
+
+            mediaDataFilter.onSmartspaceMediaDataLoaded(KEY_MEDIA_SMARTSPACE, mediaRecommendations)
+
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
+
+            recommendationsLoadingModel = SmartspaceMediaLoadingModel.Removed(KEY_MEDIA_SMARTSPACE)
+
+            mediaDataFilter.onSmartspaceMediaDataRemoved(KEY_MEDIA_SMARTSPACE)
+
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
+        }
+
     companion object {
+        private const val KEY = "key"
+        private const val KEY_2 = "key2"
+        private const val USER_ID = 0
         private const val KEY_MEDIA_SMARTSPACE = "MEDIA_SMARTSPACE_ID"
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/data/repository/MediaFilterRepository.kt b/packages/SystemUI/src/com/android/systemui/media/controls/data/repository/MediaFilterRepository.kt
index df34169..9dc5900 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/data/repository/MediaFilterRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/data/repository/MediaFilterRepository.kt
@@ -19,7 +19,9 @@
 import com.android.internal.logging.InstanceId
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.media.controls.shared.model.MediaData
+import com.android.systemui.media.controls.shared.model.MediaDataLoadingModel
 import com.android.systemui.media.controls.shared.model.SmartspaceMediaData
+import com.android.systemui.media.controls.shared.model.SmartspaceMediaLoadingModel
 import javax.inject.Inject
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
@@ -46,6 +48,16 @@
         MutableStateFlow(LinkedHashMap())
     val allUserEntries: StateFlow<Map<String, MediaData>> = _allUserEntries.asStateFlow()
 
+    private val _mediaDataLoadedStates: MutableStateFlow<List<MediaDataLoadingModel>> =
+        MutableStateFlow(mutableListOf())
+    val mediaDataLoadedStates: StateFlow<List<MediaDataLoadingModel>> =
+        _mediaDataLoadedStates.asStateFlow()
+
+    private val _recommendationsLoadingState: MutableStateFlow<SmartspaceMediaLoadingModel> =
+        MutableStateFlow(SmartspaceMediaLoadingModel.Unknown)
+    val recommendationsLoadingState: StateFlow<SmartspaceMediaLoadingModel> =
+        _recommendationsLoadingState.asStateFlow()
+
     fun addMediaEntry(key: String, data: MediaData) {
         val entries = LinkedHashMap<String, MediaData>(_allUserEntries.value)
         entries[key] = data
@@ -110,4 +122,25 @@
     fun setReactivatedId(instanceId: InstanceId?) {
         _reactivatedId.value = instanceId
     }
+
+    fun addMediaDataLoadingState(mediaDataLoadingModel: MediaDataLoadingModel) {
+        // Filter out previous loading state that has same [InstanceId].
+        val loadedStates =
+            _mediaDataLoadedStates.value.filter { loadedModel ->
+                loadedModel !is MediaDataLoadingModel.Loaded ||
+                    !loadedModel.equalInstanceIds(mediaDataLoadingModel)
+            }
+
+        _mediaDataLoadedStates.value =
+            loadedStates +
+                if (mediaDataLoadingModel is MediaDataLoadingModel.Loaded) {
+                    listOf(mediaDataLoadingModel)
+                } else {
+                    emptyList()
+                }
+    }
+
+    fun setRecommedationsLoadingState(smartspaceMediaLoadingModel: SmartspaceMediaLoadingModel) {
+        _recommendationsLoadingState.value = smartspaceMediaLoadingModel
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImpl.kt b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImpl.kt
index d40069c..a30e582 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImpl.kt
@@ -28,7 +28,9 @@
 import com.android.systemui.media.controls.data.repository.MediaFilterRepository
 import com.android.systemui.media.controls.shared.model.EXTRA_KEY_TRIGGER_RESUME
 import com.android.systemui.media.controls.shared.model.MediaData
+import com.android.systemui.media.controls.shared.model.MediaDataLoadingModel
 import com.android.systemui.media.controls.shared.model.SmartspaceMediaData
+import com.android.systemui.media.controls.shared.model.SmartspaceMediaLoadingModel
 import com.android.systemui.media.controls.util.MediaFlags
 import com.android.systemui.media.controls.util.MediaUiEventLogger
 import com.android.systemui.settings.UserTracker
@@ -67,9 +69,6 @@
     private val mediaFlags: MediaFlags,
     private val mediaFilterRepository: MediaFilterRepository,
 ) : MediaDataManager.Listener {
-    private val _listeners: MutableSet<Listener> = mutableSetOf()
-    val listeners: Set<Listener>
-        get() = _listeners.toSet()
     lateinit var mediaDataManager: MediaDataManager
 
     // Ensure the field (and associated reference) isn't removed during optimization.
@@ -111,8 +110,9 @@
 
         mediaFilterRepository.addSelectedUserMediaEntry(data)
 
-        // Notify listeners
-        listeners.forEach { it.onMediaDataLoaded(data.instanceId) }
+        mediaFilterRepository.addMediaDataLoadingState(
+            MediaDataLoadingModel.Loaded(data.instanceId)
+        )
     }
 
     override fun onSmartspaceMediaDataLoaded(
@@ -159,7 +159,7 @@
             // reactivate.
             if (shouldReactivate) {
                 val lastActiveId = sorted.lastKey() // most recently active id
-                // Notify listeners to consider this media active
+                // Update loading state to consider this media active
                 Log.d(TAG, "reactivating $lastActiveId instead of smartspace")
                 mediaFilterRepository.setReactivatedId(lastActiveId)
                 val mediaData = sorted[lastActiveId]!!.copy(active = true)
@@ -168,15 +168,9 @@
                     mediaData.packageName,
                     mediaData.instanceId
                 )
-                listeners.forEach {
-                    it.onMediaDataLoaded(
-                        lastActiveId,
-                        receivedSmartspaceCardLatency =
-                            (systemClock.currentTimeMillis() - data.headphoneConnectionTimeMillis)
-                                .toInt(),
-                        isSsReactivated = true
-                    )
-                }
+                mediaFilterRepository.addMediaDataLoadingState(
+                    MediaDataLoadingModel.Loaded(lastActiveId)
+                )
             }
         } else if (data.isActive) {
             // Mark to prioritize Smartspace card if no recent media.
@@ -192,15 +186,18 @@
             smartspaceMediaData.packageName,
             smartspaceMediaData.instanceId
         )
-        listeners.forEach { it.onSmartspaceMediaDataLoaded(key, shouldPrioritizeMutable) }
+        mediaFilterRepository.setRecommedationsLoadingState(
+            SmartspaceMediaLoadingModel.Loaded(key, shouldPrioritizeMutable)
+        )
     }
 
     override fun onMediaDataRemoved(key: String) {
         mediaFilterRepository.removeMediaEntry(key)?.let { mediaData ->
             val instanceId = mediaData.instanceId
             mediaFilterRepository.removeSelectedUserMediaEntry(instanceId)?.let {
-                // Only notify listeners if something actually changed
-                listeners.forEach { it.onMediaDataRemoved(instanceId) }
+                mediaFilterRepository.addMediaDataLoadingState(
+                    MediaDataLoadingModel.Removed(instanceId)
+                )
             }
         }
     }
@@ -210,11 +207,11 @@
         mediaFilterRepository.reactivatedId.value?.let { lastActiveId ->
             mediaFilterRepository.setReactivatedId(null)
             Log.d(TAG, "expiring reactivated key $lastActiveId")
-            // Notify listeners to update with actual active value
+            // Update loading state with actual active value
             mediaFilterRepository.selectedUserEntries.value[lastActiveId]?.let {
-                listeners.forEach { listener ->
-                    listener.onMediaDataLoaded(lastActiveId, immediately)
-                }
+                mediaFilterRepository.addMediaDataLoadingState(
+                    MediaDataLoadingModel.Loaded(lastActiveId, immediately)
+                )
             }
         }
 
@@ -227,7 +224,9 @@
                 )
             )
         }
-        listeners.forEach { it.onSmartspaceMediaDataRemoved(key, immediately) }
+        mediaFilterRepository.setRecommedationsLoadingState(
+            SmartspaceMediaLoadingModel.Removed(key, immediately)
+        )
     }
 
     @VisibleForTesting
@@ -238,29 +237,37 @@
                 // Only remove media when the profile is unavailable.
                 if (DEBUG) Log.d(TAG, "Removing $key after profile change")
                 mediaFilterRepository.removeSelectedUserMediaEntry(data.instanceId, data)
-                listeners.forEach { listener -> listener.onMediaDataRemoved(data.instanceId) }
+                mediaFilterRepository.addMediaDataLoadingState(
+                    MediaDataLoadingModel.Removed(data.instanceId)
+                )
             }
         }
     }
 
     @VisibleForTesting
     internal fun handleUserSwitched() {
-        // If the user changes, remove all current MediaData objects and inform listeners
-        val listenersCopy = listeners
+        // If the user changes, remove all current MediaData objects.
         val keyCopy = mediaFilterRepository.selectedUserEntries.value.keys.toMutableList()
-        // Clear the list first, to make sure callbacks from listeners if we have any entries
-        // are up to date
+        // Clear the list first and update loading state to remove media from UI.
         mediaFilterRepository.clearSelectedUserMedia()
         keyCopy.forEach { instanceId ->
             if (DEBUG) Log.d(TAG, "Removing $instanceId after user change")
-            listenersCopy.forEach { listener -> listener.onMediaDataRemoved(instanceId) }
+            mediaFilterRepository.addMediaDataLoadingState(
+                MediaDataLoadingModel.Removed(instanceId)
+            )
         }
 
         mediaFilterRepository.allUserEntries.value.forEach { (key, data) ->
             if (lockscreenUserManager.isCurrentProfile(data.userId)) {
-                if (DEBUG) Log.d(TAG, "Re-adding $key after user change")
+                if (DEBUG)
+                    Log.d(
+                        TAG,
+                        "Re-adding $key with instanceId=${data.instanceId} after user change"
+                    )
                 mediaFilterRepository.addSelectedUserMediaEntry(data)
-                listenersCopy.forEach { listener -> listener.onMediaDataLoaded(data.instanceId) }
+                mediaFilterRepository.addMediaDataLoadingState(
+                    MediaDataLoadingModel.Loaded(data.instanceId)
+                )
             }
         }
     }
@@ -310,12 +317,6 @@
         }
     }
 
-    /** Add a listener for filtered [MediaData] changes */
-    fun addListener(listener: Listener) = _listeners.add(listener)
-
-    /** Remove a listener that was registered with addListener */
-    fun removeListener(listener: Listener) = _listeners.remove(listener)
-
     /**
      * Return the time since last active for the most-recent media.
      *
@@ -335,48 +336,6 @@
         return sortedEntries[lastActiveInstanceId]?.let { now - it.lastActive } ?: Long.MAX_VALUE
     }
 
-    interface Listener {
-        /**
-         * Called whenever there's new MediaData Loaded for the consumption in views.
-         *
-         * @param immediately indicates should apply the UI changes immediately, otherwise wait
-         *   until the next refresh-round before UI becomes visible. True by default to take in
-         *   place immediately.
-         * @param receivedSmartspaceCardLatency is the latency between headphone connects and sysUI
-         *   displays Smartspace media targets. Will be 0 if the data is not activated by Smartspace
-         *   signal.
-         * @param isSsReactivated indicates resume media card is reactivated by Smartspace
-         *   recommendation signal
-         */
-        fun onMediaDataLoaded(
-            instanceId: InstanceId,
-            immediately: Boolean = true,
-            receivedSmartspaceCardLatency: Int = 0,
-            isSsReactivated: Boolean = false,
-        )
-
-        /**
-         * Called whenever there's new Smartspace media data loaded.
-         *
-         * @param shouldPrioritize indicates the sorting priority of the Smartspace card. If true,
-         *   it will be prioritized as the first card. Otherwise, it will show up as the last card
-         *   as default.
-         */
-        fun onSmartspaceMediaDataLoaded(key: String, shouldPrioritize: Boolean = false)
-
-        /** Called whenever a previously existing Media notification was removed. */
-        fun onMediaDataRemoved(instanceId: InstanceId)
-
-        /**
-         * Called whenever a previously existing Smartspace media data was removed.
-         *
-         * @param immediately indicates should apply the UI changes immediately, otherwise wait
-         *   until the next refresh-round before UI becomes visible. True by default to take in
-         *   place immediately.
-         */
-        fun onSmartspaceMediaDataRemoved(key: String, immediately: Boolean = true)
-    }
-
     companion object {
         /**
          * Maximum age of a media control to re-activate on smartspace signal. If there is no media
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/interactor/MediaCarouselInteractor.kt b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/interactor/MediaCarouselInteractor.kt
index 7dbca0a..cdcf363 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/interactor/MediaCarouselInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/interactor/MediaCarouselInteractor.kt
@@ -34,11 +34,14 @@
 import com.android.systemui.media.controls.domain.pipeline.MediaSessionBasedFilter
 import com.android.systemui.media.controls.domain.pipeline.MediaTimeoutListener
 import com.android.systemui.media.controls.domain.resume.MediaResumeListener
+import com.android.systemui.media.controls.shared.model.MediaDataLoadingModel
+import com.android.systemui.media.controls.shared.model.SmartspaceMediaLoadingModel
 import com.android.systemui.media.controls.util.MediaFlags
 import java.io.PrintWriter
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
@@ -109,6 +112,14 @@
             .distinctUntilChanged()
             .stateIn(applicationScope, SharingStarted.WhileSubscribed(), false)
 
+    /** The most recent list of loaded media controls. */
+    val mediaDataLoadedStates: Flow<List<MediaDataLoadingModel>> =
+        mediaFilterRepository.mediaDataLoadedStates
+
+    /** The most recent change to loaded media recommendations. */
+    val recommendationsLoadingState: Flow<SmartspaceMediaLoadingModel> =
+        mediaFilterRepository.recommendationsLoadingState
+
     override fun start() {
         if (!mediaFlags.isMediaControlsRefactorEnabled()) {
             return
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/shared/model/MediaDataLoadingModel.kt b/packages/SystemUI/src/com/android/systemui/media/controls/shared/model/MediaDataLoadingModel.kt
new file mode 100644
index 0000000..bd42a4d
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/shared/model/MediaDataLoadingModel.kt
@@ -0,0 +1,46 @@
+/*
+ * 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.systemui.media.controls.shared.model
+
+import com.android.internal.logging.InstanceId
+
+/** Models media data loading state. */
+sealed class MediaDataLoadingModel {
+    /** The initial loading state when no media data has yet loaded. */
+    data object Unknown : MediaDataLoadingModel()
+
+    /** Media data has been loaded. */
+    data class Loaded(
+        val instanceId: InstanceId,
+        val immediatelyUpdateUi: Boolean = true,
+    ) : MediaDataLoadingModel() {
+
+        /** Returns true if [other] has the same instance id, false otherwise. */
+        fun equalInstanceIds(other: MediaDataLoadingModel): Boolean {
+            return when (other) {
+                is Loaded -> other.instanceId == instanceId
+                is Removed -> other.instanceId == instanceId
+                Unknown -> false
+            }
+        }
+    }
+
+    /** Media data has been removed. */
+    data class Removed(
+        val instanceId: InstanceId,
+    ) : MediaDataLoadingModel()
+}
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/shared/model/SmartspaceMediaLoadingModel.kt b/packages/SystemUI/src/com/android/systemui/media/controls/shared/model/SmartspaceMediaLoadingModel.kt
new file mode 100644
index 0000000..6c1e536
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/shared/model/SmartspaceMediaLoadingModel.kt
@@ -0,0 +1,35 @@
+/*
+ * 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.systemui.media.controls.shared.model
+
+/** Models smartspace media loading state. */
+sealed class SmartspaceMediaLoadingModel {
+    /** The initial loading state when no smartspace media has yet loaded. */
+    data object Unknown : SmartspaceMediaLoadingModel()
+
+    /** Smartspace media has been loaded. */
+    data class Loaded(
+        val key: String,
+        val isPrioritized: Boolean = false,
+    ) : SmartspaceMediaLoadingModel()
+
+    /** Smartspace media has been removed. */
+    data class Removed(
+        val key: String,
+        val immediatelyUpdateUi: Boolean = true,
+    ) : SmartspaceMediaLoadingModel()
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImplTest.kt
index b70cc30..fe8fdc0 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataFilterImplTest.kt
@@ -29,7 +29,9 @@
 import com.android.systemui.media.controls.data.repository.MediaFilterRepository
 import com.android.systemui.media.controls.shared.model.EXTRA_KEY_TRIGGER_RESUME
 import com.android.systemui.media.controls.shared.model.MediaData
+import com.android.systemui.media.controls.shared.model.MediaDataLoadingModel
 import com.android.systemui.media.controls.shared.model.SmartspaceMediaData
+import com.android.systemui.media.controls.shared.model.SmartspaceMediaLoadingModel
 import com.android.systemui.media.controls.ui.controller.MediaPlayerData
 import com.android.systemui.media.controls.util.MediaFlags
 import com.android.systemui.media.controls.util.MediaUiEventLogger
@@ -48,12 +50,10 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.ArgumentMatchers.anyBoolean
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.ArgumentMatchers.anyLong
 import org.mockito.Mock
 import org.mockito.Mockito.never
-import org.mockito.Mockito.reset
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
@@ -76,7 +76,6 @@
 @TestableLooper.RunWithLooper
 class MediaDataFilterImplTest : SysuiTestCase() {
 
-    @Mock private lateinit var listener: MediaDataFilterImpl.Listener
     @Mock private lateinit var userTracker: UserTracker
     @Mock private lateinit var broadcastSender: BroadcastSender
     @Mock private lateinit var mediaDataManager: MediaDataManager
@@ -89,7 +88,7 @@
     @Mock private lateinit var cardAction: SmartspaceAction
 
     private lateinit var mediaDataFilter: MediaDataFilterImpl
-    private lateinit var mediaFilterRepository: MediaFilterRepository
+    private lateinit var repository: MediaFilterRepository
     private lateinit var testScope: TestScope
     private lateinit var dataMain: MediaData
     private lateinit var dataGuest: MediaData
@@ -102,7 +101,7 @@
         MediaPlayerData.clear()
         whenever(mediaFlags.isPersistentSsCardEnabled()).thenReturn(false)
         testScope = TestScope()
-        mediaFilterRepository = MediaFilterRepository()
+        repository = MediaFilterRepository()
         mediaDataFilter =
             MediaDataFilterImpl(
                 context,
@@ -113,10 +112,9 @@
                 clock,
                 logger,
                 mediaFlags,
-                mediaFilterRepository,
+                repository,
             )
         mediaDataFilter.mediaDataManager = mediaDataManager
-        mediaDataFilter.addListener(listener)
 
         // Start all tests as main user
         setUser(USER_MAIN)
@@ -162,91 +160,114 @@
     }
 
     @Test
-    fun testOnDataLoadedForCurrentUser_callsListener() {
-        // GIVEN a media for main user
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
+    fun onDataLoadedForCurrentUser_updatesLoadedStates() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val mediaDataLoadingModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
 
-        // THEN we should tell the listener
-        verify(listener).onMediaDataLoaded(eq(dataMain.instanceId), eq(true), eq(0), eq(false))
-    }
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaDataLoadingModel)
+        }
 
     @Test
-    fun testOnDataLoadedForGuest_doesNotCallListener() {
-        // GIVEN a media for guest user
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataGuest)
+    fun onDataLoadedForGuest_doesNotUpdateLoadedStates() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
 
-        // THEN we should NOT tell the listener
-        verify(listener, never()).onMediaDataLoaded(any(), anyBoolean(), anyInt(), anyBoolean())
-    }
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataGuest)
+
+            assertThat(mediaDataLoadedStates).isNotEqualTo(mediaLoadedStatesModel)
+        }
 
     @Test
-    fun testOnRemovedForCurrent_callsListener() {
-        // GIVEN a media was removed for main user
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
-        mediaDataFilter.onMediaDataRemoved(KEY)
+    fun onRemovedForCurrent_updatesLoadedStates() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val mediaLoadedStatesModel =
+                mutableListOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
 
-        // THEN we should tell the listener
-        verify(listener).onMediaDataRemoved(eq(dataMain.instanceId))
-    }
+            // GIVEN a media was removed for main user
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
+
+            mediaLoadedStatesModel.remove(MediaDataLoadingModel.Loaded(dataMain.instanceId))
+            mediaDataFilter.onMediaDataRemoved(KEY)
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
+        }
 
     @Test
-    fun testOnRemovedForGuest_doesNotCallListener() {
-        // GIVEN a media was removed for guest user
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataGuest)
-        mediaDataFilter.onMediaDataRemoved(KEY)
+    fun onRemovedForGuest_doesNotUpdateLoadedStates() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
 
-        // THEN we should NOT tell the listener
-        verify(listener, never()).onMediaDataRemoved(eq(dataGuest.instanceId))
-    }
+            // GIVEN a media was removed for guest user
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataGuest)
+            mediaDataFilter.onMediaDataRemoved(KEY)
+
+            assertThat(mediaDataLoadedStates).isEmpty()
+        }
 
     @Test
-    fun testOnUserSwitched_removesOldUserControls() {
-        // GIVEN that we have a media loaded for main user
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
+    fun onUserSwitched_removesOldUserControls() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
 
-        // and we switch to guest user
-        setUser(USER_GUEST)
+            // GIVEN that we have a media loaded for main user
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
 
-        // THEN we should remove the main user's media
-        verify(listener).onMediaDataRemoved(eq(dataMain.instanceId))
-    }
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
+
+            // and we switch to guest user
+            setUser(USER_GUEST)
+
+            // THEN we should remove the main user's media
+            assertThat(mediaDataLoadedStates).isEmpty()
+        }
 
     @Test
-    fun testOnUserSwitched_addsNewUserControls() {
-        // GIVEN that we had some media for both users
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
-        mediaDataFilter.onMediaDataLoaded(KEY_ALT, null, dataGuest)
-        reset(listener)
+    fun onUserSwitched_addsNewUserControls() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val guestLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataGuest.instanceId))
+            val mainLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
 
-        // and we switch to guest user
-        setUser(USER_GUEST)
+            // GIVEN that we had some media for both users
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
+            mediaDataFilter.onMediaDataLoaded(KEY_ALT, null, dataGuest)
 
-        // THEN we should add back the guest user media
-        verify(listener).onMediaDataLoaded(eq(dataGuest.instanceId), eq(true), eq(0), eq(false))
+            // and we switch to guest user
+            setUser(USER_GUEST)
 
-        // but not the main user's
-        verify(listener, never())
-            .onMediaDataLoaded(eq(dataMain.instanceId), anyBoolean(), anyInt(), anyBoolean())
-    }
+            assertThat(mediaDataLoadedStates).isEqualTo(guestLoadedStatesModel)
+            assertThat(mediaDataLoadedStates).isNotEqualTo(mainLoadedStatesModel)
+        }
 
     @Test
-    fun testOnProfileChanged_profileUnavailable_loadControls() {
-        // GIVEN that we had some media for both profiles
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
-        mediaDataFilter.onMediaDataLoaded(KEY_ALT, null, dataPrivateProfile)
-        reset(listener)
+    fun onProfileChanged_profileUnavailable_updateStates() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
 
-        // and we change profile status
-        setPrivateProfileUnavailable()
+            // GIVEN that we had some media for both profiles
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
+            mediaDataFilter.onMediaDataLoaded(KEY_ALT, null, dataPrivateProfile)
 
-        // THEN we should add the private profile media
-        verify(listener).onMediaDataRemoved(eq(dataPrivateProfile.instanceId))
-    }
+            // and we change profile status
+            setPrivateProfileUnavailable()
+
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
+            // THEN we should remove the private profile media
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
+        }
 
     @Test
     fun hasAnyMedia_mediaSet_returnsTrue() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
             mediaDataFilter.onMediaDataLoaded(KEY, oldKey = null, data = dataMain)
 
             assertThat(hasAnyMedia(selectedUserEntries)).isTrue()
@@ -255,7 +276,7 @@
     @Test
     fun hasAnyMedia_recommendationSet_returnsFalse() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
             assertThat(hasAnyMedia(selectedUserEntries)).isFalse()
@@ -264,8 +285,8 @@
     @Test
     fun hasAnyMediaOrRecommendation_mediaSet_returnsTrue() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
             mediaDataFilter.onMediaDataLoaded(KEY, oldKey = null, data = dataMain)
 
             assertThat(hasAnyMediaOrRecommendation(selectedUserEntries, smartspaceMediaData))
@@ -275,8 +296,8 @@
     @Test
     fun hasAnyMediaOrRecommendation_recommendationSet_returnsTrue() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
             assertThat(hasAnyMediaOrRecommendation(selectedUserEntries, smartspaceMediaData))
@@ -286,7 +307,7 @@
     @Test
     fun hasActiveMedia_inactiveMediaSet_returnsFalse() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
 
             val data = dataMain.copy(active = false)
             mediaDataFilter.onMediaDataLoaded(KEY, oldKey = null, data = data)
@@ -297,7 +318,7 @@
     @Test
     fun hasActiveMedia_activeMediaSet_returnsTrue() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
             val data = dataMain.copy(active = true)
             mediaDataFilter.onMediaDataLoaded(KEY, oldKey = null, data = data)
 
@@ -307,9 +328,9 @@
     @Test
     fun hasActiveMediaOrRecommendation_inactiveMediaSet_returnsFalse() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
             val data = dataMain.copy(active = false)
             mediaDataFilter.onMediaDataLoaded(KEY, oldKey = null, data = data)
 
@@ -326,9 +347,9 @@
     @Test
     fun hasActiveMediaOrRecommendation_activeMediaSet_returnsTrue() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
             val data = dataMain.copy(active = true)
             mediaDataFilter.onMediaDataLoaded(KEY, oldKey = null, data = data)
 
@@ -345,9 +366,9 @@
     @Test
     fun hasActiveMediaOrRecommendation_inactiveRecommendationSet_returnsFalse() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
             whenever(smartspaceData.isActive).thenReturn(false)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
@@ -364,9 +385,9 @@
     @Test
     fun hasActiveMediaOrRecommendation_invalidRecommendationSet_returnsFalse() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
             whenever(smartspaceData.isValid()).thenReturn(false)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
@@ -383,9 +404,9 @@
     @Test
     fun hasActiveMediaOrRecommendation_activeAndValidRecommendationSet_returnsTrue() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
             whenever(smartspaceData.isActive).thenReturn(true)
             whenever(smartspaceData.isValid()).thenReturn(true)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
@@ -401,10 +422,10 @@
         }
 
     @Test
-    fun testHasAnyMediaOrRecommendation_onlyCurrentUser() =
+    fun hasAnyMediaOrRecommendation_onlyCurrentUser() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
             assertThat(hasAnyMediaOrRecommendation(selectedUserEntries, smartspaceMediaData))
                 .isFalse()
 
@@ -415,11 +436,11 @@
         }
 
     @Test
-    fun testHasActiveMediaOrRecommendation_onlyCurrentUser() =
+    fun hasActiveMediaOrRecommendation_onlyCurrentUser() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -443,10 +464,10 @@
         }
 
     @Test
-    fun testOnNotificationRemoved_doesNotHaveMedia() =
+    fun onNotificationRemoved_doesNotHaveMedia() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
 
             mediaDataFilter.onMediaDataLoaded(KEY, oldKey = null, data = dataMain)
             mediaDataFilter.onMediaDataRemoved(KEY)
@@ -456,7 +477,7 @@
         }
 
     @Test
-    fun testOnSwipeToDismiss_setsTimedOut() {
+    fun onSwipeToDismiss_setsTimedOut() {
         mediaDataFilter.onMediaDataLoaded(KEY, null, dataMain)
         mediaDataFilter.onSwipeToDismiss()
 
@@ -464,15 +485,19 @@
     }
 
     @Test
-    fun testOnSmartspaceMediaDataLoaded_noMedia_activeValidRec_prioritizesSmartspace() =
+    fun onSmartspaceMediaDataLoaded_noMedia_activeValidRec_prioritizesSmartspace() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel =
+                SmartspaceMediaLoadingModel.Loaded(SMARTSPACE_KEY, isPrioritized = true)
 
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            verify(listener).onSmartspaceMediaDataLoaded(eq(SMARTSPACE_KEY), eq(true))
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -487,18 +512,19 @@
         }
 
     @Test
-    fun testOnSmartspaceMediaDataLoaded_noMedia_inactiveRec_showsNothing() =
+    fun onSmartspaceMediaDataLoaded_noMedia_inactiveRec_showsNothing() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
 
             whenever(smartspaceData.isActive).thenReturn(false)
 
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            verify(listener, never()).onMediaDataLoaded(any(), anyBoolean(), anyInt(), anyBoolean())
-            verify(listener, never()).onSmartspaceMediaDataLoaded(any(), anyBoolean())
+            assertThat(recommendationsLoadingState).isEqualTo(SmartspaceMediaLoadingModel.Unknown)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -513,17 +539,21 @@
         }
 
     @Test
-    fun testOnSmartspaceMediaDataLoaded_noRecentMedia_activeValidRec_prioritizesSmartspace() =
+    fun onSmartspaceMediaDataLoaded_noRecentMedia_activeValidRec_prioritizesSmartspace() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel =
+                SmartspaceMediaLoadingModel.Loaded(SMARTSPACE_KEY, isPrioritized = true)
             val dataOld = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
             mediaDataFilter.onMediaDataLoaded(KEY, null, dataOld)
             clock.advanceTime(MediaDataFilterImpl.SMARTSPACE_MAX_AGE + 100)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            verify(listener).onSmartspaceMediaDataLoaded(eq(SMARTSPACE_KEY), eq(true))
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -538,11 +568,13 @@
         }
 
     @Test
-    fun testOnSmartspaceMediaDataLoaded_noRecentMedia_inactiveRec_showsNothing() =
+    fun onSmartspaceMediaDataLoaded_noRecentMedia_inactiveRec_showsNothing() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
             whenever(smartspaceData.isActive).thenReturn(false)
 
             val dataOld = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
@@ -550,7 +582,7 @@
             clock.advanceTime(MediaDataFilterImpl.SMARTSPACE_MAX_AGE + 100)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            verify(listener, never()).onSmartspaceMediaDataLoaded(any(), anyBoolean())
+            assertThat(recommendationsLoadingState).isEqualTo(SmartspaceMediaLoadingModel.Unknown)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -565,27 +597,29 @@
         }
 
     @Test
-    fun testOnSmartspaceMediaDataLoaded_hasRecentMedia_inactiveRec_showsNothing() =
+    fun onSmartspaceMediaDataLoaded_hasRecentMedia_inactiveRec_showsNothing() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
 
             whenever(smartspaceData.isActive).thenReturn(false)
 
             // WHEN we have media that was recently played, but not currently active
             val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
             mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrent.instanceId), eq(true), eq(0), eq(false))
 
-            reset(listener)
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
+
             // AND we get a smartspace signal
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            // THEN we should tell listeners to treat the media as not active instead
-            verify(listener, never()).onMediaDataLoaded(any(), anyBoolean(), anyInt(), anyBoolean())
-            verify(listener, never()).onSmartspaceMediaDataLoaded(any(), anyBoolean())
+            // THEN we should treat the media as not active instead
+            assertThat(recommendationsLoadingState).isEqualTo(SmartspaceMediaLoadingModel.Unknown)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -600,27 +634,28 @@
         }
 
     @Test
-    fun testOnSmartspaceMediaDataLoaded_hasRecentMedia_activeInvalidRec_usesMedia() =
+    fun onSmartspaceMediaDataLoaded_hasRecentMedia_activeInvalidRec_usesMedia() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
             whenever(smartspaceData.isValid()).thenReturn(false)
 
             // WHEN we have media that was recently played, but not currently active
             val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
             mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrent.instanceId), eq(true), eq(0), eq(false))
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
 
             // AND we get a smartspace signal
             runCurrent()
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            // THEN we should tell listeners to treat the media as active instead
-            val dataCurrentAndActive = dataCurrent.copy(active = true)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrentAndActive.instanceId), eq(true), eq(100), eq(true))
+            // THEN we should treat the media as active instead
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -630,31 +665,35 @@
                 )
                 .isTrue()
             // Smartspace update shouldn't be propagated for the empty rec list.
-            verify(listener, never()).onSmartspaceMediaDataLoaded(any(), anyBoolean())
+            assertThat(recommendationsLoadingState).isEqualTo(SmartspaceMediaLoadingModel.Unknown)
             verify(logger, never()).logRecommendationAdded(any(), any())
             verify(logger).logRecommendationActivated(eq(APP_UID), eq(PACKAGE), eq(INSTANCE_ID))
         }
 
     @Test
-    fun testOnSmartspaceMediaDataLoaded_hasRecentMedia_activeValidRec_usesBoth() =
+    fun onSmartspaceMediaDataLoaded_hasRecentMedia_activeValidRec_usesBoth() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
             // WHEN we have media that was recently played, but not currently active
             val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
+            val mediaDataLoadingModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
+            val recommendationsLoadingModel = SmartspaceMediaLoadingModel.Loaded(SMARTSPACE_KEY)
+
             mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrent.instanceId), eq(true), eq(0), eq(false))
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaDataLoadingModel)
 
             // AND we get a smartspace signal
             runCurrent()
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            // THEN we should tell listeners to treat the media as active instead
-            val dataCurrentAndActive = dataCurrent.copy(active = true)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrentAndActive.instanceId), eq(true), eq(100), eq(true))
+            // THEN we should treat the media as active instead
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaDataLoadingModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -664,22 +703,25 @@
                 )
                 .isTrue()
             // Smartspace update should also be propagated but not prioritized.
-            verify(listener).onSmartspaceMediaDataLoaded(eq(SMARTSPACE_KEY), eq(false))
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
             verify(logger).logRecommendationAdded(SMARTSPACE_PACKAGE, SMARTSPACE_INSTANCE_ID)
             verify(logger).logRecommendationActivated(eq(APP_UID), eq(PACKAGE), eq(INSTANCE_ID))
         }
 
     @Test
-    fun testOnSmartspaceMediaDataRemoved_usedSmartspace_clearsSmartspace() =
+    fun onSmartspaceMediaDataRemoved_usedSmartspace_clearsSmartspace() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel = SmartspaceMediaLoadingModel.Removed(SMARTSPACE_KEY)
 
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
             mediaDataFilter.onSmartspaceMediaDataRemoved(SMARTSPACE_KEY)
 
-            verify(listener).onSmartspaceMediaDataRemoved(SMARTSPACE_KEY)
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -692,26 +734,28 @@
         }
 
     @Test
-    fun testOnSmartspaceMediaDataRemoved_usedMediaAndSmartspace_clearsBoth() =
+    fun onSmartspaceMediaDataRemoved_usedMediaAndSmartspace_clearsBoth() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel = SmartspaceMediaLoadingModel.Removed(SMARTSPACE_KEY)
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
             val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
             mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrent.instanceId), eq(true), eq(0), eq(false))
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
 
             runCurrent()
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            val dataCurrentAndActive = dataCurrent.copy(active = true)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrentAndActive.instanceId), eq(true), eq(100), eq(true))
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
 
             mediaDataFilter.onSmartspaceMediaDataRemoved(SMARTSPACE_KEY)
 
-            verify(listener).onSmartspaceMediaDataRemoved(SMARTSPACE_KEY)
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -724,17 +768,20 @@
         }
 
     @Test
-    fun testOnSmartspaceLoaded_persistentEnabled_isInactive_notifiesListeners() =
+    fun onSmartspaceLoaded_persistentEnabled_isInactive() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel = SmartspaceMediaLoadingModel.Loaded(SMARTSPACE_KEY)
             whenever(mediaFlags.isPersistentSsCardEnabled()).thenReturn(true)
             whenever(smartspaceData.isActive).thenReturn(false)
 
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            verify(listener).onSmartspaceMediaDataLoaded(eq(SMARTSPACE_KEY), eq(false))
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -748,11 +795,16 @@
         }
 
     @Test
-    fun testOnSmartspaceLoaded_persistentEnabled_inactive_hasRecentMedia_staysInactive() =
+    fun onSmartspaceLoaded_persistentEnabled_inactive_hasRecentMedia_staysInactive() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel = SmartspaceMediaLoadingModel.Loaded(SMARTSPACE_KEY)
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
 
             whenever(mediaFlags.isPersistentSsCardEnabled()).thenReturn(true)
             whenever(smartspaceData.isActive).thenReturn(false)
@@ -760,16 +812,14 @@
             // If there is media that was recently played but inactive
             val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
             mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrent.instanceId), eq(true), eq(0), eq(false))
 
-            reset(listener)
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
+
             // And an inactive recommendation is loaded
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
             // Smartspace is loaded but the media stays inactive
-            verify(listener).onSmartspaceMediaDataLoaded(eq(SMARTSPACE_KEY), eq(false))
-            verify(listener, never()).onMediaDataLoaded(any(), anyBoolean(), anyInt(), anyBoolean())
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -783,7 +833,7 @@
         }
 
     @Test
-    fun testOnSwipeToDismiss_persistentEnabled_recommendationSetInactive() {
+    fun onSwipeToDismiss_persistentEnabled_recommendationSetInactive() {
         whenever(mediaFlags.isPersistentSsCardEnabled()).thenReturn(true)
 
         val data =
@@ -802,16 +852,21 @@
     }
 
     @Test
-    fun testSmartspaceLoaded_shouldTriggerResume_doesTrigger() =
+    fun smartspaceLoaded_shouldTriggerResume_doesTrigger() =
         testScope.runTest {
-            val selectedUserEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
-            val smartspaceMediaData by collectLastValue(mediaFilterRepository.smartspaceMediaData)
-            val reactivatedKey by collectLastValue(mediaFilterRepository.reactivatedId)
+            val selectedUserEntries by collectLastValue(repository.selectedUserEntries)
+            val smartspaceMediaData by collectLastValue(repository.smartspaceMediaData)
+            val reactivatedKey by collectLastValue(repository.reactivatedId)
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel = SmartspaceMediaLoadingModel.Loaded(SMARTSPACE_KEY)
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
             // WHEN we have media that was recently played, but not currently active
             val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
             mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrent.instanceId), eq(true), eq(0), eq(false))
+
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
 
             // AND we get a smartspace signal with extra to trigger resume
             runCurrent()
@@ -819,10 +874,8 @@
             whenever(cardAction.extras).thenReturn(extras)
             mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
 
-            // THEN we should tell listeners to treat the media as active instead
-            val dataCurrentAndActive = dataCurrent.copy(active = true)
-            verify(listener)
-                .onMediaDataLoaded(eq(dataCurrentAndActive.instanceId), eq(true), eq(100), eq(true))
+            // THEN we should treat the media as active instead
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
             assertThat(
                     hasActiveMediaOrRecommendation(
                         selectedUserEntries,
@@ -831,27 +884,33 @@
                     )
                 )
                 .isTrue()
-            // And send the smartspace data, but not prioritized
-            verify(listener).onSmartspaceMediaDataLoaded(eq(SMARTSPACE_KEY), eq(false))
+            // And update the smartspace data state, but not prioritized
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
         }
 
     @Test
-    fun testSmartspaceLoaded_notShouldTriggerResume_doesNotTrigger() {
-        // WHEN we have media that was recently played, but not currently active
-        val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
-        mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
-        verify(listener).onMediaDataLoaded(eq(dataCurrent.instanceId), eq(true), eq(0), eq(false))
+    fun smartspaceLoaded_notShouldTriggerResume_doesNotTrigger() =
+        testScope.runTest {
+            val mediaDataLoadedStates by collectLastValue(repository.mediaDataLoadedStates)
+            val recommendationsLoadingState by
+                collectLastValue(repository.recommendationsLoadingState)
+            val recommendationsLoadingModel = SmartspaceMediaLoadingModel.Loaded(SMARTSPACE_KEY)
+            val mediaLoadedStatesModel = listOf(MediaDataLoadingModel.Loaded(dataMain.instanceId))
 
-        // AND we get a smartspace signal with extra to not trigger resume
-        val extras = Bundle().apply { putBoolean(EXTRA_KEY_TRIGGER_RESUME, false) }
-        whenever(cardAction.extras).thenReturn(extras)
-        mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
+            // WHEN we have media that was recently played, but not currently active
+            val dataCurrent = dataMain.copy(active = false, lastActive = clock.elapsedRealtime())
+            mediaDataFilter.onMediaDataLoaded(KEY, null, dataCurrent)
 
-        // THEN listeners are not updated to show media
-        verify(listener, never()).onMediaDataLoaded(any(), eq(true), eq(100), eq(true))
-        // But the smartspace update is still propagated
-        verify(listener).onSmartspaceMediaDataLoaded(eq(SMARTSPACE_KEY), eq(false))
-    }
+            assertThat(mediaDataLoadedStates).isEqualTo(mediaLoadedStatesModel)
+
+            // AND we get a smartspace signal with extra to not trigger resume
+            val extras = Bundle().apply { putBoolean(EXTRA_KEY_TRIGGER_RESUME, false) }
+            whenever(cardAction.extras).thenReturn(extras)
+            mediaDataFilter.onSmartspaceMediaDataLoaded(SMARTSPACE_KEY, smartspaceData)
+
+            // But the smartspace update is still propagated
+            assertThat(recommendationsLoadingState).isEqualTo(recommendationsLoadingModel)
+        }
 
     private fun hasActiveMediaOrRecommendation(
         entries: Map<InstanceId, MediaData>?,