[Media TTT] Add TransferToThisDeviceSucceeded callback.

Bug: 203800643
Bug: 203800347
Test: verify `adb shell cmd statusbar media-ttt-chip-add-sender Device
TransferToThisDeviceSucceeded` shows the chip
Test: media.taptotransfer tests

Change-Id: Iba98b91617e25126e03c99d03e0a628994e6516a
diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/mediattt/IDeviceSenderService.aidl b/packages/SystemUI/shared/src/com/android/systemui/shared/mediattt/IDeviceSenderService.aidl
index 3b995093..67259f4bb 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/mediattt/IDeviceSenderService.aidl
+++ b/packages/SystemUI/shared/src/com/android/systemui/shared/mediattt/IDeviceSenderService.aidl
@@ -102,6 +102,22 @@
         in IUndoTransferCallback undoCallback);
 
     /**
+     * Invoke to notify System UI that a media transfer from the receiver and back to this device
+     * (the sender) has finished successfully.
+     *
+     * Important notes:
+     *   - This callback is for *ending* a cast. It should be used when media was previously being
+     *     played on the receiver device and has been successfully transferred to play locally on
+     *     this device instead.
+     *
+     * @param undoCallback will be invoked if the user chooses to undo this transfer.
+     */
+    oneway void transferToThisDeviceSucceeded(
+        in MediaRoute2Info mediaInfo,
+        in DeviceInfo otherDeviceInfo,
+        in IUndoTransferCallback undoCallback);
+
+    /**
      * Invoke to notify System UI that the attempted transfer has failed.
      *
      * This callback will be used for both the transfer that should've *started* playing the media
diff --git a/packages/SystemUI/src/com/android/systemui/media/taptotransfer/MediaTttCommandLineHelper.kt b/packages/SystemUI/src/com/android/systemui/media/taptotransfer/MediaTttCommandLineHelper.kt
index a2ab281..da93e92 100644
--- a/packages/SystemUI/src/com/android/systemui/media/taptotransfer/MediaTttCommandLineHelper.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/taptotransfer/MediaTttCommandLineHelper.kt
@@ -36,6 +36,7 @@
 import com.android.systemui.media.taptotransfer.sender.MoveCloserToStartCast
 import com.android.systemui.media.taptotransfer.sender.TransferFailed
 import com.android.systemui.media.taptotransfer.sender.TransferToReceiverTriggered
+import com.android.systemui.media.taptotransfer.sender.TransferToThisDeviceSucceeded
 import com.android.systemui.media.taptotransfer.sender.TransferToThisDeviceTriggered
 import com.android.systemui.media.taptotransfer.sender.TransferToReceiverSucceeded
 import com.android.systemui.shared.mediattt.DeviceInfo
@@ -108,7 +109,7 @@
                 TRANSFER_TO_RECEIVER_SUCCEEDED_COMMAND_NAME -> {
                     val undoCallback = object : IUndoTransferCallback.Stub() {
                         override fun onUndoTriggered() {
-                            Log.i(TAG, "Undo callback triggered")
+                            Log.i(TAG, "Undo transfer to receiver callback triggered")
                             // The external services that implement this callback would kick off a
                             // transfer back to this device, so mimic that here.
                             runOnService { senderService ->
@@ -122,6 +123,23 @@
                             .transferToReceiverSucceeded(mediaInfo, otherDeviceInfo, undoCallback)
                     }
                 }
+                TRANSFER_TO_THIS_DEVICE_SUCCEEDED_COMMAND_NAME -> {
+                    val undoCallback = object : IUndoTransferCallback.Stub() {
+                        override fun onUndoTriggered() {
+                            Log.i(TAG, "Undo transfer to this device callback triggered")
+                            // The external services that implement this callback would kick off a
+                            // transfer back to the receiver, so mimic that here.
+                            runOnService { senderService ->
+                                senderService
+                                    .transferToReceiverTriggered(mediaInfo, otherDeviceInfo)
+                            }
+                        }
+                    }
+                    runOnService { senderService ->
+                        senderService
+                            .transferToThisDeviceSucceeded(mediaInfo, otherDeviceInfo, undoCallback)
+                    }
+                }
                 TRANSFER_FAILED_COMMAND_NAME -> {
                     runOnService { senderService ->
                         senderService.transferFailed(mediaInfo, otherDeviceInfo)
@@ -134,6 +152,7 @@
                             "$TRANSFER_TO_RECEIVER_TRIGGERED_COMMAND_NAME, " +
                             "$TRANSFER_TO_THIS_DEVICE_TRIGGERED_COMMAND_NAME, " +
                             "$TRANSFER_TO_RECEIVER_SUCCEEDED_COMMAND_NAME, " +
+                            "$TRANSFER_TO_THIS_DEVICE_SUCCEEDED_COMMAND_NAME, " +
                             TRANSFER_FAILED_COMMAND_NAME
                     )
                 }
@@ -245,6 +264,9 @@
 @VisibleForTesting
 val TRANSFER_TO_RECEIVER_SUCCEEDED_COMMAND_NAME = TransferToReceiverSucceeded::class.simpleName!!
 @VisibleForTesting
+val TRANSFER_TO_THIS_DEVICE_SUCCEEDED_COMMAND_NAME =
+    TransferToThisDeviceSucceeded::class.simpleName!!
+@VisibleForTesting
 val TRANSFER_FAILED_COMMAND_NAME = TransferFailed::class.simpleName!!
 
 private const val APP_ICON_CONTENT_DESCRIPTION = "Fake media app icon"
diff --git a/packages/SystemUI/src/com/android/systemui/media/taptotransfer/sender/ChipStateSender.kt b/packages/SystemUI/src/com/android/systemui/media/taptotransfer/sender/ChipStateSender.kt
index 47969f2..c656df2 100644
--- a/packages/SystemUI/src/com/android/systemui/media/taptotransfer/sender/ChipStateSender.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/taptotransfer/sender/ChipStateSender.kt
@@ -18,6 +18,7 @@
 
 import android.content.Context
 import android.graphics.drawable.Drawable
+import android.view.View
 import com.android.systemui.R
 import com.android.systemui.media.taptotransfer.common.MediaTttChipState
 import com.android.systemui.shared.mediattt.IUndoTransferCallback
@@ -37,7 +38,18 @@
     abstract fun getChipTextString(context: Context): String
 
     /** Returns true if the loading icon should be displayed and false otherwise. */
-    abstract fun showLoading(): Boolean
+    open fun showLoading(): Boolean = false
+
+    /**
+     * Returns a click listener for the undo button on the chip. Returns null if this chip state
+     * doesn't have an undo button.
+     *
+     * @param controllerSender passed as a parameter in case we want to display a new chip state
+     *   when undo is clicked.
+     */
+    open fun undoClickListener(
+        controllerSender: MediaTttChipControllerSender
+    ): View.OnClickListener? = null
 }
 
 /**
@@ -55,8 +67,6 @@
     override fun getChipTextString(context: Context): String {
         return context.getString(R.string.media_move_closer_to_start_cast, otherDeviceName)
     }
-
-    override fun showLoading() = false
 }
 
 /**
@@ -74,8 +84,6 @@
     override fun getChipTextString(context: Context): String {
         return context.getString(R.string.media_move_closer_to_end_cast, otherDeviceName)
     }
-
-    override fun showLoading() = false
 }
 
 /**
@@ -128,7 +136,66 @@
         return context.getString(R.string.media_transfer_playing_different_device, otherDeviceName)
     }
 
-    override fun showLoading() = false
+    override fun undoClickListener(
+        controllerSender: MediaTttChipControllerSender
+    ): View.OnClickListener? {
+        if (undoCallback == null) {
+            return null
+        }
+
+        return View.OnClickListener {
+            this.undoCallback.onUndoTriggered()
+            // The external service should eventually send us a TransferToThisDeviceTriggered state,
+            // but that may take too long to go through the binder and the user may be confused as
+            // to why the UI hasn't changed yet. So, we immediately change the UI here.
+            controllerSender.displayChip(
+                TransferToThisDeviceTriggered(
+                    this.appIconDrawable,
+                    this.appIconContentDescription
+                )
+            )
+        }
+    }
+}
+
+/**
+ * A state representing that a transfer back to this device has been successfully completed.
+ *
+ * @property otherDeviceName the name of the other device involved in the transfer.
+ * @property undoCallback if present, the callback that should be called when the user clicks the
+ *   undo button. The undo button will only be shown if this is non-null.
+ */
+class TransferToThisDeviceSucceeded(
+    appIconDrawable: Drawable,
+    appIconContentDescription: String,
+    private val otherDeviceName: String,
+    val undoCallback: IUndoTransferCallback? = null
+) : ChipStateSender(appIconDrawable, appIconContentDescription) {
+    override fun getChipTextString(context: Context): String {
+        return context.getString(R.string.media_transfer_playing_this_device)
+    }
+
+    override fun undoClickListener(
+        controllerSender: MediaTttChipControllerSender
+    ): View.OnClickListener? {
+        if (undoCallback == null) {
+            return null
+        }
+
+        return View.OnClickListener {
+            this.undoCallback.onUndoTriggered()
+            // The external service should eventually send us a TransferToReceiverTriggered state,
+            // but that may take too long to go through the binder and the user may be confused as
+            // to why the UI hasn't changed yet. So, we immediately change the UI here.
+            controllerSender.displayChip(
+                TransferToReceiverTriggered(
+                    this.appIconDrawable,
+                    this.appIconContentDescription,
+                    this.otherDeviceName
+                )
+            )
+        }
+    }
 }
 
 /** A state representing that a transfer has failed. */
@@ -139,6 +206,4 @@
     override fun getChipTextString(context: Context): String {
         return context.getString(R.string.media_transfer_failed)
     }
-
-    override fun showLoading() = false
 }
diff --git a/packages/SystemUI/src/com/android/systemui/media/taptotransfer/sender/MediaTttChipControllerSender.kt b/packages/SystemUI/src/com/android/systemui/media/taptotransfer/sender/MediaTttChipControllerSender.kt
index 30a2809..453e3d6 100644
--- a/packages/SystemUI/src/com/android/systemui/media/taptotransfer/sender/MediaTttChipControllerSender.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/taptotransfer/sender/MediaTttChipControllerSender.kt
@@ -53,30 +53,10 @@
             if (chipState.showLoading()) { View.VISIBLE } else { View.GONE }
 
         // Undo
-        val undoClickListener: View.OnClickListener? =
-            if (chipState is TransferToReceiverSucceeded && chipState.undoCallback != null)
-                View.OnClickListener {
-                    chipState.undoCallback.onUndoTriggered()
-                    // The external service should eventually send us a
-                    // TransferToThisDeviceTriggered state, but that may take too long to go through
-                    // the binder and the user may be confused as to why the UI hasn't changed yet.
-                    // So, we immediately change the UI here.
-                    displayChip(
-                        TransferToThisDeviceTriggered(
-                            chipState.appIconDrawable,
-                            chipState.appIconContentDescription
-                        )
-                    )
-                }
-            else
-                null
         val undoView = currentChipView.requireViewById<View>(R.id.undo)
-        undoView.visibility = if (undoClickListener != null) {
-            View.VISIBLE
-        } else {
-            View.GONE
-        }
+        val undoClickListener = chipState.undoClickListener(this)
         undoView.setOnClickListener(undoClickListener)
+        undoView.visibility = if (undoClickListener != null) { View.VISIBLE } else { View.GONE }
 
         // Failure
         val showFailure = chipState is TransferFailed
diff --git a/packages/SystemUI/src/com/android/systemui/media/taptotransfer/sender/MediaTttSenderService.kt b/packages/SystemUI/src/com/android/systemui/media/taptotransfer/sender/MediaTttSenderService.kt
index 8431e25..8d9d7a9 100644
--- a/packages/SystemUI/src/com/android/systemui/media/taptotransfer/sender/MediaTttSenderService.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/taptotransfer/sender/MediaTttSenderService.kt
@@ -78,6 +78,16 @@
                 mediaInfo, otherDeviceInfo, undoCallback
             )
         }
+
+        override fun transferToThisDeviceSucceeded(
+            mediaInfo: MediaRoute2Info,
+            otherDeviceInfo: DeviceInfo,
+            undoCallback: IUndoTransferCallback
+        ) {
+            this@MediaTttSenderService.transferToThisDeviceSucceeded(
+                mediaInfo, otherDeviceInfo, undoCallback
+            )
+        }
     }
 
     // TODO(b/203800643): Use the app icon from the media info instead of a fake one.
@@ -146,4 +156,16 @@
         )
         controller.displayChip(chipState)
     }
+
+    private fun transferToThisDeviceSucceeded(
+        mediaInfo: MediaRoute2Info, otherDeviceInfo: DeviceInfo, undoCallback: IUndoTransferCallback
+    ) {
+        val chipState = TransferToThisDeviceSucceeded(
+            appIconDrawable = fakeAppIconDrawable,
+            appIconContentDescription = mediaInfo.name.toString(),
+            otherDeviceName = otherDeviceInfo.name,
+            undoCallback = undoCallback
+        )
+        controller.displayChip(chipState)
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/MediaTttCommandLineHelperTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/MediaTttCommandLineHelperTest.kt
index f2e2247..9b16305 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/MediaTttCommandLineHelperTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/MediaTttCommandLineHelperTest.kt
@@ -167,6 +167,18 @@
     }
 
     @Test
+    fun sender_transferToThisDeviceSucceeded_chipDisplayWithCorrectState() {
+        commandRegistry.onShellCommand(pw, getTransferToThisDeviceSucceededCommand())
+
+        assertThat(context.isBound(mediaSenderServiceComponentName)).isTrue()
+
+        val deviceInfoCaptor = argumentCaptor<DeviceInfo>()
+        verify(mediaSenderService)
+            .transferToThisDeviceSucceeded(any(), capture(deviceInfoCaptor), any())
+        assertThat(deviceInfoCaptor.value!!.name).isEqualTo(DEVICE_NAME)
+    }
+
+    @Test
     fun sender_transferFailed_serviceCallbackCalled() {
         commandRegistry.onShellCommand(pw, getTransferFailedCommand())
 
@@ -230,6 +242,13 @@
             TRANSFER_TO_RECEIVER_SUCCEEDED_COMMAND_NAME
         )
 
+    private fun getTransferToThisDeviceSucceededCommand(): Array<String> =
+        arrayOf(
+            ADD_CHIP_COMMAND_SENDER_TAG,
+            DEVICE_NAME,
+            TRANSFER_TO_THIS_DEVICE_SUCCEEDED_COMMAND_NAME
+        )
+
     private fun getTransferFailedCommand(): Array<String> =
         arrayOf(
             ADD_CHIP_COMMAND_SENDER_TAG,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/sender/MediaTttChipControllerSenderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/sender/MediaTttChipControllerSenderTest.kt
index 319b692..509ae33 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/sender/MediaTttChipControllerSenderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/sender/MediaTttChipControllerSenderTest.kt
@@ -170,6 +170,67 @@
     }
 
     @Test
+    fun transferToThisDeviceSucceeded_appIcon_deviceName_noLoadingIcon_noFailureIcon() {
+        val state = transferToThisDeviceSucceeded()
+        controllerSender.displayChip(state)
+
+        val chipView = getChipView()
+        assertThat(chipView.getAppIconView().drawable).isEqualTo(appIconDrawable)
+        assertThat(chipView.getAppIconView().contentDescription).isEqualTo(APP_ICON_CONTENT_DESC)
+        assertThat(chipView.getChipText()).isEqualTo(state.getChipTextString(context))
+        assertThat(chipView.getLoadingIconVisibility()).isEqualTo(View.GONE)
+        assertThat(chipView.getFailureIcon().visibility).isEqualTo(View.GONE)
+    }
+
+    @Test
+    fun transferToThisDeviceSucceeded_nullUndoRunnable_noUndo() {
+        controllerSender.displayChip(transferToThisDeviceSucceeded(undoCallback = null))
+
+        val chipView = getChipView()
+        assertThat(chipView.getUndoButton().visibility).isEqualTo(View.GONE)
+    }
+
+    @Test
+    fun transferToThisDeviceSucceeded_withUndoRunnable_undoWithClick() {
+        val undoCallback = object : IUndoTransferCallback.Stub() {
+            override fun onUndoTriggered() {}
+        }
+        controllerSender.displayChip(transferToThisDeviceSucceeded(undoCallback))
+
+        val chipView = getChipView()
+        assertThat(chipView.getUndoButton().visibility).isEqualTo(View.VISIBLE)
+        assertThat(chipView.getUndoButton().hasOnClickListeners()).isTrue()
+    }
+
+    @Test
+    fun transferToThisDeviceSucceeded_withUndoRunnable_undoButtonClickRunsRunnable() {
+        var undoCallbackCalled = false
+        val undoCallback = object : IUndoTransferCallback.Stub() {
+            override fun onUndoTriggered() {
+                undoCallbackCalled = true
+            }
+        }
+
+        controllerSender.displayChip(transferToThisDeviceSucceeded(undoCallback))
+        getChipView().getUndoButton().performClick()
+
+        assertThat(undoCallbackCalled).isTrue()
+    }
+
+    @Test
+    fun transferToThisDeviceSucceeded_undoButtonClick_switchesToTransferToReceiverTriggered() {
+        val undoCallback = object : IUndoTransferCallback.Stub() {
+            override fun onUndoTriggered() {}
+        }
+        controllerSender.displayChip(transferToThisDeviceSucceeded(undoCallback))
+
+        getChipView().getUndoButton().performClick()
+
+        assertThat(getChipView().getChipText())
+            .isEqualTo(transferToReceiverTriggered().getChipTextString(context))
+    }
+
+    @Test
     fun transferFailed_appIcon_noDeviceName_noLoadingIcon_noUndo_failureIcon() {
         val state = transferFailed()
         controllerSender.displayChip(state)
@@ -270,6 +331,12 @@
         )
 
     /** Helper method providing default parameters to not clutter up the tests. */
+    private fun transferToThisDeviceSucceeded(undoCallback: IUndoTransferCallback? = null) =
+        TransferToThisDeviceSucceeded(
+            appIconDrawable, APP_ICON_CONTENT_DESC, DEVICE_NAME, undoCallback
+        )
+
+    /** Helper method providing default parameters to not clutter up the tests. */
     private fun transferFailed() = TransferFailed(appIconDrawable, APP_ICON_CONTENT_DESC)
 }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/sender/MediaTttSenderServiceTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/sender/MediaTttSenderServiceTest.kt
index 177a1c2..e7304d4 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/sender/MediaTttSenderServiceTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/taptotransfer/sender/MediaTttSenderServiceTest.kt
@@ -95,6 +95,20 @@
     }
 
     @Test
+    fun transferToThisDeviceSucceeded_controllerTriggeredWithCorrectState() {
+        val undoCallback = object : IUndoTransferCallback.Stub() {
+            override fun onUndoTriggered() {}
+        }
+        service.transferToThisDeviceSucceeded(mediaInfo, DeviceInfo("name"), undoCallback)
+
+        val chipStateCaptor = argumentCaptor<TransferToThisDeviceSucceeded>()
+        verify(controller).displayChip(capture(chipStateCaptor))
+
+        val chipState = chipStateCaptor.value!!
+        assertThat(chipState.undoCallback).isEqualTo(undoCallback)
+    }
+
+    @Test
     fun transferFailed_controllerTriggeredWithTransferFailedState() {
         service.transferFailed(mediaInfo, DeviceInfo("Fake name"))