Update media3 button tests and error handling

- Use a fake Media3ActionFactory in kosmos
- Wrap media3 handler calls in a try-finally, to ensure that the
  controller disconnects
- Minor edits from code review

Flag: com.android.systemui.media_controls_button_media3
Bug: 360196209
Test: atest Media3ActionFactoryTest MediaDataLoaderTest
Change-Id: Ic91401af5b030da5bfa4ff188d1b55b69b72be5f

5 files changed, 97 insertions, 43 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/pipeline/Media3ActionFactoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/pipeline/Media3ActionFactoryTest.kt
index 9e3fdf377b83..8e67e602abd9 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/pipeline/Media3ActionFactoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/pipeline/Media3ActionFactoryTest.kt
@@ -20,7 +20,6 @@ import android.media.session.MediaSession
 import android.os.Bundle
 import android.os.Handler
 import android.os.looper
-import android.testing.TestableLooper
 import android.testing.TestableLooper.RunWithLooper
 import androidx.media.utils.MediaConstants
 import androidx.media3.common.Player
@@ -69,10 +68,9 @@ class Media3ActionFactoryTest : SysuiTestCase() {
     private val testScope = kosmos.testScope
     private val controllerFactory = kosmos.fakeMediaControllerFactory
     private val tokenFactory = kosmos.fakeSessionTokenFactory
-    private lateinit var testableLooper: TestableLooper
 
-    private var commandCaptor = argumentCaptor<SessionCommand>()
-    private var runnableCaptor = argumentCaptor<Runnable>()
+    private val commandCaptor = argumentCaptor<SessionCommand>()
+    private val runnableCaptor = argumentCaptor<Runnable>()
 
     private val legacyToken = MediaSession.Token(1, null)
     private val token = mock<SessionToken>()
@@ -97,8 +95,6 @@ class Media3ActionFactoryTest : SysuiTestCase() {
 
     @Before
     fun setup() {
-        testableLooper = TestableLooper.get(this)
-
         underTest =
             Media3ActionFactory(
                 context,
@@ -246,7 +242,6 @@ class Media3ActionFactoryTest : SysuiTestCase() {
             assertThat(actions.custom0!!.contentDescription).isEqualTo("$CUSTOM_ACTION_NAME 2")
             actions.custom0!!.action!!.run()
             runCurrent()
-            testableLooper.processAllMessages()
             verify(media3Controller).sendCustomCommand(commandCaptor.capture(), any<Bundle>())
             assertThat(commandCaptor.lastValue.customAction).isEqualTo("$CUSTOM_ACTION_COMMAND 2")
             verify(media3Controller).release()
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataLoaderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataLoaderTest.kt
index 1a7265b09aae..cf503bbf1310 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataLoaderTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataLoaderTest.kt
@@ -82,7 +82,7 @@ class MediaDataLoaderTest : SysuiTestCase() {
     private val fakeFeatureFlags = kosmos.fakeFeatureFlagsClassic
     private val mediaFlags = kosmos.mediaFlags
     private val mediaControllerFactory = kosmos.fakeMediaControllerFactory
-    private val media3ActionFactory = kosmos.media3ActionFactory
+    private lateinit var media3ActionFactory: Media3ActionFactory
     private val session = MediaSession(context, "MediaDataLoaderTestSession")
     private val metadataBuilder =
         MediaMetadata.Builder().apply {
@@ -94,6 +94,7 @@ class MediaDataLoaderTest : SysuiTestCase() {
 
     @Before
     fun setUp() {
+        media3ActionFactory = kosmos.media3ActionFactory
         mediaControllerFactory.setControllerForToken(session.sessionToken, mediaController)
         whenever(mediaController.sessionToken).thenReturn(session.sessionToken)
         whenever(mediaController.metadata).then { metadataBuilder.build() }
@@ -311,7 +312,6 @@ class MediaDataLoaderTest : SysuiTestCase() {
                     }
                     build()
                 }
-
             val result = underTest.loadMediaData(KEY, mediaNotification)
 
             assertThat(result).isNotNull()
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/Media3ActionFactory.kt b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/Media3ActionFactory.kt
index a33685b61237..d38e507082b9 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/Media3ActionFactory.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/Media3ActionFactory.kt
@@ -28,6 +28,7 @@ import androidx.annotation.WorkerThread
 import androidx.media.utils.MediaConstants
 import androidx.media3.common.Player
 import androidx.media3.session.CommandButton
+import androidx.media3.session.MediaController as Media3Controller
 import androidx.media3.session.SessionCommand
 import androidx.media3.session.SessionToken
 import com.android.systemui.dagger.SysUISingleton
@@ -82,10 +83,19 @@ constructor(
         // Build button info
         val buttons = suspendCancellableCoroutine { continuation ->
             // Media3Controller methods must always be called from a specific looper
-            handler.post {
-                val result = getMedia3Actions(packageName, m3controller, token)
-                m3controller.release()
-                continuation.resumeWith(Result.success(result))
+            val runnable = Runnable {
+                try {
+                    val result = getMedia3Actions(packageName, m3controller, token)
+                    continuation.resumeWith(Result.success(result))
+                } finally {
+                    m3controller.release()
+                }
+            }
+            handler.post(runnable)
+            continuation.invokeOnCancellation {
+                // Ensure controller is released, even if loading was cancelled partway through
+                handler.post(m3controller::release)
+                handler.removeCallbacks(runnable)
             }
         }
         return buttons
@@ -95,7 +105,7 @@ constructor(
     @WorkerThread
     private fun getMedia3Actions(
         packageName: String,
-        m3controller: androidx.media3.session.MediaController,
+        m3controller: Media3Controller,
         token: SessionToken,
     ): MediaButton? {
         Assert.isNotMainThread()
@@ -197,7 +207,7 @@ constructor(
      * @return A [MediaAction] representing the first supported command, or null if not supported
      */
     private fun getStandardAction(
-        controller: androidx.media3.session.MediaController,
+        controller: Media3Controller,
         token: SessionToken,
         vararg commands: @Player.Command Int,
     ): MediaAction? {
@@ -304,37 +314,40 @@ constructor(
         bgScope.launch {
             val controller = controllerFactory.create(token, looper)
             handler.post {
-                when (command) {
-                    Player.COMMAND_PLAY_PAUSE -> {
-                        if (controller.isPlaying) controller.pause() else controller.play()
-                    }
+                try {
+                    when (command) {
+                        Player.COMMAND_PLAY_PAUSE -> {
+                            if (controller.isPlaying) controller.pause() else controller.play()
+                        }
 
-                    Player.COMMAND_SEEK_TO_PREVIOUS -> controller.seekToPrevious()
-                    Player.COMMAND_SEEK_TO_PREVIOUS_MEDIA_ITEM ->
-                        controller.seekToPreviousMediaItem()
+                        Player.COMMAND_SEEK_TO_PREVIOUS -> controller.seekToPrevious()
+                        Player.COMMAND_SEEK_TO_PREVIOUS_MEDIA_ITEM ->
+                            controller.seekToPreviousMediaItem()
 
-                    Player.COMMAND_SEEK_TO_NEXT -> controller.seekToNext()
-                    Player.COMMAND_SEEK_TO_NEXT_MEDIA_ITEM -> controller.seekToNextMediaItem()
-                    Player.COMMAND_INVALID -> {
-                        if (
-                            customAction != null &&
-                                customAction!!.sessionCommand != null &&
-                                controller.isSessionCommandAvailable(
-                                    customAction!!.sessionCommand!!
-                                )
-                        ) {
-                            controller.sendCustomCommand(
-                                customAction!!.sessionCommand!!,
-                                customAction!!.extras,
-                            )
-                        } else {
-                            logger.logMedia3UnsupportedCommand("$command, action $customAction")
+                        Player.COMMAND_SEEK_TO_NEXT -> controller.seekToNext()
+                        Player.COMMAND_SEEK_TO_NEXT_MEDIA_ITEM -> controller.seekToNextMediaItem()
+                        Player.COMMAND_INVALID -> {
+                            if (customAction?.sessionCommand != null) {
+                                val sessionCommand = customAction.sessionCommand!!
+                                if (controller.isSessionCommandAvailable(sessionCommand)) {
+                                    controller.sendCustomCommand(
+                                        sessionCommand,
+                                        customAction.extras,
+                                    )
+                                } else {
+                                    logger.logMedia3UnsupportedCommand(
+                                        "$sessionCommand, action $customAction"
+                                    )
+                                }
+                            } else {
+                                logger.logMedia3UnsupportedCommand("$command, action $customAction")
+                            }
                         }
+                        else -> logger.logMedia3UnsupportedCommand(command.toString())
                     }
-
-                    else -> logger.logMedia3UnsupportedCommand(command.toString())
+                } finally {
+                    controller.release()
                 }
-                controller.release()
             }
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/util/MediaControllerFactory.kt b/packages/SystemUI/src/com/android/systemui/media/controls/util/MediaControllerFactory.kt
index 741f52998782..d815852b790f 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/util/MediaControllerFactory.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/util/MediaControllerFactory.kt
@@ -35,7 +35,12 @@ open class MediaControllerFactory @Inject constructor(private val context: Conte
         return MediaController(context, token)
     }
 
-    /** Creates a new [Media3Controller] from a [SessionToken] */
+    /**
+     * Creates a new [Media3Controller] from the media3 [SessionToken].
+     *
+     * @param token The token for the session
+     * @param looper The looper that will be used for this controller's operations
+     */
     open suspend fun create(token: SessionToken, looper: Looper): Media3Controller {
         return Media3Controller.Builder(context, token)
             .setApplicationLooper(looper)
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/Media3ActionFactoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/Media3ActionFactoryKosmos.kt
index 7e7eea216584..f49e3771763a 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/Media3ActionFactoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/Media3ActionFactoryKosmos.kt
@@ -16,7 +16,48 @@
 
 package com.android.systemui.media.controls.domain.pipeline
 
+import android.content.applicationContext
+import android.os.Bundle
+import android.os.Handler
+import android.os.looper
+import androidx.media3.session.CommandButton
+import androidx.media3.session.MediaController
+import androidx.media3.session.SessionToken
+import com.android.systemui.Flags
+import com.android.systemui.graphics.imageLoader
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.media.controls.shared.mediaLogger
+import com.android.systemui.media.controls.util.fakeMediaControllerFactory
+import com.android.systemui.media.controls.util.fakeSessionTokenFactory
+import com.google.common.collect.ImmutableList
 import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
-var Kosmos.media3ActionFactory: Media3ActionFactory by Kosmos.Fixture { mock {} }
+/**
+ * Set up fake [Media3ActionFactory]. Note that tests using this fake will need to be
+ * annotated @RunWithLooper
+ */
+var Kosmos.media3ActionFactory: Media3ActionFactory by
+    Kosmos.Fixture {
+        if (Flags.mediaControlsButtonMedia3()) {
+            val customLayout = ImmutableList.of<CommandButton>()
+            val media3Controller =
+                mock<MediaController>().also {
+                    whenever(it.customLayout).thenReturn(customLayout)
+                    whenever(it.sessionExtras).thenReturn(Bundle())
+                }
+            fakeMediaControllerFactory.setMedia3Controller(media3Controller)
+            fakeSessionTokenFactory.setMedia3SessionToken(mock<SessionToken>())
+        }
+        Media3ActionFactory(
+            context = applicationContext,
+            imageLoader = imageLoader,
+            controllerFactory = fakeMediaControllerFactory,
+            tokenFactory = fakeSessionTokenFactory,
+            logger = mediaLogger,
+            looper = looper,
+            handler = Handler(looper),
+            bgScope = testScope,
+        )
+    }