Retain media controls longer

When the flag is enabled, apps which have not added support for
resumption via MediaBrowserService will still be set to a resume state when
the underlying notification (for apps using notification-based actions) or
session (for apps using PlaybackState actions) have gone away, if this
occurs after they've already timed out. These resume controls will not
have a play button, but just open the app's launcher intent when tapped.

If the control hadn't timed out yet, it will be removed from the carousel as
before. For apps which do implement long-term resumption, they will
continue to convert into resume state regardless of whether the active
player had already timed out.

Test: atest com.android.systemui.media.controls
Bug: 264691138
Change-Id: I13a43bfc2eca342000726e696afe77706947d774

6 files changed, 320 insertions, 48 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/flags/Flags.kt b/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
index fbe0cc99fc01..7779351e26e1 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
@@ -349,6 +349,9 @@ object Flags {
     val MEDIA_TAP_TO_TRANSFER_DISMISS_GESTURE =
         unreleasedFlag(912, "media_ttt_dismiss_gesture", teamfood = true)
 
+    // TODO(b/266157412): Tracking Bug
+    val MEDIA_RETAIN_SESSIONS = unreleasedFlag(913, "media_retain_sessions")
+
     // 1000 - dock
     val SIMULATE_DOCK_THROUGH_CHARGING = releasedFlag(1000, "simulate_dock_through_charging")
 
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 a13279717d05..b11f628623bb 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
@@ -303,6 +303,7 @@ class MediaDataManager(
         mediaTimeoutListener.stateCallback = { key: String, state: PlaybackState ->
             updateState(key, state)
         }
+        mediaTimeoutListener.sessionCallback = { key: String -> onSessionDestroyed(key) }
         mediaResumeListener.setManager(this)
         mediaDataFilter.mediaDataManager = this
 
@@ -1289,43 +1290,104 @@ class MediaDataManager(
 
     fun onNotificationRemoved(key: String) {
         Assert.isMainThread()
-        val removed = mediaEntries.remove(key)
-        if (useMediaResumption && removed?.resumeAction != null && removed.isLocalSession()) {
-            Log.d(TAG, "Not removing $key because resumable")
-            // Move to resume key (aka package name) if that key doesn't already exist.
-            val resumeAction = getResumeMediaAction(removed.resumeAction!!)
-            val updated =
-                removed.copy(
-                    token = null,
-                    actions = listOf(resumeAction),
-                    semanticActions = MediaButton(playOrPause = resumeAction),
-                    actionsToShowInCompact = listOf(0),
-                    active = false,
-                    resumption = true,
-                    isPlaying = false,
-                    isClearable = true
+        val removed = mediaEntries.remove(key) ?: return
+
+        if (useMediaResumption && removed.resumeAction != null && removed.isLocalSession()) {
+            convertToResumePlayer(removed)
+        } else if (mediaFlags.isRetainingPlayersEnabled()) {
+            handlePossibleRemoval(removed, notificationRemoved = true)
+        } else {
+            notifyMediaDataRemoved(key)
+            logger.logMediaRemoved(removed.appUid, removed.packageName, removed.instanceId)
+        }
+    }
+
+    private fun onSessionDestroyed(key: String) {
+        if (!mediaFlags.isRetainingPlayersEnabled()) return
+
+        if (DEBUG) Log.d(TAG, "session destroyed for $key")
+        val entry = mediaEntries.remove(key) ?: return
+        // Clear token since the session is no longer valid
+        val updated = entry.copy(token = null)
+        handlePossibleRemoval(updated)
+    }
+
+    /**
+     * Convert to resume state if the player is no longer valid and active, then notify listeners
+     * that the data was updated. Does not convert to resume state if the player is still valid, or
+     * if it was removed before becoming inactive. (Assumes that [removed] was removed from
+     * [mediaEntries] before this function was called)
+     */
+    private fun handlePossibleRemoval(removed: MediaData, notificationRemoved: Boolean = false) {
+        val key = removed.notificationKey!!
+        val hasSession = removed.token != null
+        if (hasSession && removed.semanticActions != null) {
+            // The app was using session actions, and the session is still valid: keep player
+            if (DEBUG) Log.d(TAG, "Notification removed but using session actions $key")
+            mediaEntries.put(key, removed)
+            notifyMediaDataLoaded(key, key, removed)
+        } else if (!notificationRemoved && removed.semanticActions == null) {
+            // The app was using notification actions, and notif wasn't removed yet: keep player
+            if (DEBUG) Log.d(TAG, "Session destroyed but using notification actions $key")
+            mediaEntries.put(key, removed)
+            notifyMediaDataLoaded(key, key, removed)
+        } else if (removed.active) {
+            // This player was still active - it didn't last long enough to time out: remove
+            if (DEBUG) Log.d(TAG, "Removing still-active player $key")
+            notifyMediaDataRemoved(key)
+            logger.logMediaRemoved(removed.appUid, removed.packageName, removed.instanceId)
+        } else {
+            // Convert to resume
+            if (DEBUG) {
+                Log.d(
+                    TAG,
+                    "Notification ($notificationRemoved) and/or session " +
+                        "($hasSession) gone for inactive player $key"
                 )
-            val pkg = removed.packageName
-            val migrate = mediaEntries.put(pkg, updated) == null
-            // Notify listeners of "new" controls when migrating or removed and update when not
-            if (migrate) {
-                notifyMediaDataLoaded(pkg, key, updated)
-            } else {
-                // Since packageName is used for the key of the resumption controls, it is
-                // possible that another notification has already been reused for the resumption
-                // controls of this package. In this case, rather than renaming this player as
-                // packageName, just remove it and then send a update to the existing resumption
-                // controls.
-                notifyMediaDataRemoved(key)
-                notifyMediaDataLoaded(pkg, pkg, updated)
             }
-            logger.logActiveConvertedToResume(updated.appUid, pkg, updated.instanceId)
-            return
+            convertToResumePlayer(removed)
         }
-        if (removed != null) {
+    }
+
+    /** Set the given [MediaData] as a resume state player and notify listeners */
+    private fun convertToResumePlayer(data: MediaData) {
+        val key = data.notificationKey!!
+        if (DEBUG) Log.d(TAG, "Converting $key to resume")
+        // Move to resume key (aka package name) if that key doesn't already exist.
+        val resumeAction = data.resumeAction?.let { getResumeMediaAction(it) }
+        val actions = resumeAction?.let { listOf(resumeAction) } ?: emptyList()
+        val launcherIntent =
+            context.packageManager.getLaunchIntentForPackage(data.packageName)?.let {
+                PendingIntent.getActivity(context, 0, it, PendingIntent.FLAG_IMMUTABLE)
+            }
+        val updated =
+            data.copy(
+                token = null,
+                actions = actions,
+                semanticActions = MediaButton(playOrPause = resumeAction),
+                actionsToShowInCompact = listOf(0),
+                active = false,
+                resumption = true,
+                isPlaying = false,
+                isClearable = true,
+                clickIntent = launcherIntent,
+            )
+        val pkg = data.packageName
+        val migrate = mediaEntries.put(pkg, updated) == null
+        // Notify listeners of "new" controls when migrating or removed and update when not
+        Log.d(TAG, "migrating? $migrate from $key -> $pkg")
+        if (migrate) {
+            notifyMediaDataLoaded(key = pkg, oldKey = key, info = updated)
+        } else {
+            // Since packageName is used for the key of the resumption controls, it is
+            // possible that another notification has already been reused for the resumption
+            // controls of this package. In this case, rather than renaming this player as
+            // packageName, just remove it and then send a update to the existing resumption
+            // controls.
             notifyMediaDataRemoved(key)
-            logger.logMediaRemoved(removed.appUid, removed.packageName, removed.instanceId)
+            notifyMediaDataLoaded(key = pkg, oldKey = pkg, info = updated)
         }
+        logger.logActiveConvertedToResume(updated.appUid, pkg, updated.instanceId)
     }
 
     fun setMediaResumptionEnabled(isEnabled: Boolean) {
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/pipeline/MediaTimeoutListener.kt b/packages/SystemUI/src/com/android/systemui/media/controls/pipeline/MediaTimeoutListener.kt
index 7f5c82fb5eee..a898b00790a9 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/pipeline/MediaTimeoutListener.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/pipeline/MediaTimeoutListener.kt
@@ -71,6 +71,12 @@ constructor(
      */
     lateinit var stateCallback: (String, PlaybackState) -> Unit
 
+    /**
+     * Callback representing that the [MediaSession] for an active control has been destroyed
+     * @param key Media control unique identifier
+     */
+    lateinit var sessionCallback: (String) -> Unit
+
     init {
         statusBarStateController.addCallback(
             object : StatusBarStateController.StateListener {
@@ -211,6 +217,7 @@ constructor(
             } else {
                 // For active controls, if the session is destroyed, clean up everything since we
                 // will need to recreate it if this key is updated later
+                sessionCallback.invoke(key)
                 destroy()
             }
         }
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/util/MediaFlags.kt b/packages/SystemUI/src/com/android/systemui/media/controls/util/MediaFlags.kt
index 5bc35caed515..ab03930e42ac 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/util/MediaFlags.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/util/MediaFlags.kt
@@ -45,4 +45,10 @@ class MediaFlags @Inject constructor(private val featureFlags: FeatureFlags) {
 
     /** Check whether we show explicit indicator on UMO */
     fun isExplicitIndicatorEnabled() = featureFlags.isEnabled(Flags.MEDIA_EXPLICIT_INDICATOR)
+
+    /**
+     * If true, keep active media controls for the lifetime of the MediaSession, regardless of
+     * whether the underlying notification was dismissed
+     */
+    fun isRetainingPlayersEnabled() = featureFlags.isEnabled(Flags.MEDIA_RETAIN_SESSIONS)
 }
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 1687fdc9f76c..1ac66952fd3f 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
@@ -134,7 +134,8 @@ class MediaDataManagerTest : SysuiTestCase() {
     private val clock = FakeSystemClock()
     @Mock private lateinit var tunerService: TunerService
     @Captor lateinit var tunableCaptor: ArgumentCaptor<TunerService.Tunable>
-    @Captor lateinit var callbackCaptor: ArgumentCaptor<(String, PlaybackState) -> Unit>
+    @Captor lateinit var stateCallbackCaptor: ArgumentCaptor<(String, PlaybackState) -> Unit>
+    @Captor lateinit var sessionCallbackCaptor: ArgumentCaptor<(String) -> Unit>
     @Captor lateinit var smartSpaceConfigBuilderCaptor: ArgumentCaptor<SmartspaceConfig>
 
     private val instanceIdSequence = InstanceIdSequenceFake(1 shl 20)
@@ -184,6 +185,8 @@ class MediaDataManagerTest : SysuiTestCase() {
             )
         verify(tunerService)
             .addTunable(capture(tunableCaptor), eq(Settings.Secure.MEDIA_CONTROLS_RECOMMENDATION))
+        verify(mediaTimeoutListener).stateCallback = capture(stateCallbackCaptor)
+        verify(mediaTimeoutListener).sessionCallback = capture(sessionCallbackCaptor)
         session = MediaSession(context, "MediaDataManagerTestSession")
         mediaNotification =
             SbnBuilder().run {
@@ -230,6 +233,7 @@ class MediaDataManagerTest : SysuiTestCase() {
         whenever(mediaSmartspaceTarget.creationTimeMillis).thenReturn(1234L)
         whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(false)
         whenever(mediaFlags.isExplicitIndicatorEnabled()).thenReturn(true)
+        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(false)
         whenever(logger.getNewInstanceId()).thenReturn(instanceIdSequence.newInstanceId())
     }
 
@@ -547,6 +551,7 @@ class MediaDataManagerTest : SysuiTestCase() {
         mediaDataManager.onNotificationAdded(KEY_2, mediaNotification)
         assertThat(backgroundExecutor.runAllReady()).isEqualTo(2)
         assertThat(foregroundExecutor.runAllReady()).isEqualTo(2)
+
         verify(listener)
             .onMediaDataLoaded(
                 eq(KEY),
@@ -558,9 +563,21 @@ class MediaDataManagerTest : SysuiTestCase() {
             )
         val data = mediaDataCaptor.value
         assertThat(data.resumption).isFalse()
-        val resumableData = data.copy(resumeAction = Runnable {})
-        mediaDataManager.onMediaDataLoaded(KEY, null, resumableData)
-        mediaDataManager.onMediaDataLoaded(KEY_2, null, resumableData)
+
+        verify(listener)
+            .onMediaDataLoaded(
+                eq(KEY_2),
+                eq(null),
+                capture(mediaDataCaptor),
+                eq(true),
+                eq(0),
+                eq(false)
+            )
+        val data2 = mediaDataCaptor.value
+        assertThat(data2.resumption).isFalse()
+
+        mediaDataManager.onMediaDataLoaded(KEY, null, data.copy(resumeAction = Runnable {}))
+        mediaDataManager.onMediaDataLoaded(KEY_2, null, data2.copy(resumeAction = Runnable {}))
         reset(listener)
         // WHEN the first is removed
         mediaDataManager.onNotificationRemoved(KEY)
@@ -1310,11 +1327,10 @@ class MediaDataManagerTest : SysuiTestCase() {
     fun testPlaybackStateChange_keyExists_callsListener() {
         // Notification has been added
         addNotificationAndLoad()
-        verify(mediaTimeoutListener).stateCallback = capture(callbackCaptor)
 
         // Callback gets an updated state
         val state = PlaybackState.Builder().setState(PlaybackState.STATE_PLAYING, 0L, 1f).build()
-        callbackCaptor.value.invoke(KEY, state)
+        stateCallbackCaptor.value.invoke(KEY, state)
 
         // Listener is notified of updated state
         verify(listener)
@@ -1332,11 +1348,10 @@ class MediaDataManagerTest : SysuiTestCase() {
     @Test
     fun testPlaybackStateChange_keyDoesNotExist_doesNothing() {
         val state = PlaybackState.Builder().build()
-        verify(mediaTimeoutListener).stateCallback = capture(callbackCaptor)
 
         // No media added with this key
 
-        callbackCaptor.value.invoke(KEY, state)
+        stateCallbackCaptor.value.invoke(KEY, state)
         verify(listener, never())
             .onMediaDataLoaded(eq(KEY), any(), any(), anyBoolean(), anyInt(), anyBoolean())
     }
@@ -1352,10 +1367,9 @@ class MediaDataManagerTest : SysuiTestCase() {
 
         // And then get a state update
         val state = PlaybackState.Builder().build()
-        verify(mediaTimeoutListener).stateCallback = capture(callbackCaptor)
 
         // Then no changes are made
-        callbackCaptor.value.invoke(KEY, state)
+        stateCallbackCaptor.value.invoke(KEY, state)
         verify(listener, never())
             .onMediaDataLoaded(eq(KEY), any(), any(), anyBoolean(), anyInt(), anyBoolean())
     }
@@ -1367,8 +1381,7 @@ class MediaDataManagerTest : SysuiTestCase() {
         whenever(controller.playbackState).thenReturn(state)
 
         addNotificationAndLoad()
-        verify(mediaTimeoutListener).stateCallback = capture(callbackCaptor)
-        callbackCaptor.value.invoke(KEY, state)
+        stateCallbackCaptor.value.invoke(KEY, state)
 
         verify(listener)
             .onMediaDataLoaded(
@@ -1410,8 +1423,7 @@ class MediaDataManagerTest : SysuiTestCase() {
         backgroundExecutor.runAllReady()
         foregroundExecutor.runAllReady()
 
-        verify(mediaTimeoutListener).stateCallback = capture(callbackCaptor)
-        callbackCaptor.value.invoke(PACKAGE_NAME, state)
+        stateCallbackCaptor.value.invoke(PACKAGE_NAME, state)
 
         verify(listener)
             .onMediaDataLoaded(
@@ -1436,8 +1448,7 @@ class MediaDataManagerTest : SysuiTestCase() {
                 .build()
 
         addNotificationAndLoad()
-        verify(mediaTimeoutListener).stateCallback = capture(callbackCaptor)
-        callbackCaptor.value.invoke(KEY, state)
+        stateCallbackCaptor.value.invoke(KEY, state)
 
         verify(listener)
             .onMediaDataLoaded(
@@ -1485,6 +1496,177 @@ class MediaDataManagerTest : SysuiTestCase() {
         assertThat(mediaDataCaptor.value.isClearable).isFalse()
     }
 
+    @Test
+    fun testRetain_notifPlayer_notifRemoved_setToResume() {
+        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
+
+        // When a media control based on notification is added, times out, and then removed
+        addNotificationAndLoad()
+        mediaDataManager.setTimedOut(KEY, timedOut = true)
+        assertThat(mediaDataCaptor.value.active).isFalse()
+        mediaDataManager.onNotificationRemoved(KEY)
+
+        // It is converted to a resume player
+        verify(listener)
+            .onMediaDataLoaded(
+                eq(PACKAGE_NAME),
+                eq(KEY),
+                capture(mediaDataCaptor),
+                eq(true),
+                eq(0),
+                eq(false)
+            )
+        assertThat(mediaDataCaptor.value.resumption).isTrue()
+        assertThat(mediaDataCaptor.value.active).isFalse()
+        verify(logger)
+            .logActiveConvertedToResume(
+                anyInt(),
+                eq(PACKAGE_NAME),
+                eq(mediaDataCaptor.value.instanceId)
+            )
+    }
+
+    @Test
+    fun testRetain_notifPlayer_sessionDestroyed_doesNotChange() {
+        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
+
+        // When a media control based on notification is added and times out
+        addNotificationAndLoad()
+        mediaDataManager.setTimedOut(KEY, timedOut = true)
+        assertThat(mediaDataCaptor.value.active).isFalse()
+
+        // and then the session is destroyed
+        sessionCallbackCaptor.value.invoke(KEY)
+
+        // It remains as a regular player
+        verify(listener, never()).onMediaDataRemoved(eq(KEY))
+        verify(listener, never())
+            .onMediaDataLoaded(eq(PACKAGE_NAME), any(), any(), anyBoolean(), anyInt(), anyBoolean())
+    }
+
+    @Test
+    fun testRetain_notifPlayer_removeWhileActive_fullyRemoved() {
+        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
+
+        // When a media control based on notification is added and then removed, without timing out
+        addNotificationAndLoad()
+        val data = mediaDataCaptor.value
+        assertThat(data.active).isTrue()
+        mediaDataManager.onNotificationRemoved(KEY)
+
+        // It is fully removed
+        verify(listener).onMediaDataRemoved(eq(KEY))
+        verify(logger).logMediaRemoved(anyInt(), eq(PACKAGE_NAME), eq(data.instanceId))
+        verify(listener, never())
+            .onMediaDataLoaded(eq(PACKAGE_NAME), any(), any(), anyBoolean(), anyInt(), anyBoolean())
+    }
+
+    @Test
+    fun testRetain_canResume_removeWhileActive_setToResume() {
+        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
+
+        // When a media control that supports resumption is added
+        addNotificationAndLoad()
+        val dataResumable = mediaDataCaptor.value.copy(resumeAction = Runnable {})
+        mediaDataManager.onMediaDataLoaded(KEY, null, dataResumable)
+
+        // And then removed while still active
+        mediaDataManager.onNotificationRemoved(KEY)
+
+        // It is converted to a resume player
+        verify(listener)
+            .onMediaDataLoaded(
+                eq(PACKAGE_NAME),
+                eq(KEY),
+                capture(mediaDataCaptor),
+                eq(true),
+                eq(0),
+                eq(false)
+            )
+        assertThat(mediaDataCaptor.value.resumption).isTrue()
+        assertThat(mediaDataCaptor.value.active).isFalse()
+        verify(logger)
+            .logActiveConvertedToResume(
+                anyInt(),
+                eq(PACKAGE_NAME),
+                eq(mediaDataCaptor.value.instanceId)
+            )
+    }
+
+    @Test
+    fun testRetain_sessionPlayer_notifRemoved_doesNotChange() {
+        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
+        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        addPlaybackStateAction()
+
+        // When a media control with PlaybackState actions is added, times out,
+        // and then the notification is removed
+        addNotificationAndLoad()
+        val data = mediaDataCaptor.value
+        assertThat(data.active).isTrue()
+        mediaDataManager.setTimedOut(KEY, timedOut = true)
+        mediaDataManager.onNotificationRemoved(KEY)
+
+        // It remains as a regular player
+        verify(listener, never()).onMediaDataRemoved(eq(KEY))
+        verify(listener, never())
+            .onMediaDataLoaded(eq(PACKAGE_NAME), any(), any(), anyBoolean(), anyInt(), anyBoolean())
+    }
+
+    @Test
+    fun testRetain_sessionPlayer_sessionDestroyed_setToResume() {
+        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
+        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        addPlaybackStateAction()
+
+        // When a media control with PlaybackState actions is added, times out,
+        // and then the session is destroyed
+        addNotificationAndLoad()
+        val data = mediaDataCaptor.value
+        assertThat(data.active).isTrue()
+        mediaDataManager.setTimedOut(KEY, timedOut = true)
+        sessionCallbackCaptor.value.invoke(KEY)
+
+        // It is converted to a resume player
+        verify(listener)
+            .onMediaDataLoaded(
+                eq(PACKAGE_NAME),
+                eq(KEY),
+                capture(mediaDataCaptor),
+                eq(true),
+                eq(0),
+                eq(false)
+            )
+        assertThat(mediaDataCaptor.value.resumption).isTrue()
+        assertThat(mediaDataCaptor.value.active).isFalse()
+        verify(logger)
+            .logActiveConvertedToResume(
+                anyInt(),
+                eq(PACKAGE_NAME),
+                eq(mediaDataCaptor.value.instanceId)
+            )
+    }
+
+    @Test
+    fun testRetain_sessionPlayer_destroyedWhileActive_fullyRemoved() {
+        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
+        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        addPlaybackStateAction()
+
+        // When a media control using session actions is added, and then the session is destroyed
+        // without timing out first
+        addNotificationAndLoad()
+        val data = mediaDataCaptor.value
+        assertThat(data.active).isTrue()
+        sessionCallbackCaptor.value.invoke(KEY)
+
+        // It is fully removed
+        verify(listener).onMediaDataRemoved(eq(KEY))
+        verify(logger).logMediaRemoved(anyInt(), eq(PACKAGE_NAME), eq(data.instanceId))
+        verify(listener, never())
+            .onMediaDataLoaded(eq(PACKAGE_NAME), any(), any(), anyBoolean(), anyInt(), anyBoolean())
+    }
+
     /** Helper function to add a media notification and capture the resulting MediaData */
     private fun addNotificationAndLoad() {
         mediaDataManager.onNotificationAdded(KEY, mediaNotification)
@@ -1500,4 +1682,12 @@ class MediaDataManagerTest : SysuiTestCase() {
                 eq(false)
             )
     }
+
+    /** Helper function to set up a PlaybackState with action */
+    private fun addPlaybackStateAction() {
+        val stateActions = PlaybackState.ACTION_PLAY_PAUSE
+        val stateBuilder = PlaybackState.Builder().setActions(stateActions)
+        stateBuilder.setState(PlaybackState.STATE_PAUSED, 0, 1.0f)
+        whenever(controller.playbackState).thenReturn(stateBuilder.build())
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/controls/pipeline/MediaTimeoutListenerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/controls/pipeline/MediaTimeoutListenerTest.kt
index 344dffafb448..92bf84ce285c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/controls/pipeline/MediaTimeoutListenerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/controls/pipeline/MediaTimeoutListenerTest.kt
@@ -72,6 +72,7 @@ class MediaTimeoutListenerTest : SysuiTestCase() {
     private lateinit var executor: FakeExecutor
     @Mock private lateinit var timeoutCallback: (String, Boolean) -> Unit
     @Mock private lateinit var stateCallback: (String, PlaybackState) -> Unit
+    @Mock private lateinit var sessionCallback: (String) -> Unit
     @Captor private lateinit var mediaCallbackCaptor: ArgumentCaptor<MediaController.Callback>
     @Captor
     private lateinit var dozingCallbackCaptor:
@@ -99,6 +100,7 @@ class MediaTimeoutListenerTest : SysuiTestCase() {
             )
         mediaTimeoutListener.timeoutCallback = timeoutCallback
         mediaTimeoutListener.stateCallback = stateCallback
+        mediaTimeoutListener.sessionCallback = sessionCallback
 
         // Create a media session and notification for testing.
         metadataBuilder =
@@ -284,6 +286,7 @@ class MediaTimeoutListenerTest : SysuiTestCase() {
         verify(mediaController).unregisterCallback(anyObject())
         assertThat(executor.numPending()).isEqualTo(0)
         verify(logger).logSessionDestroyed(eq(KEY))
+        verify(sessionCallback).invoke(eq(KEY))
     }
 
     @Test
@@ -322,6 +325,7 @@ class MediaTimeoutListenerTest : SysuiTestCase() {
         // THEN the controller is unregistered, but the timeout is still scheduled
         verify(mediaController).unregisterCallback(anyObject())
         assertThat(executor.numPending()).isEqualTo(1)
+        verify(sessionCallback, never()).invoke(eq(KEY))
     }
 
     @Test