Ensure connect device button is shown

This CL fixes a bug in the logic of attachConnectNewDeviceItemIfNeeded
that caused the "Connect a device" button to disappear after the output
device was changed, if the input routing flag is enabled.

Before:
https://screencast.googleplex.com/cast/NTU5NTQzMDEzNzEwMjMzNnw0NWY1Mzk3YS0wZA

After:
https://screencast.googleplex.com/cast/NTQwMDMwMjcyNjQxNDMzNnw2NjA3MjM4Ni02MA

Test: atest MediaSwitchingControllerTest
Bug: b/379902375
Flag: com.android.media.flags.enable_audio_input_device_routing_and_volume_control
Change-Id: Ic16f0d18a5cac7e3279bbaaabbfdb89eb7b5b2e2

2 files changed, 70 insertions, 1 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaSwitchingController.java b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaSwitchingController.java
index b69b25dbddf2..8fbbb8bc4872 100644
--- a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaSwitchingController.java
+++ b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaSwitchingController.java
@@ -813,8 +813,15 @@ public class MediaSwitchingController
     }
 
     private void attachConnectNewDeviceItemIfNeeded(List<MediaItem> mediaItems) {
+        boolean isSelectedDeviceNotAGroup = getSelectedMediaDevice().size() == 1;
+        if (enableInputRouting()) {
+            // When input routing is enabled, there are expected to be at least 2 total selected
+            // devices: one output device and one input device.
+            isSelectedDeviceNotAGroup = getSelectedMediaDevice().size() <= 2;
+        }
+
         // Attach "Connect a device" item only when current output is not remote and not a group
-        if (!isCurrentConnectedDeviceRemote() && getSelectedMediaDevice().size() == 1) {
+        if (!isCurrentConnectedDeviceRemote() && isSelectedDeviceNotAGroup) {
             mediaItems.add(MediaItem.createPairNewDeviceMediaItem());
         }
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaSwitchingControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaSwitchingControllerTest.java
index eb1b44b75247..7ba797c03a0d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaSwitchingControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaSwitchingControllerTest.java
@@ -1467,4 +1467,66 @@ public class MediaSwitchingControllerTest extends SysuiTestCase {
         verify(mInputRouteManager, never()).selectDevice(outputMediaDevice);
         verify(mLocalMediaManager, atLeastOnce()).connectDevice(outputMediaDevice);
     }
+
+    @DisableFlags(Flags.FLAG_ENABLE_AUDIO_INPUT_DEVICE_ROUTING_AND_VOLUME_CONTROL)
+    @Test
+    public void connectDeviceButton_presentAtAllTimesForNonGroupOutputs() {
+        mMediaSwitchingController.start(mCb);
+        reset(mCb);
+
+        // Mock the selected output device.
+        doReturn(Collections.singletonList(mMediaDevice1))
+                .when(mLocalMediaManager)
+                .getSelectedMediaDevice();
+
+        // Verify that there is initially one "Connect a device" button present.
+        assertThat(getNumberOfConnectDeviceButtons()).isEqualTo(1);
+
+        // Change the selected device, and verify that there is still one "Connect a device" button
+        // present.
+        doReturn(Collections.singletonList(mMediaDevice2))
+                .when(mLocalMediaManager)
+                .getSelectedMediaDevice();
+        mMediaSwitchingController.onDeviceListUpdate(mMediaDevices);
+
+        assertThat(getNumberOfConnectDeviceButtons()).isEqualTo(1);
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_AUDIO_INPUT_DEVICE_ROUTING_AND_VOLUME_CONTROL)
+    @Test
+    public void connectDeviceButton_presentAtAllTimesForNonGroupOutputs_inputRoutingEnabled() {
+        mMediaSwitchingController.start(mCb);
+        reset(mCb);
+
+        // Mock the selected output device.
+        doReturn(Collections.singletonList(mMediaDevice1))
+                .when(mLocalMediaManager)
+                .getSelectedMediaDevice();
+
+        // Mock the selected input media device.
+        final MediaDevice selectedInputMediaDevice = mock(MediaDevice.class);
+        doReturn(selectedInputMediaDevice).when(mInputRouteManager).getSelectedInputDevice();
+
+        // Verify that there is initially one "Connect a device" button present.
+        assertThat(getNumberOfConnectDeviceButtons()).isEqualTo(1);
+
+        // Change the selected device, and verify that there is still one "Connect a device" button
+        // present.
+        doReturn(Collections.singletonList(mMediaDevice2))
+                .when(mLocalMediaManager)
+                .getSelectedMediaDevice();
+        mMediaSwitchingController.onDeviceListUpdate(mMediaDevices);
+
+        assertThat(getNumberOfConnectDeviceButtons()).isEqualTo(1);
+    }
+
+    private int getNumberOfConnectDeviceButtons() {
+        int numberOfConnectDeviceButtons = 0;
+        for (MediaItem item : mMediaSwitchingController.getMediaItemList()) {
+            if (item.getMediaItemType() == MediaItem.MediaItemType.TYPE_PAIR_NEW_DEVICE) {
+                numberOfConnectDeviceButtons++;
+            }
+        }
+        return numberOfConnectDeviceButtons;
+    }
 }