Isolate changes for output switcher grouping behind flag

The changes in b/31912901 seem to be causing crashes. This change looks
like it should've been a no-op for the non-flagged changed, but moved it
behind a flag anyway to isolate the issue.

Fixes: 399104884
Test: local test, atest
Flag: com.android.media.flags.enable_output_switcher_device_grouping
Change-Id: Ibf4fbf9bd696b4aadc7633b2da602e3a544b7ba0

2 files changed, 40 insertions, 0 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 9d375809786a..02dce406bbee 100644
--- a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaSwitchingController.java
+++ b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaSwitchingController.java
@@ -698,6 +698,27 @@ public class MediaSwitchingController
             List<MediaItem> finalMediaItems = targetMediaDevices.stream()
                     .map(MediaItem::createDeviceMediaItem)
                     .collect(Collectors.toList());
+
+            boolean shouldAddFirstSeenSelectedDevice =
+                    com.android.media.flags.Flags.enableOutputSwitcherDeviceGrouping();
+
+            if (shouldAddFirstSeenSelectedDevice) {
+                finalMediaItems.clear();
+                Set<String> selectedDevicesIds = getSelectedMediaDevice().stream()
+                        .map(MediaDevice::getId)
+                        .collect(Collectors.toSet());
+                for (MediaDevice targetMediaDevice : targetMediaDevices) {
+                    if (shouldAddFirstSeenSelectedDevice
+                            && selectedDevicesIds.contains(targetMediaDevice.getId())) {
+                        finalMediaItems.add(MediaItem.createDeviceMediaItem(
+                                targetMediaDevice, /* isFirstDeviceInGroup */ true));
+                        shouldAddFirstSeenSelectedDevice = false;
+                    } else {
+                        finalMediaItems.add(MediaItem.createDeviceMediaItem(
+                                targetMediaDevice, /* isFirstDeviceInGroup */ false));
+                    }
+                }
+            }
             dividerItems.forEach(finalMediaItems::add);
             attachConnectNewDeviceItemIfNeeded(finalMediaItems);
             return finalMediaItems;
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 88c2697fe2f2..5c26dac5eb30 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
@@ -1573,6 +1573,25 @@ public class MediaSwitchingControllerTest extends SysuiTestCase {
         assertThat(items.get(1).isFirstDeviceInGroup()).isFalse();
     }
 
+    @EnableFlags(Flags.FLAG_ENABLE_OUTPUT_SWITCHER_DEVICE_GROUPING)
+    @Test
+    public void deviceListUpdateWithDifferentDevices_firstSelectedDeviceIsFirstDeviceInGroup() {
+        when(mLocalMediaManager.isPreferenceRouteListingExist()).thenReturn(true);
+        doReturn(mMediaDevices)
+                .when(mLocalMediaManager)
+                .getSelectedMediaDevice();
+        mMediaSwitchingController.start(mCb);
+        reset(mCb);
+        mMediaSwitchingController.getMediaItemList().clear();
+        mMediaSwitchingController.onDeviceListUpdate(mMediaDevices);
+        mMediaDevices.clear();
+        mMediaDevices.add(mMediaDevice2);
+        mMediaSwitchingController.onDeviceListUpdate(mMediaDevices);
+
+        List<MediaItem> items = mMediaSwitchingController.getMediaItemList();
+        assertThat(items.get(0).isFirstDeviceInGroup()).isTrue();
+    }
+
     private int getNumberOfConnectDeviceButtons() {
         int numberOfConnectDeviceButtons = 0;
         for (MediaItem item : mMediaSwitchingController.getMediaItemList()) {