diff --git a/src/com/android/settings/media/MediaDeviceUpdateWorker.java b/src/com/android/settings/media/MediaDeviceUpdateWorker.java
index 52249ffdbac..92cc0e37003 100644
--- a/src/com/android/settings/media/MediaDeviceUpdateWorker.java
+++ b/src/com/android/settings/media/MediaDeviceUpdateWorker.java
@@ -174,6 +174,15 @@ public class MediaDeviceUpdateWorker extends SliceBackgroundWorker
         return mLocalMediaManager.getSelectedMediaDevice();
     }
 
+    boolean isSelectedMediaDevice(MediaDevice device) {
+        for (MediaDevice selectedMediaDevice : getSelectedMediaDevice()) {
+            if (TextUtils.equals(selectedMediaDevice.getId(), device.getId())) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     void adjustSessionVolume(String sessionId, int volume) {
         mLocalMediaManager.adjustSessionVolume(sessionId, volume);
     }
diff --git a/src/com/android/settings/media/MediaOutputGroupSlice.java b/src/com/android/settings/media/MediaOutputGroupSlice.java
index be6c2984867..64ed8177e8e 100644
--- a/src/com/android/settings/media/MediaOutputGroupSlice.java
+++ b/src/com/android/settings/media/MediaOutputGroupSlice.java
@@ -238,7 +238,7 @@ public class MediaOutputGroupSlice implements CustomSliceable {
                             + ") is unavailable");
                     return;
                 }
-                if (TextUtils.equals(device.getClientPackageName(), getWorker().getPackageName())) {
+                if (getWorker().isSelectedMediaDevice(device)) {
                     getWorker().removeDeviceFromPlayMedia(device);
                 } else {
                     getWorker().addDeviceToPlayMedia(device);
diff --git a/tests/robotests/src/com/android/settings/media/MediaOutputGroupSliceTest.java b/tests/robotests/src/com/android/settings/media/MediaOutputGroupSliceTest.java
index 12e066bb485..9da18d50cde 100644
--- a/tests/robotests/src/com/android/settings/media/MediaOutputGroupSliceTest.java
+++ b/tests/robotests/src/com/android/settings/media/MediaOutputGroupSliceTest.java
@@ -212,11 +212,9 @@ public class MediaOutputGroupSliceTest {
     }
 
     @Test
-    public void onNotifyChange_sessionOperation_differentClient_verifyAddSession() {
-        mSelectableDevices.add(mDevice1);
+    public void onNotifyChange_sendSelectableDevice_verifyAddSession() {
         mSelectableDevices.add(mDevice2);
         mSelectedDevices.add(mDevice1);
-        when(mDevice2.getClientPackageName()).thenReturn(TEST_PACKAGE_NAME2);
         when(mLocalMediaManager.getMediaDeviceById(mSelectableDevices, TEST_DEVICE_2_ID))
                 .thenReturn(mDevice2);
         sMediaDeviceUpdateWorker.onDeviceListUpdate(mSelectableDevices);
@@ -229,16 +227,13 @@ public class MediaOutputGroupSliceTest {
 
         verify(sMediaDeviceUpdateWorker).addDeviceToPlayMedia(mDevice2);
     }
-
     @Test
-    public void onNotifyChange_sessionOperation_sameClient_verifyRemoveSession() {
-        mSelectableDevices.add(mDevice1);
-        mSelectableDevices.add(mDevice2);
+    public void onNotifyChange_sendSelectedDevice_verifyRemoveSession() {
         mSelectedDevices.add(mDevice1);
-        when(mDevice2.getClientPackageName()).thenReturn(TEST_PACKAGE_NAME);
-        when(mLocalMediaManager.getMediaDeviceById(mSelectableDevices, TEST_DEVICE_2_ID))
+        mSelectedDevices.add(mDevice2);
+        when(mLocalMediaManager.getMediaDeviceById(mSelectedDevices, TEST_DEVICE_2_ID))
                 .thenReturn(mDevice2);
-        sMediaDeviceUpdateWorker.onDeviceListUpdate(mSelectableDevices);
+        sMediaDeviceUpdateWorker.onDeviceListUpdate(mSelectedDevices);
         when(sMediaDeviceUpdateWorker.getSelectedMediaDevice()).thenReturn(mSelectedDevices);
         final Intent intent = new Intent();
         intent.putExtra(MEDIA_DEVICE_ID, TEST_DEVICE_2_ID);