2 files changed, 15 insertions, 2 deletions
diff --git a/packages/SystemUI/src/com/android/systemui/media/MediaDataCombineLatest.kt b/packages/SystemUI/src/com/android/systemui/media/MediaDataCombineLatest.kt
index 0904ebccd414..e8f0e069c98d 100644
--- a/packages/SystemUI/src/com/android/systemui/media/MediaDataCombineLatest.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/MediaDataCombineLatest.kt
@@ -62,9 +62,9 @@ class MediaDataCombineLatest @Inject constructor(
      */
     fun getData(): Map<String, MediaData> {
         return entries.filter {
-            (key, pair) -> pair.first != null
+            (key, pair) -> pair.first != null && pair.second != null
         }.mapValues {
-            (key, pair) -> pair.first!!
+            (key, pair) -> pair.first!!.copy(device = pair.second)
         }
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/MediaDataCombineLatestTest.java b/packages/SystemUI/tests/src/com/android/systemui/media/MediaDataCombineLatestTest.java
index 5d4693d3ccf8..dbc5596d9f4e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/MediaDataCombineLatestTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/MediaDataCombineLatestTest.java
@@ -39,6 +39,7 @@ import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 
 import java.util.ArrayList;
+import java.util.Map;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
@@ -159,6 +160,18 @@ public class MediaDataCombineLatestTest extends SysuiTestCase {
         verify(mListener).onMediaDataLoaded(eq("NEW_KEY"), any(), captor.capture());
     }
 
+    @Test
+    public void getDataIncludesDevice() {
+        // GIVEN that device and media events have been received
+        mDeviceListener.onMediaDeviceChanged(KEY, mDeviceData);
+        mDataListener.onMediaDataLoaded(KEY, null, mMediaData);
+
+        // THEN the result of getData includes device info
+        Map<String, MediaData> results = mManager.getData();
+        assertThat(results.get(KEY)).isNotNull();
+        assertThat(results.get(KEY).getDevice()).isEqualTo(mDeviceData);
+    }
+
     private MediaDataManager.Listener captureDataListener() {
         ArgumentCaptor<MediaDataManager.Listener> captor = ArgumentCaptor.forClass(
                 MediaDataManager.Listener.class);