Use BluetoothLeBroadcastAssistant#getSourceMetadata to retrieve broadcast name.

Test: atest
Bug: 381944659
Flag: com.android.settingslib.flags.enable_le_audio_sharing
Change-Id: I6e4c83a0858717727066de708fbde88e4b03ed8e

src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsHelper.java

@@ -21,6 +21,8 @@ import static com.android.settings.connecteddevice.audiosharing.audiostreams.Aud
 import static com.android.settings.connecteddevice.audiosharing.audiostreams.AudioStreamMediaService.DEVICES;
 
 import static java.util.Collections.emptyList;
+import static java.util.Collections.emptyMap;
+import static java.util.stream.Collectors.toMap;
 
 import android.bluetooth.BluetoothDevice;
 import android.bluetooth.BluetoothLeAudioContentMetadata;
@@ -48,7 +50,9 @@ import com.google.common.base.Strings;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
+import java.util.function.Function;
 import java.util.stream.Stream;
 
 import javax.annotation.Nullable;
@@ -149,6 +153,16 @@ public class AudioStreamsHelper {
                 .toList();
     }
 
+    /** Retrieves a list of all LE broadcast receive states keyed by each active device. */
+    public Map<BluetoothDevice, List<BluetoothLeBroadcastReceiveState>> getAllSourcesByDevice() {
+        if (mLeBroadcastAssistant == null) {
+            Log.w(TAG, "getAllSourcesByDevice(): LeBroadcastAssistant is null!");
+            return emptyMap();
+        }
+        return getConnectedBluetoothDevices(mBluetoothManager, /* inSharingOnly= */ true).stream()
+                .collect(toMap(Function.identity(), mLeBroadcastAssistant::getAllSources));
+    }
+
     /** Retrieves a list of all LE broadcast receive states from sinks with source present. */
     @VisibleForTesting
     public List<BluetoothLeBroadcastReceiveState> getAllPresentSources() {

src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsProgressCategoryCallback.java

@@ -16,7 +16,6 @@
 
 package com.android.settings.connecteddevice.audiosharing.audiostreams;
 
-
 import android.bluetooth.BluetoothDevice;
 import android.bluetooth.BluetoothLeBroadcastMetadata;
 import android.bluetooth.BluetoothLeBroadcastReceiveState;
@@ -43,13 +42,13 @@ public class AudioStreamsProgressCategoryCallback extends AudioStreamsBroadcastA
         super.onReceiveStateChanged(sink, sourceId, state);
 
         if (AudioStreamsHelper.isConnected(state)) {
-            mCategoryController.handleSourceConnected(state);
+            mCategoryController.handleSourceConnected(sink, state);
         } else if (AudioStreamsHelper.isBadCode(state)) {
             mCategoryController.handleSourceConnectBadCode(state);
         } else if (BluetoothUtils.isAudioSharingHysteresisModeFixAvailable(mContext)
                 && AudioStreamsHelper.hasSourcePresent(state)) {
             // Keep this check as the last, source might also present in above states
-            mCategoryController.handleSourcePresent(state);
+            mCategoryController.handleSourcePresent(sink, state);
         }
     }
 

src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsProgressCategoryController.java

@@ -17,10 +17,12 @@
 package com.android.settings.connecteddevice.audiosharing.audiostreams;
 
 import static java.util.Collections.emptyList;
+import static java.util.stream.Collectors.toMap;
 
 import android.app.AlertDialog;
 import android.app.settings.SettingsEnums;
 import android.bluetooth.BluetoothAdapter;
+import android.bluetooth.BluetoothDevice;
 import android.bluetooth.BluetoothLeBroadcastMetadata;
 import android.bluetooth.BluetoothLeBroadcastReceiveState;
 import android.bluetooth.BluetoothProfile;
@@ -49,6 +51,8 @@ import com.android.settingslib.utils.ThreadUtils;
 
 import java.util.Comparator;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
@@ -391,34 +395,19 @@ public class AudioStreamsProgressCategoryController extends BasePreferenceContro
     // Expect one of the following:
     // 1) No preference existed, create new preference with state SOURCE_ADDED
     // 2) Any other state, move to SOURCE_ADDED
-    void handleSourceConnected(BluetoothLeBroadcastReceiveState receiveState) {
+    void handleSourceConnected(
+            BluetoothDevice device, BluetoothLeBroadcastReceiveState receiveState) {
         if (DEBUG) {
             Log.d(TAG, "handleSourceConnected()");
         }
         if (!AudioStreamsHelper.isConnected(receiveState)) {
             return;
         }
-
         var broadcastIdConnected = receiveState.getBroadcastId();
-        if (mSourceFromQrCode != null && mSourceFromQrCode.getBroadcastId() == UNSET_BROADCAST_ID) {
+        Optional<BluetoothLeBroadcastMetadata> metadata =
-            // mSourceFromQrCode could have no broadcast Id, we fill in the broadcast Id from the
+                getMetadataMatchingByBroadcastId(
-            // connected source receiveState.
+                        device, receiveState.getSourceId(), broadcastIdConnected);
-            if (DEBUG) {
+        handleQrCodeWithUnsetBroadcastIdIfNeeded(metadata, receiveState);
-                Log.d(
-                        TAG,
-                        "handleSourceConnected() : processing mSourceFromQrCode with broadcastId"
-                                + " unset");
-            }
-            boolean updated =
-                    maybeUpdateId(
-                            AudioStreamsHelper.getBroadcastName(receiveState),
-                            receiveState.getBroadcastId());
-            if (updated && mBroadcastIdToPreferenceMap.containsKey(UNSET_BROADCAST_ID)) {
-                var preference = mBroadcastIdToPreferenceMap.remove(UNSET_BROADCAST_ID);
-                mBroadcastIdToPreferenceMap.put(receiveState.getBroadcastId(), preference);
-            }
-        }
-
         mBroadcastIdToPreferenceMap.compute(
                 broadcastIdConnected,
                 (k, existingPreference) -> {
@@ -428,7 +417,12 @@ public class AudioStreamsProgressCategoryController extends BasePreferenceContro
                         // we retrieves the connected source during onStart() from
                         // AudioStreamsHelper#getAllConnectedSources() even before the source is
                         // founded by scanning.
-                        return addNewPreference(receiveState, AudioStreamState.SOURCE_ADDED);
+                        return metadata.isPresent()
+                                ? addNewPreference(
+                                        metadata.get(),
+                                        AudioStreamState.SOURCE_ADDED,
+                                        SourceOriginForLogging.UNKNOWN)
+                                : addNewPreference(receiveState, AudioStreamState.SOURCE_ADDED);
                     }
                     if (existingPreference.getAudioStreamState() == AudioStreamState.WAIT_FOR_SYNC
                             && existingPreference.getAudioStreamBroadcastId() == UNSET_BROADCAST_ID
@@ -473,7 +467,8 @@ public class AudioStreamsProgressCategoryController extends BasePreferenceContro
 
     // Find preference by receiveState and decide next state.
     // Expect one preference existed, move to SOURCE_PRESENT
-    void handleSourcePresent(BluetoothLeBroadcastReceiveState receiveState) {
+    void handleSourcePresent(
+            BluetoothDevice device, BluetoothLeBroadcastReceiveState receiveState) {
         if (DEBUG) {
             Log.d(TAG, "handleSourcePresent()");
         }
@@ -482,25 +477,10 @@ public class AudioStreamsProgressCategoryController extends BasePreferenceContro
         }
 
         var broadcastIdConnected = receiveState.getBroadcastId();
-        if (mSourceFromQrCode != null && mSourceFromQrCode.getBroadcastId() == UNSET_BROADCAST_ID) {
+        Optional<BluetoothLeBroadcastMetadata> metadata =
-            // mSourceFromQrCode could have no broadcast Id, we fill in the broadcast Id from the
+                getMetadataMatchingByBroadcastId(
-            // connected source receiveState.
+                        device, receiveState.getSourceId(), broadcastIdConnected);
-            if (DEBUG) {
+        handleQrCodeWithUnsetBroadcastIdIfNeeded(metadata, receiveState);
-                Log.d(
-                        TAG,
-                        "handleSourcePresent() : processing mSourceFromQrCode with broadcastId"
-                                + " unset");
-            }
-            boolean updated =
-                    maybeUpdateId(
-                            AudioStreamsHelper.getBroadcastName(receiveState),
-                            receiveState.getBroadcastId());
-            if (updated && mBroadcastIdToPreferenceMap.containsKey(UNSET_BROADCAST_ID)) {
-                var preference = mBroadcastIdToPreferenceMap.remove(UNSET_BROADCAST_ID);
-                mBroadcastIdToPreferenceMap.put(receiveState.getBroadcastId(), preference);
-            }
-        }
-
         mBroadcastIdToPreferenceMap.compute(
                 broadcastIdConnected,
                 (k, existingPreference) -> {
@@ -511,7 +491,12 @@ public class AudioStreamsProgressCategoryController extends BasePreferenceContro
                         // we retrieves the connected source during onStart() from
                         // AudioStreamsHelper#getAllPresentSources() even before the source is
                         // founded by scanning.
-                        return addNewPreference(receiveState, AudioStreamState.SOURCE_PRESENT);
+                        return metadata.isPresent()
+                                ? addNewPreference(
+                                        metadata.get(),
+                                        AudioStreamState.SOURCE_PRESENT,
+                                        SourceOriginForLogging.UNKNOWN)
+                                : addNewPreference(receiveState, AudioStreamState.SOURCE_PRESENT);
                     }
                     if (existingPreference.getAudioStreamState() == AudioStreamState.WAIT_FOR_SYNC
                             && existingPreference.getAudioStreamBroadcastId() == UNSET_BROADCAST_ID
@@ -598,28 +583,85 @@ public class AudioStreamsProgressCategoryController extends BasePreferenceContro
                     // Handle QR code scan, display currently connected streams then start scanning
                     // sequentially
                     handleSourceFromQrCodeIfExists();
+                    Map<BluetoothDevice, List<BluetoothLeBroadcastReceiveState>> sources =
+                            mAudioStreamsHelper.getAllSourcesByDevice();
+                    Map<BluetoothDevice, List<BluetoothLeBroadcastReceiveState>> connectedSources =
+                            getConnectedSources(sources);
                     if (isAudioSharingHysteresisModeFixAvailable(mContext)) {
                         // With hysteresis mode, we prioritize showing connected sources first.
                         // If no connected sources are found, we then show present sources.
-                        List<BluetoothLeBroadcastReceiveState> sources =
+                        if (!connectedSources.isEmpty()) {
-                                mAudioStreamsHelper.getAllConnectedSources();
+                            connectedSources.forEach(
-                        if (!sources.isEmpty()) {
+                                    (device, stateList) ->
-                            sources.forEach(this::handleSourceConnected);
+                                            stateList.forEach(
+                                                    state -> handleSourceConnected(device, state)));
                         } else {
-                            mAudioStreamsHelper
+                            Map<BluetoothDevice, List<BluetoothLeBroadcastReceiveState>>
-                                    .getAllPresentSources()
+                                    presentSources = getPresentSources(sources);
-                                    .forEach(this::handleSourcePresent);
+                            presentSources.forEach(
+                                    (device, stateList) ->
+                                            stateList.forEach(
+                                                    state -> handleSourcePresent(device, state)));
                         }
                     } else {
-                        mAudioStreamsHelper
+                        connectedSources.forEach(
-                                .getAllConnectedSources()
+                                (device, stateList) ->
-                                .forEach(this::handleSourceConnected);
+                                        stateList.forEach(
+                                                state -> handleSourceConnected(device, state)));
                     }
                     mLeBroadcastAssistant.startSearchingForSources(emptyList());
                     mMediaControlHelper.start();
                 });
     }
 
+    private Map<BluetoothDevice, List<BluetoothLeBroadcastReceiveState>> getConnectedSources(
+            Map<BluetoothDevice, List<BluetoothLeBroadcastReceiveState>> sources) {
+        return sources.entrySet().stream()
+                .filter(
+                        entry ->
+                                entry.getValue().stream().anyMatch(AudioStreamsHelper::isConnected))
+                .collect(toMap(Map.Entry::getKey, Map.Entry::getValue));
+    }
+
+    private Map<BluetoothDevice, List<BluetoothLeBroadcastReceiveState>> getPresentSources(
+            Map<BluetoothDevice, List<BluetoothLeBroadcastReceiveState>> sources) {
+        return sources.entrySet().stream()
+                .filter(
+                        entry ->
+                                entry.getValue().stream()
+                                        .anyMatch(AudioStreamsHelper::hasSourcePresent))
+                .collect(toMap(Map.Entry::getKey, Map.Entry::getValue));
+    }
+
+    private Optional<BluetoothLeBroadcastMetadata> getMetadataMatchingByBroadcastId(
+            BluetoothDevice device, int sourceId, int broadcastId) {
+        return Optional.ofNullable(
+                        mLeBroadcastAssistant != null
+                                ? mLeBroadcastAssistant.getSourceMetadata(device, sourceId)
+                                : null)
+                .filter(m -> m.getBroadcastId() == broadcastId);
+    }
+
+    private void handleQrCodeWithUnsetBroadcastIdIfNeeded(
+            Optional<BluetoothLeBroadcastMetadata> metadata,
+            BluetoothLeBroadcastReceiveState receiveState) {
+        if (mSourceFromQrCode != null && mSourceFromQrCode.getBroadcastId() == UNSET_BROADCAST_ID) {
+            if (DEBUG) {
+                Log.d(TAG, "Processing mSourceFromQrCode with unset broadcastId");
+            }
+            boolean updated =
+                    maybeUpdateId(
+                            metadata.isPresent()
+                                    ? AudioStreamsHelper.getBroadcastName(metadata.get())
+                                    : AudioStreamsHelper.getBroadcastName(receiveState),
+                            receiveState.getBroadcastId());
+            if (updated && mBroadcastIdToPreferenceMap.containsKey(UNSET_BROADCAST_ID)) {
+                var preference = mBroadcastIdToPreferenceMap.remove(UNSET_BROADCAST_ID);
+                mBroadcastIdToPreferenceMap.put(receiveState.getBroadcastId(), preference);
+            }
+        }
+    }
+
     private void stopScanning() {
         if (mLeBroadcastAssistant == null) {
             Log.w(TAG, "stopScanning(): LeBroadcastAssistant is null!");

tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsProgressCategoryCallbackTest.java

@@ -54,7 +54,7 @@ import java.util.List;
 @RunWith(RobolectricTestRunner.class)
 @Config(
         shadows = {
-                ShadowBluetoothAdapter.class,
+            ShadowBluetoothAdapter.class,
         })
 public class AudioStreamsProgressCategoryCallbackTest {
     @Rule public final MockitoRule mMockitoRule = MockitoJUnit.rule();
@@ -70,8 +70,8 @@ public class AudioStreamsProgressCategoryCallbackTest {
     @Before
     public void setUp() {
         mSetFlagsRule.disableFlags(FLAG_AUDIO_SHARING_HYSTERESIS_MODE_FIX);
-        ShadowBluetoothAdapter shadowBluetoothAdapter = Shadow.extract(
+        ShadowBluetoothAdapter shadowBluetoothAdapter =
-                BluetoothAdapter.getDefaultAdapter());
+                Shadow.extract(BluetoothAdapter.getDefaultAdapter());
         shadowBluetoothAdapter.setEnabled(true);
         shadowBluetoothAdapter.setIsLeAudioBroadcastSourceSupported(
                 BluetoothStatusCodes.FEATURE_SUPPORTED);
@@ -87,7 +87,7 @@ public class AudioStreamsProgressCategoryCallbackTest {
         when(mState.getBisSyncState()).thenReturn(bisSyncState);
         mCallback.onReceiveStateChanged(mDevice, /* sourceId= */ 0, mState);
 
-        verify(mController).handleSourceConnected(any());
+        verify(mController).handleSourceConnected(any(), any());
     }
 
     @Test
@@ -102,7 +102,7 @@ public class AudioStreamsProgressCategoryCallbackTest {
         when(mSourceDevice.getAddress()).thenReturn(address);
         mCallback.onReceiveStateChanged(mDevice, /* sourceId= */ 0, mState);
 
-        verify(mController).handleSourcePresent(any());
+        verify(mController).handleSourcePresent(any(), any());
     }
 
     @Test

tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsProgressCategoryControllerTest.java

@@ -32,6 +32,7 @@ import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -41,7 +42,7 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static org.robolectric.Shadows.shadowOf;
 
-import static java.util.Collections.emptyList;
+import static java.util.Collections.emptyMap;
 
 import android.bluetooth.BluetoothAdapter;
 import android.bluetooth.BluetoothDevice;
@@ -93,6 +94,7 @@ import org.robolectric.shadows.androidx.fragment.FragmentController;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 
 @RunWith(RobolectricTestRunner.class)
 @Config(
@@ -134,8 +136,8 @@ public class AudioStreamsProgressCategoryControllerTest {
 
     @Before
     public void setUp() {
-        ShadowBluetoothAdapter shadowBluetoothAdapter = Shadow.extract(
+        ShadowBluetoothAdapter shadowBluetoothAdapter =
-                BluetoothAdapter.getDefaultAdapter());
+                Shadow.extract(BluetoothAdapter.getDefaultAdapter());
         shadowBluetoothAdapter.setEnabled(true);
         shadowBluetoothAdapter.setIsLeAudioBroadcastSourceSupported(
                 BluetoothStatusCodes.FEATURE_SUPPORTED);
@@ -143,7 +145,7 @@ public class AudioStreamsProgressCategoryControllerTest {
                 BluetoothStatusCodes.FEATURE_SUPPORTED);
         ShadowAudioStreamsHelper.setUseMock(mAudioStreamsHelper);
         when(mAudioStreamsHelper.getLeBroadcastAssistant()).thenReturn(mLeBroadcastAssistant);
-        when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(emptyList());
+        when(mAudioStreamsHelper.getAllSourcesByDevice()).thenReturn(emptyMap());
         mSetFlagsRule.disableFlags(FLAG_AUDIO_SHARING_HYSTERESIS_MODE_FIX);
 
         ShadowBluetoothUtils.sLocalBluetoothManager = mLocalBtManager;
@@ -310,14 +312,12 @@ public class AudioStreamsProgressCategoryControllerTest {
         // Setup a device
         ShadowAudioStreamsHelper.setCachedBluetoothDeviceInSharingOrLeConnected(mDevice);
 
-        List<BluetoothLeBroadcastReceiveState> connectedList = new ArrayList<>();
         // Empty connected device list
-        when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(connectedList);
+        when(mAudioStreamsHelper.getAllSourcesByDevice()).thenReturn(emptyMap());
 
         mController.onStart(mLifecycleOwner);
         shadowOf(Looper.getMainLooper()).idle();
 
-        verify(mAudioStreamsHelper).getAllPresentSources();
         verify(mLeBroadcastAssistant).startSearchingForSources(any());
 
         var dialog = ShadowAlertDialog.getLatestAlertDialog();
@@ -355,7 +355,7 @@ public class AudioStreamsProgressCategoryControllerTest {
     }
 
     @Test
-    public void testOnStart_handleSourceAlreadyConnected() {
+    public void testOnStart_handleSourceAlreadyConnected_useNameFromMetadata() {
         // Setup a device
         ShadowAudioStreamsHelper.setCachedBluetoothDeviceInSharingOrLeConnected(mDevice);
 
@@ -363,8 +363,14 @@ public class AudioStreamsProgressCategoryControllerTest {
         BluetoothLeBroadcastReceiveState connected =
                 createConnectedMock(ALREADY_CONNECTED_BROADCAST_ID);
         List<BluetoothLeBroadcastReceiveState> list = new ArrayList<>();
+        var data = mock(BluetoothLeAudioContentMetadata.class);
+        when(connected.getSubgroupMetadata()).thenReturn(ImmutableList.of(data));
+        when(data.getProgramInfo()).thenReturn(BROADCAST_NAME_1);
         list.add(connected);
-        when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(list);
+        when(mMetadata.getBroadcastId()).thenReturn(ALREADY_CONNECTED_BROADCAST_ID);
+        when(mMetadata.getBroadcastName()).thenReturn(BROADCAST_NAME_2);
+        when(mLeBroadcastAssistant.getSourceMetadata(any(), anyInt())).thenReturn(mMetadata);
+        when(mAudioStreamsHelper.getAllSourcesByDevice()).thenReturn(Map.of(mSourceDevice, list));
 
         // Handle already connected source in onStart
         mController.displayPreference(mScreen);
@@ -382,6 +388,7 @@ public class AudioStreamsProgressCategoryControllerTest {
         assertThat(preference.getValue()).isNotNull();
         assertThat(preference.getValue().getAudioStreamBroadcastId())
                 .isEqualTo(ALREADY_CONNECTED_BROADCAST_ID);
+        assertThat(preference.getValue().getTitle()).isEqualTo(BROADCAST_NAME_2);
         assertThat(state.getValue()).isEqualTo(SOURCE_ADDED);
     }
 
@@ -409,7 +416,8 @@ public class AudioStreamsProgressCategoryControllerTest {
         var data = mock(BluetoothLeAudioContentMetadata.class);
         when(connected.getSubgroupMetadata()).thenReturn(ImmutableList.of(data));
         when(data.getProgramInfo()).thenReturn(BROADCAST_NAME_1);
-        when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(ImmutableList.of(connected));
+        when(mAudioStreamsHelper.getAllSourcesByDevice())
+                .thenReturn(Map.of(mSourceDevice, ImmutableList.of(connected)));
 
         // Handle both source from qr code and already connected source in onStart
         mController.displayPreference(mScreen);
@@ -578,8 +586,8 @@ public class AudioStreamsProgressCategoryControllerTest {
         // Setup source already connected
         BluetoothLeBroadcastReceiveState connected =
                 createConnectedMock(ALREADY_CONNECTED_BROADCAST_ID);
-        when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(ImmutableList.of(connected));
+        when(mAudioStreamsHelper.getAllSourcesByDevice())
-
+                .thenReturn(Map.of(mSourceDevice, List.of(connected)));
         // Handle source already connected in onStart
         mController.displayPreference(mScreen);
         mController.onStart(mLifecycleOwner);
@@ -687,7 +695,8 @@ public class AudioStreamsProgressCategoryControllerTest {
         // Setup already connected source
         BluetoothLeBroadcastReceiveState connected =
                 createConnectedMock(ALREADY_CONNECTED_BROADCAST_ID);
-        when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(ImmutableList.of(connected));
+        when(mAudioStreamsHelper.getAllSourcesByDevice())
+                .thenReturn(Map.of(mSourceDevice, List.of(connected)));
 
         // Handle connected source in onStart
         mController.displayPreference(mScreen);
@@ -695,7 +704,7 @@ public class AudioStreamsProgressCategoryControllerTest {
         shadowOf(Looper.getMainLooper()).idle();
 
         // The connect source is no longer connected
-        when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(emptyList());
+        when(mAudioStreamsHelper.getAllSourcesByDevice()).thenReturn(emptyMap());
         mController.handleSourceRemoved();
         shadowOf(Looper.getMainLooper()).idle();
 
@@ -728,7 +737,8 @@ public class AudioStreamsProgressCategoryControllerTest {
         // Setup a connected source
         BluetoothLeBroadcastReceiveState connected =
                 createConnectedMock(ALREADY_CONNECTED_BROADCAST_ID);
-        when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(ImmutableList.of(connected));
+        when(mAudioStreamsHelper.getAllSourcesByDevice())
+                .thenReturn(Map.of(mSourceDevice, List.of(connected)));
 
         // Handle connected source in onStart
         mController.displayPreference(mScreen);
@@ -834,7 +844,7 @@ public class AudioStreamsProgressCategoryControllerTest {
         when(receiveState.getBisSyncState()).thenReturn(bisSyncState);
 
         // The new found source is identified as failed to connect
-        mController.handleSourcePresent(receiveState);
+        mController.handleSourcePresent(mSourceDevice, receiveState);
         shadowOf(Looper.getMainLooper()).idle();
 
         ArgumentCaptor<AudioStreamPreference> preference =

tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/testshadows/ShadowAudioStreamsHelper.java

@@ -16,6 +16,7 @@
 
 package com.android.settings.connecteddevice.audiosharing.audiostreams.testshadows;
 
+import android.bluetooth.BluetoothDevice;
 import android.bluetooth.BluetoothLeBroadcastMetadata;
 import android.bluetooth.BluetoothLeBroadcastReceiveState;
 
@@ -31,6 +32,7 @@ import org.robolectric.annotation.Implements;
 import org.robolectric.annotation.Resetter;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
 @Implements(value = AudioStreamsHelper.class, callThroughByDefault = true)
@@ -59,6 +61,11 @@ public class ShadowAudioStreamsHelper {
         return sMockHelper.getAllConnectedSources();
     }
 
+    @Implementation
+    public Map<BluetoothDevice, List<BluetoothLeBroadcastReceiveState>> getAllSourcesByDevice() {
+        return sMockHelper.getAllSourcesByDevice();
+    }
+
     @Implementation
     public List<BluetoothLeBroadcastReceiveState> getAllPresentSources() {
         return sMockHelper.getAllPresentSources();