Merge "Fix grouped items not appearing if total items changed" into main

2 files changed, 40 insertions, 6 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..1f2f571496bd 100644
--- a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaSwitchingController.java
+++ b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaSwitchingController.java
@@ -648,6 +648,10 @@ public class MediaSwitchingController
             final MediaDevice connectedMediaDevice =
                     needToHandleMutingExpectedDevice ? null
                             : getCurrentConnectedMediaDevice();
+
+            Set<String> selectedDevicesIds = getSelectedMediaDevice().stream()
+                    .map(MediaDevice::getId)
+                    .collect(Collectors.toSet());
             if (oldMediaItems.isEmpty()) {
                 if (connectedMediaDevice == null) {
                     if (DEBUG) {
@@ -656,12 +660,14 @@ public class MediaSwitchingController
                     return categorizeMediaItemsLocked(
                             /* connectedMediaDevice */ null,
                             devices,
+                            selectedDevicesIds,
                             needToHandleMutingExpectedDevice);
                 } else {
                     // selected device exist
                     return categorizeMediaItemsLocked(
                             connectedMediaDevice,
                             devices,
+                            selectedDevicesIds,
                             /* needToHandleMutingExpectedDevice */ false);
                 }
             }
@@ -695,9 +701,20 @@ public class MediaSwitchingController
                 devices.removeAll(targetMediaDevices);
                 targetMediaDevices.addAll(devices);
             }
-            List<MediaItem> finalMediaItems = targetMediaDevices.stream()
-                    .map(MediaItem::createDeviceMediaItem)
-                    .collect(Collectors.toList());
+            List<MediaItem> finalMediaItems = new ArrayList<>();
+            boolean shouldAddFirstSeenSelectedDevice =
+                    com.android.media.flags.Flags.enableOutputSwitcherDeviceGrouping();
+            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;
@@ -724,11 +741,9 @@ public class MediaSwitchingController
     @GuardedBy("mMediaDevicesLock")
     private List<MediaItem> categorizeMediaItemsLocked(MediaDevice connectedMediaDevice,
             List<MediaDevice> devices,
+            Set<String> selectedDevicesIds,
             boolean needToHandleMutingExpectedDevice) {
         List<MediaItem> finalMediaItems = new ArrayList<>();
-        Set<String> selectedDevicesIds = getSelectedMediaDevice().stream()
-                .map(MediaDevice::getId)
-                .collect(Collectors.toSet());
         if (connectedMediaDevice != null) {
             selectedDevicesIds.add(connectedMediaDevice.getId());
         }
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()) {