Merge "Group session MediaItems together in OutputSwitcher" into main
diff --git a/media/java/android/media/flags/media_better_together.aconfig b/media/java/android/media/flags/media_better_together.aconfig
index 94454cf..405d292 100644
--- a/media/java/android/media/flags/media_better_together.aconfig
+++ b/media/java/android/media/flags/media_better_together.aconfig
@@ -125,6 +125,13 @@
 }
 
 flag {
+    name: "enable_output_switcher_session_grouping"
+    namespace: "media_better_together"
+    description: "Enables selected items in Output Switcher to be grouped together."
+    bug: "388347018"
+}
+
+flag {
     name: "enable_prevention_of_keep_alive_route_providers"
     namespace: "media_solutions"
     description: "Enables mechanisms to prevent route providers from keeping malicious apps alive."
diff --git a/packages/SettingsLib/src/com/android/settingslib/media/InfoMediaManager.java b/packages/SettingsLib/src/com/android/settingslib/media/InfoMediaManager.java
index ad196b8..4ee9ff0 100644
--- a/packages/SettingsLib/src/com/android/settingslib/media/InfoMediaManager.java
+++ b/packages/SettingsLib/src/com/android/settingslib/media/InfoMediaManager.java
@@ -658,12 +658,9 @@
         if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.UPSIDE_DOWN_CAKE) {
             RouteListingPreference routeListingPreference = getRouteListingPreference();
             if (routeListingPreference != null) {
-                final List<RouteListingPreference.Item> preferenceRouteListing =
-                        Api34Impl.composePreferenceRouteListing(
-                                routeListingPreference);
                 availableRoutes = Api34Impl.arrangeRouteListByPreference(selectedRoutes,
                         getAvailableRoutesFromRouter(),
-                                preferenceRouteListing);
+                        routeListingPreference);
             }
             return Api34Impl.filterDuplicatedIds(availableRoutes);
         } else {
@@ -760,11 +757,15 @@
         @DoNotInline
         static List<RouteListingPreference.Item> composePreferenceRouteListing(
                 RouteListingPreference routeListingPreference) {
+            boolean preferRouteListingOrdering =
+                    com.android.media.flags.Flags.enableOutputSwitcherSessionGrouping()
+                    && preferRouteListingOrdering(routeListingPreference);
             List<RouteListingPreference.Item> finalizedItemList = new ArrayList<>();
             List<RouteListingPreference.Item> itemList = routeListingPreference.getItems();
             for (RouteListingPreference.Item item : itemList) {
                 // Put suggested devices on the top first before further organization
-                if ((item.getFlags() & RouteListingPreference.Item.FLAG_SUGGESTED) != 0) {
+                if (!preferRouteListingOrdering
+                        && (item.getFlags() & RouteListingPreference.Item.FLAG_SUGGESTED) != 0) {
                     finalizedItemList.add(0, item);
                 } else {
                     finalizedItemList.add(item);
@@ -792,7 +793,7 @@
          * Returns an ordered list of available devices based on the provided {@code
          * routeListingPreferenceItems}.
          *
-         * <p>The result has the following order:
+         * <p>The resulting order if enableOutputSwitcherSessionGrouping is disabled is:
          *
          * <ol>
          *   <li>Selected routes.
@@ -800,22 +801,54 @@
          *   <li>Not-selected, non-system, available routes sorted by route listing preference.
          * </ol>
          *
+         * <p>The resulting order if enableOutputSwitcherSessionGrouping is enabled is:
+         *
+         * <ol>
+         *   <li>Selected routes sorted by route listing preference.
+         *   <li>Selected routes not defined by route listing preference.
+         *   <li>Not-selected system routes.
+         *   <li>Not-selected, non-system, available routes sorted by route listing preference.
+         * </ol>
+         *
+         *
          * @param selectedRoutes List of currently selected routes.
          * @param availableRoutes List of available routes that match the app's requested route
          *     features.
-         * @param routeListingPreferenceItems Ordered list of {@link RouteListingPreference.Item} to
-         *     sort routes with.
+         * @param routeListingPreference Preferences provided by the app to determine route order.
          */
         @DoNotInline
         static List<MediaRoute2Info> arrangeRouteListByPreference(
                 List<MediaRoute2Info> selectedRoutes,
                 List<MediaRoute2Info> availableRoutes,
-                List<RouteListingPreference.Item> routeListingPreferenceItems) {
+                RouteListingPreference routeListingPreference) {
+            final List<RouteListingPreference.Item> routeListingPreferenceItems =
+                    Api34Impl.composePreferenceRouteListing(routeListingPreference);
+
             Set<String> sortedRouteIds = new LinkedHashSet<>();
 
+            boolean addSelectedRlpItemsFirst =
+                    com.android.media.flags.Flags.enableOutputSwitcherSessionGrouping()
+                    && preferRouteListingOrdering(routeListingPreference);
+            Set<String> selectedRouteIds = new HashSet<>();
+
+            if (addSelectedRlpItemsFirst) {
+                // Add selected RLP items first
+                for (MediaRoute2Info selectedRoute : selectedRoutes) {
+                    selectedRouteIds.add(selectedRoute.getId());
+                }
+                for (RouteListingPreference.Item item: routeListingPreferenceItems) {
+                    if (selectedRouteIds.contains(item.getRouteId())) {
+                        sortedRouteIds.add(item.getRouteId());
+                    }
+                }
+            }
+
             // Add selected routes first.
-            for (MediaRoute2Info selectedRoute : selectedRoutes) {
-                sortedRouteIds.add(selectedRoute.getId());
+            if (com.android.media.flags.Flags.enableOutputSwitcherSessionGrouping()
+                    && sortedRouteIds.size() != selectedRoutes.size()) {
+                for (MediaRoute2Info selectedRoute : selectedRoutes) {
+                    sortedRouteIds.add(selectedRoute.getId());
+                }
             }
 
             // Add not-yet-added system routes.
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InfoMediaManagerTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InfoMediaManagerTest.java
index e1447dc..1a83f0a 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InfoMediaManagerTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InfoMediaManagerTest.java
@@ -48,15 +48,20 @@
 import android.media.RoutingSessionInfo;
 import android.media.session.MediaSessionManager;
 import android.os.Build;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.SetFlagsRule;
 
+import com.android.media.flags.Flags;
 import com.android.settingslib.bluetooth.CachedBluetoothDevice;
 import com.android.settingslib.bluetooth.CachedBluetoothDeviceManager;
 import com.android.settingslib.bluetooth.LocalBluetoothManager;
+import com.android.settingslib.media.InfoMediaManager.Api34Impl;
 import com.android.settingslib.testutils.shadow.ShadowRouter2Manager;
 
 import com.google.common.collect.ImmutableList;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -122,6 +127,8 @@
                     .addFeature(MediaRoute2Info.FEATURE_LIVE_AUDIO)
                     .build();
 
+    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
     @Mock
     private MediaRouter2Manager mRouterManager;
     @Mock
@@ -377,21 +384,26 @@
     }
 
     private RouteListingPreference setUpPreferenceList(String packageName) {
+        return setUpPreferenceList(packageName, false);
+    }
+
+    private RouteListingPreference setUpPreferenceList(
+                String packageName, boolean useSystemOrdering) {
         ReflectionHelpers.setStaticField(Build.VERSION.class, "SDK_INT",
                 Build.VERSION_CODES.UPSIDE_DOWN_CAKE);
         final List<RouteListingPreference.Item> preferenceItemList = new ArrayList<>();
-        RouteListingPreference.Item item1 =
+        RouteListingPreference.Item item1 = new RouteListingPreference.Item.Builder(
+                TEST_ID_3).build();
+        RouteListingPreference.Item item2 =
                 new RouteListingPreference.Item.Builder(TEST_ID_4)
                         .setFlags(RouteListingPreference.Item.FLAG_SUGGESTED)
                         .build();
-        RouteListingPreference.Item item2 = new RouteListingPreference.Item.Builder(
-                TEST_ID_3).build();
         preferenceItemList.add(item1);
         preferenceItemList.add(item2);
 
         RouteListingPreference routeListingPreference =
                 new RouteListingPreference.Builder().setItems(
-                        preferenceItemList).setUseSystemOrdering(false).build();
+                        preferenceItemList).setUseSystemOrdering(useSystemOrdering).build();
         when(mRouterManager.getRouteListingPreference(packageName))
                 .thenReturn(routeListingPreference);
         return routeListingPreference;
@@ -908,4 +920,66 @@
         assertThat(device.getState()).isEqualTo(STATE_SELECTED);
         assertThat(mInfoMediaManager.getCurrentConnectedDevice()).isEqualTo(device);
     }
+
+    @EnableFlags(Flags.FLAG_ENABLE_OUTPUT_SWITCHER_SESSION_GROUPING)
+    @Test
+    public void composePreferenceRouteListing_useSystemOrderingIsFalse() {
+        RouteListingPreference routeListingPreference =
+                setUpPreferenceList(TEST_PACKAGE_NAME, false);
+
+        List<RouteListingPreference.Item> routeOrder =
+                Api34Impl.composePreferenceRouteListing(routeListingPreference);
+
+        assertThat(routeOrder.get(0).getRouteId()).isEqualTo(TEST_ID_3);
+        assertThat(routeOrder.get(1).getRouteId()).isEqualTo(TEST_ID_4);
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_OUTPUT_SWITCHER_SESSION_GROUPING)
+    @Test
+    public void composePreferenceRouteListing_useSystemOrderingIsTrue() {
+        RouteListingPreference routeListingPreference =
+                setUpPreferenceList(TEST_PACKAGE_NAME, true);
+
+        List<RouteListingPreference.Item> routeOrder =
+                Api34Impl.composePreferenceRouteListing(routeListingPreference);
+
+        assertThat(routeOrder.get(0).getRouteId()).isEqualTo(TEST_ID_4);
+        assertThat(routeOrder.get(1).getRouteId()).isEqualTo(TEST_ID_3);
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_OUTPUT_SWITCHER_SESSION_GROUPING)
+    @Test
+    public void arrangeRouteListByPreference_useSystemOrderingIsFalse() {
+        RouteListingPreference routeListingPreference =
+                setUpPreferenceList(TEST_PACKAGE_NAME, false);
+        List<MediaRoute2Info> routes = setAvailableRoutesList(TEST_PACKAGE_NAME);
+        when(mRouterManager.getSelectedRoutes(any())).thenReturn(routes);
+
+        List<MediaRoute2Info> routeOrder =
+                Api34Impl.arrangeRouteListByPreference(
+                        routes, routes, routeListingPreference);
+
+        assertThat(routeOrder.get(0).getId()).isEqualTo(TEST_ID_3);
+        assertThat(routeOrder.get(1).getId()).isEqualTo(TEST_ID_4);
+        assertThat(routeOrder.get(2).getId()).isEqualTo(TEST_ID_2);
+        assertThat(routeOrder.get(3).getId()).isEqualTo(TEST_ID_1);
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_OUTPUT_SWITCHER_SESSION_GROUPING)
+    @Test
+    public void arrangeRouteListByPreference_useSystemOrderingIsTrue() {
+        RouteListingPreference routeListingPreference =
+                setUpPreferenceList(TEST_PACKAGE_NAME, true);
+        List<MediaRoute2Info> routes = setAvailableRoutesList(TEST_PACKAGE_NAME);
+        when(mRouterManager.getSelectedRoutes(any())).thenReturn(routes);
+
+        List<MediaRoute2Info> routeOrder =
+                Api34Impl.arrangeRouteListByPreference(
+                        routes, routes, routeListingPreference);
+
+        assertThat(routeOrder.get(0).getId()).isEqualTo(TEST_ID_2);
+        assertThat(routeOrder.get(1).getId()).isEqualTo(TEST_ID_3);
+        assertThat(routeOrder.get(2).getId()).isEqualTo(TEST_ID_4);
+        assertThat(routeOrder.get(3).getId()).isEqualTo(TEST_ID_1);
+    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/media/dialog/MediaOutputAdapterTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/media/dialog/MediaOutputAdapterTest.java
index 7478464..57ac906 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/media/dialog/MediaOutputAdapterTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/media/dialog/MediaOutputAdapterTest.java
@@ -117,8 +117,8 @@
                 LocalMediaManager.MediaDeviceState.STATE_DISCONNECTED);
         mMediaDevices.add(mMediaDevice1);
         mMediaDevices.add(mMediaDevice2);
-        mMediaItems.add(MediaItem.createDeviceMediaItem(mMediaDevice1));
-        mMediaItems.add(MediaItem.createDeviceMediaItem(mMediaDevice2));
+        mMediaItems.add(MediaItem.createDeviceMediaItem(mMediaDevice1, true));
+        mMediaItems.add(MediaItem.createDeviceMediaItem(mMediaDevice2, false));
 
         mMediaOutputAdapter = new MediaOutputAdapter(mMediaSwitchingController);
         mMediaOutputAdapter.updateItems();
@@ -779,4 +779,120 @@
                 mViewHolder.getDrawableId(false /* isInputDevice */, false /* isMutedVolumeIcon */))
                 .isEqualTo(R.drawable.media_output_icon_volume);
     }
+
+    @EnableFlags(Flags.FLAG_ENABLE_OUTPUT_SWITCHER_SESSION_GROUPING)
+    @Test
+    public void multipleSelectedDevices_verifySessionView() {
+        initializeSession();
+
+        mViewHolder = (MediaOutputAdapter.MediaDeviceViewHolder) mMediaOutputAdapter
+                .onCreateViewHolder(
+                        new LinearLayout(mContext), MediaItem.MediaItemType.TYPE_DEVICE);
+        mMediaOutputAdapter.onBindViewHolder(mViewHolder, 0);
+
+        assertThat(mViewHolder.mSeekBar.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mViewHolder.mTitleText.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mViewHolder.mCheckBox.getVisibility()).isEqualTo(View.GONE);
+        assertThat(mViewHolder.mEndTouchArea.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mViewHolder.mEndClickIcon.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mViewHolder.mTitleText.getText().toString()).isEqualTo(TEST_SESSION_NAME);
+        assertThat(mViewHolder.mSeekBar.getVolume()).isEqualTo(TEST_CURRENT_VOLUME);
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_OUTPUT_SWITCHER_SESSION_GROUPING)
+    @Test
+    public void multipleSelectedDevices_verifyCollapsedView() {
+        initializeSession();
+
+        mViewHolder = (MediaOutputAdapter.MediaDeviceViewHolder) mMediaOutputAdapter
+                .onCreateViewHolder(
+                        new LinearLayout(mContext), MediaItem.MediaItemType.TYPE_DEVICE);
+        mMediaOutputAdapter.onBindViewHolder(mViewHolder, 1);
+
+        assertThat(mViewHolder.mItemLayout.getVisibility()).isEqualTo(View.GONE);
+        assertThat(mViewHolder.mEndTouchArea.getVisibility()).isEqualTo(View.GONE);
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_OUTPUT_SWITCHER_SESSION_GROUPING)
+    @Test
+    public void multipleSelectedDevices_expandIconClicked_verifyInitialView() {
+        initializeSession();
+        mViewHolder = (MediaOutputAdapter.MediaDeviceViewHolder) mMediaOutputAdapter
+                .onCreateViewHolder(
+                        new LinearLayout(mContext), MediaItem.MediaItemType.TYPE_DEVICE);
+        mMediaOutputAdapter.onBindViewHolder(mViewHolder, 0);
+
+        mViewHolder.mEndTouchArea.performClick();
+        mViewHolder = (MediaOutputAdapter.MediaDeviceViewHolder) mMediaOutputAdapter
+                .onCreateViewHolder(
+                        new LinearLayout(mContext), MediaItem.MediaItemType.TYPE_DEVICE);
+        mMediaOutputAdapter.onBindViewHolder(mViewHolder, 0);
+
+        assertThat(mViewHolder.mSeekBar.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mViewHolder.mTitleText.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mViewHolder.mCheckBox.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mViewHolder.mEndTouchArea.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mViewHolder.mEndClickIcon.getVisibility()).isEqualTo(View.GONE);
+        assertThat(mViewHolder.mTitleText.getText().toString()).isEqualTo(TEST_DEVICE_NAME_1);
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_OUTPUT_SWITCHER_SESSION_GROUPING)
+    @Test
+    public void multipleSelectedDevices_expandIconClicked_verifyCollapsedView() {
+        initializeSession();
+        mViewHolder = (MediaOutputAdapter.MediaDeviceViewHolder) mMediaOutputAdapter
+                .onCreateViewHolder(
+                        new LinearLayout(mContext), MediaItem.MediaItemType.TYPE_DEVICE);
+        mMediaOutputAdapter.onBindViewHolder(mViewHolder, 0);
+
+        mViewHolder.mEndTouchArea.performClick();
+        mViewHolder = (MediaOutputAdapter.MediaDeviceViewHolder) mMediaOutputAdapter
+                .onCreateViewHolder(
+                        new LinearLayout(mContext), MediaItem.MediaItemType.TYPE_DEVICE);
+        mMediaOutputAdapter.onBindViewHolder(mViewHolder, 1);
+
+        assertThat(mViewHolder.mSeekBar.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mViewHolder.mTitleText.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mViewHolder.mCheckBox.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mViewHolder.mEndTouchArea.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mViewHolder.mEndClickIcon.getVisibility()).isEqualTo(View.GONE);
+        assertThat(mViewHolder.mTitleText.getText().toString()).isEqualTo(TEST_DEVICE_NAME_2);
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_OUTPUT_SWITCHER_SESSION_GROUPING)
+    @Test
+    public void deviceCanNotBeDeselected_verifyView() {
+        List<MediaDevice> selectedDevices = new ArrayList<>();
+        selectedDevices.add(mMediaDevice1);
+        when(mMediaSwitchingController.getSelectableMediaDevice()).thenReturn(selectedDevices);
+        when(mMediaSwitchingController.getSelectedMediaDevice()).thenReturn(selectedDevices);
+        when(mMediaSwitchingController.getDeselectableMediaDevice()).thenReturn(new ArrayList<>());
+
+        mViewHolder = (MediaOutputAdapter.MediaDeviceViewHolder) mMediaOutputAdapter
+                .onCreateViewHolder(
+                        new LinearLayout(mContext), MediaItem.MediaItemType.TYPE_DEVICE);
+        mMediaOutputAdapter.onBindViewHolder(mViewHolder, 0);
+
+        assertThat(mViewHolder.mSeekBar.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mViewHolder.mTitleText.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(mViewHolder.mCheckBox.getVisibility()).isEqualTo(View.GONE);
+        assertThat(mViewHolder.mEndTouchArea.getVisibility()).isEqualTo(View.GONE);
+        assertThat(mViewHolder.mEndClickIcon.getVisibility()).isEqualTo(View.GONE);
+        assertThat(mViewHolder.mTitleText.getText().toString()).isEqualTo(TEST_DEVICE_NAME_1);
+    }
+
+    private void initializeSession() {
+        when(mMediaSwitchingController.getSessionVolumeMax()).thenReturn(TEST_MAX_VOLUME);
+        when(mMediaSwitchingController.getSessionVolume()).thenReturn(TEST_CURRENT_VOLUME);
+        when(mMediaSwitchingController.getSessionName()).thenReturn(TEST_SESSION_NAME);
+
+        List<MediaDevice> selectedDevices = new ArrayList<>();
+        selectedDevices.add(mMediaDevice1);
+        selectedDevices.add(mMediaDevice2);
+        when(mMediaSwitchingController.getSelectableMediaDevice()).thenReturn(selectedDevices);
+        when(mMediaSwitchingController.getSelectedMediaDevice()).thenReturn(selectedDevices);
+        when(mMediaSwitchingController.getDeselectableMediaDevice()).thenReturn(selectedDevices);
+
+        mMediaOutputAdapter.updateItems();
+    }
 }
diff --git a/packages/SystemUI/res/drawable/media_output_item_expand_group.xml b/packages/SystemUI/res/drawable/media_output_item_expand_group.xml
new file mode 100644
index 0000000..833843d
--- /dev/null
+++ b/packages/SystemUI/res/drawable/media_output_item_expand_group.xml
@@ -0,0 +1,26 @@
+<!--
+  ~ Copyright (C) 2025 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="24"
+    android:viewportHeight="24"
+    android:tint="?attr/colorControlNormal">
+    <path
+        android:fillColor="@android:color/white"
+        android:pathData="M12,15.4 L6,9.4l1.4,-1.4 4.6,4.6 4.6,-4.6 1.4,1.4 -6,6Z" />
+</vector>
diff --git a/packages/SystemUI/res/values/strings.xml b/packages/SystemUI/res/values/strings.xml
index 84c859c..fc9635b 100644
--- a/packages/SystemUI/res/values/strings.xml
+++ b/packages/SystemUI/res/values/strings.xml
@@ -588,6 +588,12 @@
     <!-- Content description of the cast label showing what we are connected to. [CHAR LIMIT=NONE] -->
     <string name="accessibility_cast_name">Connected to <xliff:g id="cast" example="TV">%s</xliff:g>.</string>
 
+    <!-- Content description of the button to expand the group of devices. [CHAR LIMIT=NONE] -->
+    <string name="accessibility_expand_group">Expand group.</string>
+
+    <!-- Content description of the button to open the application . [CHAR LIMIT=NONE] -->
+    <string name="accessibility_open_application">Open application.</string>
+
     <!-- Content description of an item with no signal and no connection for accessibility (not shown on the screen) [CHAR LIMIT=NONE] -->
     <string name="accessibility_not_connected">Not connected.</string>
     <!-- Content description of the roaming data connection type. [CHAR LIMIT=NONE] -->
diff --git a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaItem.java b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaItem.java
index 4496b25..7b1c62e2 100644
--- a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaItem.java
+++ b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaItem.java
@@ -36,6 +36,7 @@
     private final String mTitle;
     @MediaItemType
     private final int mMediaItemType;
+    private final boolean mIsFirstDeviceInGroup;
 
     @Retention(RetentionPolicy.SOURCE)
     @IntDef({
@@ -54,7 +55,18 @@
      * name.
      */
     public static MediaItem createDeviceMediaItem(@NonNull MediaDevice device) {
-        return new MediaItem(device, device.getName(), MediaItemType.TYPE_DEVICE);
+        return new MediaItem(device, device.getName(), MediaItemType.TYPE_DEVICE, false);
+    }
+
+    /**
+     * Returns a new {@link MediaItemType#TYPE_DEVICE} {@link MediaItem} with its {@link
+     * #getMediaDevice() media device} set to {@code device} and its title set to {@code device}'s
+     * name.
+     */
+    public static MediaItem createDeviceMediaItem(
+            @NonNull MediaDevice device, boolean isFirstDeviceInGroup) {
+        return new MediaItem(
+            device, device.getName(), MediaItemType.TYPE_DEVICE, isFirstDeviceInGroup);
     }
 
     /**
@@ -63,7 +75,10 @@
      */
     public static MediaItem createPairNewDeviceMediaItem() {
         return new MediaItem(
-                /* device */ null, /* title */ null, MediaItemType.TYPE_PAIR_NEW_DEVICE);
+                /* device */ null,
+                /* title */ null,
+                MediaItemType.TYPE_PAIR_NEW_DEVICE,
+                /* mIsFirstDeviceInGroup */ false);
     }
 
     /**
@@ -71,14 +86,22 @@
      * title and a {@code null} {@link #getMediaDevice() media device}.
      */
     public static MediaItem createGroupDividerMediaItem(@Nullable String title) {
-        return new MediaItem(/* device */ null, title, MediaItemType.TYPE_GROUP_DIVIDER);
+        return new MediaItem(
+            /* device */ null,
+            title,
+            MediaItemType.TYPE_GROUP_DIVIDER,
+            /* misFirstDeviceInGroup */ false);
     }
 
     private MediaItem(
-            @Nullable MediaDevice device, @Nullable String title, @MediaItemType int type) {
+            @Nullable MediaDevice device,
+            @Nullable String title,
+            @MediaItemType int type,
+            boolean isFirstDeviceInGroup) {
         this.mMediaDeviceOptional = Optional.ofNullable(device);
         this.mTitle = title;
         this.mMediaItemType = type;
+        this.mIsFirstDeviceInGroup = isFirstDeviceInGroup;
     }
 
     public Optional<MediaDevice> getMediaDevice() {
@@ -106,4 +129,8 @@
     public int getMediaItemType() {
         return mMediaItemType;
     }
+
+    public boolean isFirstDeviceInGroup() {
+        return mIsFirstDeviceInGroup;
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputAdapter.java b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputAdapter.java
index 53f3b3a..52b3c3ec 100644
--- a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputAdapter.java
+++ b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputAdapter.java
@@ -21,6 +21,7 @@
 import static com.android.settingslib.media.MediaDevice.SelectionBehavior.SELECTION_BEHAVIOR_TRANSFER;
 
 import android.annotation.DrawableRes;
+import android.annotation.StringRes;
 import android.content.Context;
 import android.content.res.ColorStateList;
 import android.graphics.drawable.AnimatedVectorDrawable;
@@ -38,6 +39,7 @@
 import androidx.recyclerview.widget.RecyclerView;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.media.flags.Flags;
 import com.android.settingslib.media.LocalMediaManager.MediaDeviceState;
 import com.android.settingslib.media.MediaDevice;
 import com.android.systemui.res.R;
@@ -55,6 +57,7 @@
     private static final float DEVICE_DISCONNECTED_ALPHA = 0.5f;
     private static final float DEVICE_CONNECTED_ALPHA = 1f;
     protected List<MediaItem> mMediaItemList = new CopyOnWriteArrayList<>();
+    private boolean mShouldGroupSelectedMediaItems = Flags.enableOutputSwitcherSessionGrouping();
 
     public MediaOutputAdapter(MediaSwitchingController controller) {
         super(controller);
@@ -65,6 +68,12 @@
     public void updateItems() {
         mMediaItemList.clear();
         mMediaItemList.addAll(mController.getMediaItemList());
+        if (mShouldGroupSelectedMediaItems) {
+            if (mController.getSelectedMediaDevice().size() == 1) {
+                // Don't group devices if initially there isn't more than one selected.
+                mShouldGroupSelectedMediaItems = false;
+            }
+        }
         notifyDataSetChanged();
     }
 
@@ -101,7 +110,7 @@
                 break;
             case MediaItem.MediaItemType.TYPE_DEVICE:
                 ((MediaDeviceViewHolder) viewHolder).onBind(
-                        currentMediaItem.getMediaDevice().get(),
+                        currentMediaItem,
                         position);
                 break;
             default:
@@ -141,8 +150,8 @@
             super(view);
         }
 
-        @Override
-        void onBind(MediaDevice device, int position) {
+        void onBind(MediaItem mediaItem, int position) {
+            MediaDevice device = mediaItem.getMediaDevice().get();
             super.onBind(device, position);
             boolean isMutingExpectedDeviceExist = mController.hasMutingExpectedDevice();
             final boolean currentlyConnected = isCurrentlyConnected(device);
@@ -150,6 +159,7 @@
             if (mCurrentActivePosition == position) {
                 mCurrentActivePosition = -1;
             }
+            mItemLayout.setVisibility(View.VISIBLE);
             mStatusIcon.setVisibility(View.GONE);
             enableFocusPropertyForView(mContainerLayout);
 
@@ -174,6 +184,30 @@
                     updateFullItemClickListener(v -> onItemClick(v, device));
                     setSingleLineLayout(device.getName());
                     initFakeActiveDevice(device);
+                } else if (mShouldGroupSelectedMediaItems
+                        && mController.getSelectedMediaDevice().size() > 1
+                        && isDeviceIncluded(mController.getSelectedMediaDevice(), device)) {
+                    if (!mediaItem.isFirstDeviceInGroup()) {
+                        mItemLayout.setVisibility(View.GONE);
+                        mEndTouchArea.setVisibility(View.GONE);
+                    } else {
+                        String sessionName = mController.getSessionName().toString();
+                        updateUnmutedVolumeIcon(null);
+                        updateEndClickAreaWithIcon(
+                                v -> {
+                                    mShouldGroupSelectedMediaItems = false;
+                                    notifyDataSetChanged();
+                                },
+                                R.drawable.media_output_item_expand_group,
+                                R.string.accessibility_expand_group);
+                        disableFocusPropertyForView(mContainerLayout);
+                        setUpContentDescriptionForView(mSeekBar, mContext.getString(
+                                R.string.accessibility_cast_name, sessionName));
+                        setSingleLineLayout(sessionName, true /* showSeekBar */,
+                                false /* showProgressBar */, false /* showCheckBox */,
+                                true /* showEndTouchArea */);
+                        initGroupSeekbar(isCurrentSeekbarInvisible);
+                    }
                 } else if (device.hasSubtext()) {
                     boolean isActiveWithOngoingSession =
                             (device.hasOngoingSession() && (currentlyConnected || isDeviceIncluded(
@@ -237,6 +271,8 @@
                     // selected device in group
                     boolean isDeviceDeselectable = isDeviceIncluded(
                             mController.getDeselectableMediaDevice(), device);
+                    boolean showEndArea = !Flags.enableOutputSwitcherSessionGrouping()
+                            || isDeviceDeselectable;
                     updateUnmutedVolumeIcon(device);
                     updateGroupableCheckBox(true, isDeviceDeselectable, device);
                     updateEndClickArea(device, isDeviceDeselectable);
@@ -244,7 +280,7 @@
                     setUpContentDescriptionForView(mSeekBar, device);
                     setSingleLineLayout(device.getName(), true /* showSeekBar */,
                             false /* showProgressBar */, true /* showCheckBox */,
-                            true /* showEndTouchArea */);
+                            showEndArea /* showEndTouchArea */);
                     initSeekbar(device, isCurrentSeekbarInvisible);
                 } else if (!mController.hasAdjustVolumeUserRestriction()
                         && currentlyConnected) {
@@ -335,19 +371,29 @@
         }
 
         private void updateEndClickAreaAsSessionEditing(MediaDevice device, @DrawableRes int id) {
-            mEndClickIcon.setOnClickListener(null);
-            mEndTouchArea.setOnClickListener(null);
+            updateEndClickAreaWithIcon(
+                    v -> mController.tryToLaunchInAppRoutingIntent(device.getId(), v),
+                    id,
+                    R.string.accessibility_open_application);
+        }
+
+        private void updateEndClickAreaWithIcon(View.OnClickListener clickListener,
+                @DrawableRes int iconDrawableId,
+                @StringRes int accessibilityStringId) {
             updateEndClickAreaColor(mController.getColorSeekbarProgress());
             mEndClickIcon.setImageTintList(
                     ColorStateList.valueOf(mController.getColorItemContent()));
-            mEndClickIcon.setOnClickListener(
-                    v -> mController.tryToLaunchInAppRoutingIntent(device.getId(), v));
+            mEndClickIcon.setOnClickListener(clickListener);
             mEndTouchArea.setOnClickListener(v -> mEndClickIcon.performClick());
-            Drawable drawable = mContext.getDrawable(id);
+            Drawable drawable = mContext.getDrawable(iconDrawableId);
             mEndClickIcon.setImageDrawable(drawable);
             if (drawable instanceof AnimatedVectorDrawable) {
                 ((AnimatedVectorDrawable) drawable).start();
             }
+            if (Flags.enableOutputSwitcherSessionGrouping()) {
+                setUpContentDescriptionForView(
+                        mEndClickIcon, mContext.getString(accessibilityStringId));
+            }
         }
 
         public void updateEndClickAreaColor(int color) {
@@ -479,12 +525,17 @@
         }
 
         private void setUpContentDescriptionForView(View view, MediaDevice device) {
-            view.setContentDescription(
+            setUpContentDescriptionForView(
+                    view,
                     mContext.getString(device.getDeviceType()
                             == MediaDevice.MediaDeviceType.TYPE_BLUETOOTH_DEVICE
                             ? R.string.accessibility_bluetooth_name
                             : R.string.accessibility_cast_name, device.getName()));
         }
+
+        protected void setUpContentDescriptionForView(View view, String description) {
+            view.setContentDescription(description);
+        }
     }
 
     class MediaGroupDividerViewHolder extends RecyclerView.ViewHolder {
diff --git a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputBaseAdapter.java b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputBaseAdapter.java
index 9b24c69..ee2d8aa 100644
--- a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputBaseAdapter.java
+++ b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputBaseAdapter.java
@@ -42,6 +42,7 @@
 import androidx.annotation.VisibleForTesting;
 import androidx.recyclerview.widget.RecyclerView;
 
+import com.android.media.flags.Flags;
 import com.android.settingslib.media.InputMediaDevice;
 import com.android.settingslib.media.MediaDevice;
 import com.android.settingslib.utils.ThreadUtils;
@@ -211,6 +212,10 @@
             mTitleText.setText(title);
             mCheckBox.setVisibility(showCheckBox ? View.VISIBLE : View.GONE);
             mEndTouchArea.setVisibility(showEndTouchArea ? View.VISIBLE : View.GONE);
+            if (Flags.enableOutputSwitcherSessionGrouping()) {
+                mEndClickIcon.setVisibility(
+                        !showCheckBox && showEndTouchArea ? View.VISIBLE : View.GONE);
+            }
             ViewGroup.MarginLayoutParams params =
                     (ViewGroup.MarginLayoutParams) mItemLayout.getLayoutParams();
             params.rightMargin = showEndTouchArea ? mController.getItemMarginEndSelectable()
@@ -265,14 +270,8 @@
                             mController.getActiveRadius(), 0, 0});
         }
 
-        void initSeekbar(MediaDevice device, boolean isCurrentSeekbarInvisible) {
-            if (!mController.isVolumeControlEnabled(device)) {
-                disableSeekBar();
-            } else {
-                enableSeekBar(device);
-            }
-            mSeekBar.setMaxVolume(device.getMaxVolume());
-            final int currentVolume = device.getCurrentVolume();
+        private void initializeSeekbarVolume(
+                MediaDevice device, int currentVolume, boolean isCurrentSeekbarInvisible) {
             if (!mIsDragging) {
                 if (mSeekBar.getVolume() != currentVolume && (mLatestUpdateVolume == -1
                         || currentVolume == mLatestUpdateVolume)) {
@@ -307,54 +306,75 @@
             if (mIsInitVolumeFirstTime) {
                 mIsInitVolumeFirstTime = false;
             }
-            mSeekBar.setOnSeekBarChangeListener(new SeekBar.OnSeekBarChangeListener() {
-                boolean mStartFromMute = false;
+        }
+
+        void initSeekbar(MediaDevice device, boolean isCurrentSeekbarInvisible) {
+            SeekBarVolumeControl volumeControl = new SeekBarVolumeControl() {
                 @Override
-                public void onProgressChanged(SeekBar seekBar, int progress, boolean fromUser) {
-                    if (device == null || !fromUser) {
-                        return;
-                    }
-
-                    final String percentageString = mContext.getResources().getString(
-                            R.string.media_output_dialog_volume_percentage,
-                            mSeekBar.getPercentage());
-                    mVolumeValueText.setText(percentageString);
-
-                    if (mStartFromMute) {
-                        updateUnmutedVolumeIcon(device);
-                        mStartFromMute = false;
-                    }
-                    int seekBarVolume = MediaOutputSeekbar.scaleProgressToVolume(progress);
-                    if (seekBarVolume != device.getCurrentVolume()) {
-                        mLatestUpdateVolume = seekBarVolume;
-                        mController.adjustVolume(device, seekBarVolume);
-                    }
+                public int getVolume() {
+                    return device.getCurrentVolume();
+                }
+                @Override
+                public void setVolume(int volume) {
+                    mController.adjustVolume(device, volume);
                 }
 
                 @Override
-                public void onStartTrackingTouch(SeekBar seekBar) {
-                    mTitleIcon.setVisibility(View.INVISIBLE);
-                    mVolumeValueText.setVisibility(View.VISIBLE);
-                    int currentVolume = MediaOutputSeekbar.scaleProgressToVolume(
-                            seekBar.getProgress());
-                    mStartFromMute = (currentVolume == 0);
-                    mIsDragging = true;
+                public void onMute() {
+                    mController.logInteractionUnmuteDevice(device);
                 }
+            };
 
+            if (!mController.isVolumeControlEnabled(device)) {
+                disableSeekBar();
+            } else {
+                enableSeekBar(volumeControl);
+            }
+            mSeekBar.setMaxVolume(device.getMaxVolume());
+            final int currentVolume = device.getCurrentVolume();
+            initializeSeekbarVolume(device, currentVolume, isCurrentSeekbarInvisible);
+
+            mSeekBar.setOnSeekBarChangeListener(new MediaSeekBarChangedListener(
+                    device, volumeControl) {
                 @Override
-                public void onStopTrackingTouch(SeekBar seekBar) {
-                    int currentVolume = MediaOutputSeekbar.scaleProgressToVolume(
-                            seekBar.getProgress());
-                    if (currentVolume == 0) {
-                        seekBar.setProgress(0);
-                        updateMutedVolumeIcon(device);
-                    } else {
-                        updateUnmutedVolumeIcon(device);
-                    }
-                    mTitleIcon.setVisibility(View.VISIBLE);
-                    mVolumeValueText.setVisibility(View.GONE);
+                public void onStopTrackingTouch(SeekBar seekbar) {
+                    super.onStopTrackingTouch(seekbar);
                     mController.logInteractionAdjustVolume(device);
-                    mIsDragging = false;
+                }
+            });
+        }
+
+        // Initializes the seekbar for a group of devices.
+        void initGroupSeekbar(boolean isCurrentSeekbarInvisible) {
+            SeekBarVolumeControl volumeControl = new SeekBarVolumeControl() {
+                @Override
+                public int getVolume() {
+                    return mController.getSessionVolume();
+                }
+
+                @Override
+                public void setVolume(int volume) {
+                    mController.adjustSessionVolume(volume);
+                }
+
+                @Override
+                public void onMute() {}
+            };
+
+            if (!mController.isVolumeControlEnabledForSession()) {
+                disableSeekBar();
+            } else {
+                enableSeekBar(volumeControl);
+            }
+            mSeekBar.setMaxVolume(mController.getSessionVolumeMax());
+
+            final int currentVolume = mController.getSessionVolume();
+            initializeSeekbarVolume(null, currentVolume, isCurrentSeekbarInvisible);
+            mSeekBar.setOnSeekBarChangeListener(new MediaSeekBarChangedListener(
+                    null, volumeControl) {
+                @Override
+                protected boolean shouldHandleProgressChanged() {
+                    return true;
                 }
             });
         }
@@ -385,7 +405,7 @@
         int getDrawableId(boolean isInputDevice, boolean isMutedVolumeIcon) {
             // Returns the microphone icon when the flag is enabled and the device is an input
             // device.
-            if (com.android.media.flags.Flags.enableAudioInputDeviceRoutingAndVolumeControl()
+            if (Flags.enableAudioInputDeviceRoutingAndVolumeControl()
                     && isInputDevice) {
                 return isMutedVolumeIcon ? R.drawable.ic_mic_off : R.drawable.ic_mic_26dp;
             }
@@ -452,27 +472,28 @@
             updateIconAreaClickListener(null);
         }
 
-        private void enableSeekBar(MediaDevice device) {
+        private void enableSeekBar(SeekBarVolumeControl volumeControl) {
             mSeekBar.setEnabled(true);
+
             mSeekBar.setOnTouchListener((v, event) -> false);
             updateIconAreaClickListener((v) -> {
-                if (device.getCurrentVolume() == 0) {
-                    mController.logInteractionUnmuteDevice(device);
+                if (volumeControl.getVolume() == 0) {
                     mSeekBar.setVolume(UNMUTE_DEFAULT_VOLUME);
-                    mController.adjustVolume(device, UNMUTE_DEFAULT_VOLUME);
-                    updateUnmutedVolumeIcon(device);
+                    volumeControl.setVolume(UNMUTE_DEFAULT_VOLUME);
+                    updateUnmutedVolumeIcon(null);
                     mIconAreaLayout.setOnTouchListener(((iconV, event) -> false));
                 } else {
-                    mController.logInteractionMuteDevice(device);
+                    volumeControl.onMute();
                     mSeekBar.resetVolume();
-                    mController.adjustVolume(device, 0);
-                    updateMutedVolumeIcon(device);
+                    volumeControl.setVolume(0);
+                    updateMutedVolumeIcon(null);
                     mIconAreaLayout.setOnTouchListener(((iconV, event) -> {
                         mSeekBar.dispatchTouchEvent(event);
                         return false;
                     }));
                 }
             });
+
         }
 
         protected void setUpDeviceIcon(MediaDevice device) {
@@ -488,5 +509,74 @@
                 });
             });
         }
+
+        interface SeekBarVolumeControl {
+            int getVolume();
+            void setVolume(int volume);
+            void onMute();
+        }
+
+        private abstract class MediaSeekBarChangedListener
+                implements SeekBar.OnSeekBarChangeListener {
+            boolean mStartFromMute = false;
+            private MediaDevice mMediaDevice;
+            private SeekBarVolumeControl mVolumeControl;
+
+            MediaSeekBarChangedListener(MediaDevice device, SeekBarVolumeControl volumeControl) {
+                mMediaDevice = device;
+                mVolumeControl = volumeControl;
+            }
+
+            @Override
+            public void onProgressChanged(SeekBar seekBar, int progress, boolean fromUser) {
+                if (!shouldHandleProgressChanged() || !fromUser) {
+                    return;
+                }
+
+                final String percentageString = mContext.getResources().getString(
+                        R.string.media_output_dialog_volume_percentage,
+                        mSeekBar.getPercentage());
+                mVolumeValueText.setText(percentageString);
+
+                if (mStartFromMute) {
+                    updateUnmutedVolumeIcon(mMediaDevice);
+                    mStartFromMute = false;
+                }
+
+                int seekBarVolume = MediaOutputSeekbar.scaleProgressToVolume(progress);
+                if (seekBarVolume != mVolumeControl.getVolume()) {
+                    mLatestUpdateVolume = seekBarVolume;
+                    mVolumeControl.setVolume(seekBarVolume);
+                }
+            }
+
+            @Override
+            public void onStartTrackingTouch(SeekBar seekBar) {
+                mTitleIcon.setVisibility(View.INVISIBLE);
+                mVolumeValueText.setVisibility(View.VISIBLE);
+                int currentVolume = MediaOutputSeekbar.scaleProgressToVolume(
+                        seekBar.getProgress());
+                mStartFromMute = (currentVolume == 0);
+                mIsDragging = true;
+            }
+
+            @Override
+            public void onStopTrackingTouch(SeekBar seekBar) {
+                int currentVolume = MediaOutputSeekbar.scaleProgressToVolume(
+                        seekBar.getProgress());
+                if (currentVolume == 0) {
+                    seekBar.setProgress(0);
+                    updateMutedVolumeIcon(mMediaDevice);
+                } else {
+                    updateUnmutedVolumeIcon(mMediaDevice);
+                }
+                mTitleIcon.setVisibility(View.VISIBLE);
+                mVolumeValueText.setVisibility(View.GONE);
+                mIsDragging = false;
+            }
+            protected boolean shouldHandleProgressChanged() {
+                return mMediaDevice != null;
+            }
+        };
     }
 }
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 15afd22..35c872f 100644
--- a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaSwitchingController.java
+++ b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaSwitchingController.java
@@ -760,14 +760,26 @@
         if (connectedMediaDevice != null) {
             selectedDevicesIds.add(connectedMediaDevice.getId());
         }
+        boolean groupSelectedDevices =
+                com.android.media.flags.Flags.enableOutputSwitcherSessionGrouping();
+        int nextSelectedItemIndex = 0;
         boolean suggestedDeviceAdded = false;
         boolean displayGroupAdded = false;
+        boolean selectedDeviceAdded = false;
         for (MediaDevice device : devices) {
             if (needToHandleMutingExpectedDevice && device.isMutingExpectedDevice()) {
                 finalMediaItems.add(0, MediaItem.createDeviceMediaItem(device));
+                nextSelectedItemIndex++;
             } else if (!needToHandleMutingExpectedDevice && selectedDevicesIds.contains(
                     device.getId())) {
-                finalMediaItems.add(0, MediaItem.createDeviceMediaItem(device));
+                if (groupSelectedDevices) {
+                    finalMediaItems.add(
+                            nextSelectedItemIndex++,
+                            MediaItem.createDeviceMediaItem(device, !selectedDeviceAdded));
+                    selectedDeviceAdded = true;
+                } else {
+                    finalMediaItems.add(0, MediaItem.createDeviceMediaItem(device));
+                }
             } else {
                 if (device.isSuggestedDevice() && !suggestedDeviceAdded) {
                     addSuggestedDeviceGroupDivider(finalMediaItems);
@@ -1331,6 +1343,10 @@
         return !device.isVolumeFixed();
     }
 
+    boolean isVolumeControlEnabledForSession() {
+        return mLocalMediaManager.isMediaSessionAvailableForVolumeControl();
+    }
+
     private void startActivity(Intent intent, ActivityTransitionAnimator.Controller controller) {
         // Media Output dialog can be shown from the volume panel. This makes sure the panel is
         // closed when navigating to another activity, so it doesn't stays on top of it
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 86063ac..2715cb3 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
@@ -1512,6 +1512,60 @@
         assertThat(getNumberOfConnectDeviceButtons()).isEqualTo(1);
     }
 
+    @EnableFlags(Flags.FLAG_ENABLE_OUTPUT_SWITCHER_SESSION_GROUPING)
+    @Test
+    public void selectedDevicesAddedInSameOrder() {
+        when(mLocalMediaManager.isPreferenceRouteListingExist()).thenReturn(true);
+        doReturn(mMediaDevices)
+                .when(mLocalMediaManager)
+                .getSelectedMediaDevice();
+        mMediaSwitchingController.start(mCb);
+        reset(mCb);
+        mMediaSwitchingController.getMediaItemList().clear();
+
+        mMediaSwitchingController.onDeviceListUpdate(mMediaDevices);
+
+        List<MediaItem> items = mMediaSwitchingController.getMediaItemList();
+        assertThat(items.get(0).getMediaDevice().get()).isEqualTo(mMediaDevice1);
+        assertThat(items.get(1).getMediaDevice().get()).isEqualTo(mMediaDevice2);
+    }
+
+    @DisableFlags(Flags.FLAG_ENABLE_OUTPUT_SWITCHER_SESSION_GROUPING)
+    @Test
+    public void selectedDevicesAddedInReverseOrder() {
+        when(mLocalMediaManager.isPreferenceRouteListingExist()).thenReturn(true);
+        doReturn(mMediaDevices)
+                .when(mLocalMediaManager)
+                .getSelectedMediaDevice();
+        mMediaSwitchingController.start(mCb);
+        reset(mCb);
+        mMediaSwitchingController.getMediaItemList().clear();
+
+        mMediaSwitchingController.onDeviceListUpdate(mMediaDevices);
+
+        List<MediaItem> items = mMediaSwitchingController.getMediaItemList();
+        assertThat(items.get(0).getMediaDevice().get()).isEqualTo(mMediaDevice2);
+        assertThat(items.get(1).getMediaDevice().get()).isEqualTo(mMediaDevice1);
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_OUTPUT_SWITCHER_SESSION_GROUPING)
+    @Test
+    public void firstSelectedDeviceIsFirstDeviceInGroupIsTrue() {
+        when(mLocalMediaManager.isPreferenceRouteListingExist()).thenReturn(true);
+        doReturn(mMediaDevices)
+                .when(mLocalMediaManager)
+                .getSelectedMediaDevice();
+        mMediaSwitchingController.start(mCb);
+        reset(mCb);
+        mMediaSwitchingController.getMediaItemList().clear();
+
+        mMediaSwitchingController.onDeviceListUpdate(mMediaDevices);
+
+        List<MediaItem> items = mMediaSwitchingController.getMediaItemList();
+        assertThat(items.get(0).isFirstDeviceInGroup()).isTrue();
+        assertThat(items.get(1).isFirstDeviceInGroup()).isFalse();
+    }
+
     private int getNumberOfConnectDeviceButtons() {
         int numberOfConnectDeviceButtons = 0;
         for (MediaItem item : mMediaSwitchingController.getMediaItemList()) {