Show status bar notification icons on all displays

Since Android Views can only have one parent, we need to duplicate
the StatusBarIconView for each display.

The chosen approach is to create a multi display version of the
NotificationIconViewStore, that asks IconManager to create and update
icons.

We still keep the icons for the default display in the same place as
before (IconPack in NotificationEntry), to avoid having to refactor
things too much.

Test: MultiDisplayStatusBarNotificationIconViewStoreTest
Test: More to come when approach is approved
Bug: 369337701
Flag: com.android.systemui.status_bar_connected_displays
Change-Id: I779b2d39939e44d0544c496d46fc148e9fb65c36

10 files changed, 345 insertions, 24 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/domain/interactor/ShadeDisplaysInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/domain/interactor/ShadeDisplaysInteractorTest.kt
index 982c51b8318c..80cf2f04f035 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/domain/interactor/ShadeDisplaysInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/domain/interactor/ShadeDisplaysInteractorTest.kt
@@ -53,7 +53,7 @@ class ShadeDisplaysInteractorTest : SysuiTestCase() {
     private val shadeRootview = mock<WindowRootView>()
     private val positionRepository = FakeShadeDisplayRepository()
     private val shadeContext = mock<Context>()
-    private val contextStore = FakeDisplayWindowPropertiesRepository()
+    private val contextStore = FakeDisplayWindowPropertiesRepository(context)
     private val testScope = TestScope(UnconfinedTestDispatcher())
     private val shadeWm = mock<WindowManager>()
     private val resources = mock<Resources>()
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/icon/ui/viewbinder/ConnectedDisplaysStatusBarNotificationIconViewStoreTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/icon/ui/viewbinder/ConnectedDisplaysStatusBarNotificationIconViewStoreTest.kt
new file mode 100644
index 000000000000..483c2be21fc7
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/icon/ui/viewbinder/ConnectedDisplaysStatusBarNotificationIconViewStoreTest.kt
@@ -0,0 +1,149 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+package com.android.systemui.statusbar.notification.icon.ui.viewbinder
+
+import android.app.Notification
+import android.app.NotificationChannel
+import android.app.NotificationManager
+import android.platform.test.annotations.EnableFlags
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.display.domain.interactor.displayWindowPropertiesInteractor
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.lifecycle.activateIn
+import com.android.systemui.statusbar.RankingBuilder
+import com.android.systemui.statusbar.SbnBuilder
+import com.android.systemui.statusbar.core.StatusBarConnectedDisplays
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+import com.android.systemui.statusbar.notification.collection.notifCollection
+import com.android.systemui.statusbar.notification.collection.notifPipeline
+import com.android.systemui.statusbar.notification.collection.notifcollection.NotifCollectionListener
+import com.android.systemui.statusbar.notification.icon.iconManager
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.whenever
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+@EnableFlags(StatusBarConnectedDisplays.FLAG_NAME)
+class ConnectedDisplaysStatusBarNotificationIconViewStoreTest : SysuiTestCase() {
+
+    private val kosmos = testKosmos()
+
+    private val underTest =
+        ConnectedDisplaysStatusBarNotificationIconViewStore(
+            TEST_DISPLAY_ID,
+            kosmos.notifCollection,
+            kosmos.iconManager,
+            kosmos.displayWindowPropertiesInteractor,
+            kosmos.notifPipeline,
+        )
+
+    private val notifCollectionListeners = mutableListOf<NotifCollectionListener>()
+
+    @Before
+    fun setupNoticCollectionListener() {
+        whenever(kosmos.notifPipeline.addCollectionListener(any())).thenAnswer { invocation ->
+            notifCollectionListeners.add(invocation.arguments[0] as NotifCollectionListener)
+        }
+    }
+
+    @Before
+    fun activate() {
+        underTest.activateIn(kosmos.testScope)
+    }
+
+    @Test
+    fun iconView_unknownKey_returnsNull() =
+        kosmos.testScope.runTest {
+            val unknownKey = "unknown key"
+
+            assertThat(underTest.iconView(unknownKey)).isNull()
+        }
+
+    @Test
+    fun iconView_knownKey_returnsNonNull() =
+        kosmos.testScope.runTest {
+            val entry = createEntry()
+
+            whenever(kosmos.notifCollection.getEntry(entry.key)).thenReturn(entry)
+
+            assertThat(underTest.iconView(entry.key)).isNotNull()
+        }
+
+    @Test
+    fun iconView_knownKey_calledMultipleTimes_returnsSameInstance() =
+        kosmos.testScope.runTest {
+            val entry = createEntry()
+
+            whenever(kosmos.notifCollection.getEntry(entry.key)).thenReturn(entry)
+
+            val first = underTest.iconView(entry.key)
+            val second = underTest.iconView(entry.key)
+
+            assertThat(first).isSameInstanceAs(second)
+        }
+
+    @Test
+    fun iconView_knownKey_afterNotificationRemoved_returnsNewInstance() =
+        kosmos.testScope.runTest {
+            val entry = createEntry()
+
+            whenever(kosmos.notifCollection.getEntry(entry.key)).thenReturn(entry)
+
+            val first = underTest.iconView(entry.key)
+
+            notifCollectionListeners.forEach { it.onEntryRemoved(entry, /* reason= */ 0) }
+
+            val second = underTest.iconView(entry.key)
+
+            assertThat(first).isNotSameInstanceAs(second)
+        }
+
+    private fun createEntry(): NotificationEntry {
+        val channelId = "channelId"
+        val notificationChannel =
+            NotificationChannel(channelId, "name", NotificationManager.IMPORTANCE_DEFAULT)
+        val notification =
+            Notification.Builder(context, channelId)
+                .setContentTitle("Title")
+                .setContentText("Content text")
+                .setSmallIcon(com.android.systemui.res.R.drawable.icon)
+                .build()
+        val statusBarNotification = SbnBuilder().setNotification(notification).build()
+        val ranking =
+            RankingBuilder()
+                .setChannel(notificationChannel)
+                .setKey(statusBarNotification.key)
+                .build()
+        return NotificationEntry(
+            /* sbn = */ statusBarNotification,
+            /* ranking = */ ranking,
+            /* creationTime = */ 1234L,
+        )
+    }
+
+    private companion object {
+        const val TEST_DISPLAY_ID = 1234
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/IconBuilder.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/IconBuilder.kt
index 3c8c42f6b29d..0f19d7288f6f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/IconBuilder.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/IconBuilder.kt
@@ -25,7 +25,11 @@ import javax.inject.Inject
 
 /** Testable wrapper around Context. */
 class IconBuilder @Inject constructor(private val context: Context) {
-    fun createIconView(entry: NotificationEntry): StatusBarIconView {
+    @JvmOverloads
+    fun createIconView(
+        entry: NotificationEntry,
+        context: Context = this.context,
+    ): StatusBarIconView {
         return StatusBarIconView(
             context,
             "${entry.sbn.packageName}/0x${Integer.toHexString(entry.sbn.id)}",
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/IconManager.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/IconManager.kt
index 47171948f395..98ce163b81ca 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/IconManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/IconManager.kt
@@ -19,6 +19,7 @@ package com.android.systemui.statusbar.notification.icon
 import android.app.Notification
 import android.app.Notification.MessagingStyle
 import android.app.Person
+import android.content.Context
 import android.content.pm.LauncherApps
 import android.graphics.drawable.Icon
 import android.os.Build
@@ -36,6 +37,7 @@ import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.StatusBarIconView
+import com.android.systemui.statusbar.core.StatusBarConnectedDisplays
 import com.android.systemui.statusbar.notification.InflationException
 import com.android.systemui.statusbar.notification.collection.NotificationEntry
 import com.android.systemui.statusbar.notification.collection.notifcollection.CommonNotifCollection
@@ -68,6 +70,17 @@ constructor(
     @Background private val bgCoroutineContext: CoroutineContext,
     @Main private val mainCoroutineContext: CoroutineContext,
 ) : ConversationIconManager {
+
+    /**
+     * A listener that is notified when a [NotificationEntry] has been updated and the associated
+     * icons have to be updated as well.
+     */
+    fun interface OnIconUpdateRequiredListener {
+        fun onIconUpdateRequired(entry: NotificationEntry)
+    }
+
+    private val onIconUpdateRequiredListeners = mutableSetOf<OnIconUpdateRequiredListener>()
+
     private var unimportantConversationKeys: Set<String> = emptySet()
     /**
      * A map of running jobs for fetching the person avatar from launcher. The key is the
@@ -76,6 +89,16 @@ constructor(
     private var launcherPeopleAvatarIconJobs: ConcurrentHashMap<String, Job> =
         ConcurrentHashMap<String, Job>()
 
+    fun addIconsUpdateListener(listener: OnIconUpdateRequiredListener) {
+        StatusBarConnectedDisplays.assertInNewMode()
+        onIconUpdateRequiredListeners += listener
+    }
+
+    fun removeIconsUpdateListener(listener: OnIconUpdateRequiredListener) {
+        StatusBarConnectedDisplays.assertInNewMode()
+        onIconUpdateRequiredListeners -= listener
+    }
+
     fun attach() {
         notifCollection.addCollectionListener(entryListener)
     }
@@ -112,6 +135,21 @@ constructor(
     }
 
     /**
+     * Inflate the [StatusBarIconView] for the given [NotificationEntry], using the specified
+     * [Context].
+     */
+    fun createSbIconView(context: Context, entry: NotificationEntry): StatusBarIconView =
+        traceSection("IconManager.createSbIconView") {
+            StatusBarConnectedDisplays.assertInNewMode()
+
+            val sbIcon = iconBuilder.createIconView(entry, context)
+            sbIcon.scaleType = ImageView.ScaleType.CENTER_INSIDE
+            val (normalIconDescriptor, _) = getIconDescriptors(entry)
+            setIcon(entry, normalIconDescriptor, sbIcon)
+            return sbIcon
+        }
+
+    /**
      * Inflate icon views for each icon variant and assign appropriate icons to them. Stores the
      * result in [NotificationEntry.getIcons].
      *
@@ -159,6 +197,18 @@ constructor(
             }
         }
 
+    /** Update the [StatusBarIconView] for the given [NotificationEntry]. */
+    fun updateSbIcon(entry: NotificationEntry, iconView: StatusBarIconView) =
+        traceSection("IconManager.updateSbIcon") {
+            StatusBarConnectedDisplays.assertInNewMode()
+
+            val (normalIconDescriptor, _) = getIconDescriptors(entry)
+            val notificationContentDescription =
+                entry.sbn.notification?.let { iconBuilder.getIconContentDescription(it) }
+            iconView.setNotification(entry.sbn, notificationContentDescription)
+            setIcon(entry, normalIconDescriptor, iconView)
+        }
+
     /**
      * Update the notification icons.
      *
@@ -172,6 +222,10 @@ constructor(
                 return@traceSection
             }
 
+            if (StatusBarConnectedDisplays.isEnabled) {
+                onIconUpdateRequiredListeners.onEach { it.onIconUpdateRequired(entry) }
+            }
+
             if (usingCache && !Flags.notificationsBackgroundIcons()) {
                 Log.wtf(
                     TAG,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewbinder/ConnectedDisplaysStatusBarNotificationIconViewStore.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewbinder/ConnectedDisplaysStatusBarNotificationIconViewStore.kt
new file mode 100644
index 000000000000..227a1fefb982
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewbinder/ConnectedDisplaysStatusBarNotificationIconViewStore.kt
@@ -0,0 +1,94 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+package com.android.systemui.statusbar.notification.icon.ui.viewbinder
+
+import com.android.systemui.display.domain.interactor.DisplayWindowPropertiesInteractor
+import com.android.systemui.lifecycle.Activatable
+import com.android.systemui.statusbar.StatusBarIconView
+import com.android.systemui.statusbar.notification.collection.NotifCollection
+import com.android.systemui.statusbar.notification.collection.NotifPipeline
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+import com.android.systemui.statusbar.notification.collection.notifcollection.NotifCollectionListener
+import com.android.systemui.statusbar.notification.icon.IconManager
+import com.android.systemui.statusbar.notification.icon.ui.viewbinder.NotificationIconContainerViewBinder.IconViewStore
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+import java.util.concurrent.ConcurrentHashMap
+import kotlinx.coroutines.awaitCancellation
+import kotlinx.coroutines.coroutineScope
+
+/** [IconViewStore] for the status bar on multiple displays. */
+class ConnectedDisplaysStatusBarNotificationIconViewStore
+@AssistedInject
+constructor(
+    @Assisted private val displayId: Int,
+    private val notifCollection: NotifCollection,
+    private val iconManager: IconManager,
+    private val displayWindowPropertiesInteractor: DisplayWindowPropertiesInteractor,
+    private val notifPipeline: NotifPipeline,
+) : IconViewStore, Activatable {
+
+    private val cachedIcons = ConcurrentHashMap<String, StatusBarIconView>()
+
+    private val iconUpdateRequiredListener =
+        object : IconManager.OnIconUpdateRequiredListener {
+            override fun onIconUpdateRequired(entry: NotificationEntry) {
+                val iconView = iconView(entry.key) ?: return
+                iconManager.updateSbIcon(entry, iconView)
+            }
+        }
+
+    private val notifCollectionListener =
+        object : NotifCollectionListener {
+            override fun onEntryRemoved(entry: NotificationEntry, reason: Int) {
+                cachedIcons.remove(entry.key)
+            }
+        }
+
+    override fun iconView(key: String): StatusBarIconView? {
+        val entry = notifCollection.getEntry(key) ?: return null
+        return cachedIcons.computeIfAbsent(key) {
+            val context = displayWindowPropertiesInteractor.getForStatusBar(displayId).context
+            iconManager.createSbIconView(context, entry)
+        }
+    }
+
+    override suspend fun activate() = coroutineScope {
+        start()
+        try {
+            awaitCancellation()
+        } finally {
+            stop()
+        }
+    }
+
+    private fun start() {
+        notifPipeline.addCollectionListener(notifCollectionListener)
+        iconManager.addIconsUpdateListener(iconUpdateRequiredListener)
+    }
+
+    private fun stop() {
+        notifPipeline.removeCollectionListener(notifCollectionListener)
+        iconManager.removeIconsUpdateListener(iconUpdateRequiredListener)
+    }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(displayId: Int): ConnectedDisplaysStatusBarNotificationIconViewStore
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewbinder/NotificationIconContainerStatusBarViewBinder.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewbinder/NotificationIconContainerStatusBarViewBinder.kt
index a21dabb821d4..aa81ebf22ac6 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewbinder/NotificationIconContainerStatusBarViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewbinder/NotificationIconContainerStatusBarViewBinder.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.statusbar.notification.icon.ui.viewbinder
 
+import android.view.Display
 import androidx.lifecycle.lifecycleScope
 import com.android.app.tracing.coroutines.launchTraced as launch
 import com.android.app.tracing.traceSection
@@ -29,6 +30,7 @@ import com.android.systemui.statusbar.phone.NotificationIconContainer
 import com.android.systemui.statusbar.ui.SystemBarUtilsState
 import javax.inject.Inject
 import kotlinx.coroutines.DisposableHandle
+import kotlinx.coroutines.launch
 
 /** Binds a [NotificationIconContainer] to a [NotificationIconContainerStatusBarViewModel]. */
 class NotificationIconContainerStatusBarViewBinder
@@ -38,11 +40,22 @@ constructor(
     @ShadeDisplayAware private val configuration: ConfigurationState,
     private val systemBarUtilsState: SystemBarUtilsState,
     private val failureTracker: StatusBarIconViewBindingFailureTracker,
-    private val viewStore: StatusBarNotificationIconViewStore,
+    private val defaultDisplayViewStore: StatusBarNotificationIconViewStore,
+    private val connectedDisplaysViewStoreFactory:
+        ConnectedDisplaysStatusBarNotificationIconViewStore.Factory,
 ) {
+
     fun bindWhileAttached(view: NotificationIconContainer, displayId: Int): DisposableHandle {
         return traceSection("NICStatusBar#bindWhileAttached") {
             view.repeatWhenAttached {
+                val viewStore =
+                    if (displayId == Display.DEFAULT_DISPLAY) {
+                        defaultDisplayViewStore
+                    } else {
+                        connectedDisplaysViewStoreFactory.create(displayId = displayId).also {
+                            lifecycleScope.launch { it.activate() }
+                        }
+                    }
                 lifecycleScope.launch {
                     NotificationIconContainerViewBinder.bind(
                         displayId = displayId,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/fragment/CollapsedStatusBarFragment.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/fragment/CollapsedStatusBarFragment.java
index 23b4b65bb2ac..6de4928cd0c1 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/fragment/CollapsedStatusBarFragment.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/fragment/CollapsedStatusBarFragment.java
@@ -504,12 +504,7 @@ public class CollapsedStatusBarFragment extends Fragment implements CommandQueue
                 notificationIconArea.requireViewById(R.id.notificationIcons);
         mNotificationIconAreaInner = notificationIcons;
         int displayId = mHomeStatusBarComponent.getDisplayId();
-        if (displayId == Display.DEFAULT_DISPLAY) {
-            //TODO(b/369337701): implement notification icons for all displays.
-            // Currently if we try to bind for all displays, there is a crash, because the same
-            // notification icon view can't have multiple parents.
-            mNicBindingDisposable = mNicViewBinder.bindWhileAttached(notificationIcons, displayId);
-        }
+        mNicBindingDisposable = mNicViewBinder.bindWhileAttached(notificationIcons, displayId);
 
         if (!StatusBarRootModernization.isEnabled()) {
             updateNotificationIconAreaAndOngoingActivityChip(/* animate= */ false);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/composable/StatusBarRoot.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/composable/StatusBarRoot.kt
index 1faa9f32af1f..1e8b0166409c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/composable/StatusBarRoot.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/composable/StatusBarRoot.kt
@@ -16,7 +16,6 @@
 
 package com.android.systemui.statusbar.pipeline.shared.ui.composable
 
-import android.view.Display
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
@@ -167,17 +166,11 @@ fun StatusBarRoot(
                             R.id.notificationIcons
                         )
 
-                    // TODO(b/369337701): implement notification icons for all displays.
-                    //  Currently if we try to bind for all displays, there is a crash, because the
-                    //  same notification icon view can't have multiple parents.
-                    val displayId = context.displayId
-                    if (displayId == Display.DEFAULT_DISPLAY) {
-                        scope.launch {
-                            notificationIconsBinder.bindWhileAttached(
-                                notificationIconContainer,
-                                displayId,
-                            )
-                        }
+                    scope.launch {
+                        notificationIconsBinder.bindWhileAttached(
+                            notificationIconContainer,
+                            context.displayId,
+                        )
                     }
 
                     // This binder handles everything else
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/DisplayWindowPropertiesRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/DisplayWindowPropertiesRepositoryKosmos.kt
index 65b18c102a16..5b940f93c4ee 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/DisplayWindowPropertiesRepositoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/DisplayWindowPropertiesRepositoryKosmos.kt
@@ -16,10 +16,11 @@
 
 package com.android.systemui.display.data.repository
 
+import android.content.testableContext
 import com.android.systemui.kosmos.Kosmos
 
 val Kosmos.fakeDisplayWindowPropertiesRepository by
-    Kosmos.Fixture { FakeDisplayWindowPropertiesRepository() }
+    Kosmos.Fixture { FakeDisplayWindowPropertiesRepository(testableContext) }
 
 var Kosmos.displayWindowPropertiesRepository: DisplayWindowPropertiesRepository by
     Kosmos.Fixture { fakeDisplayWindowPropertiesRepository }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayWindowPropertiesRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayWindowPropertiesRepository.kt
index 7fd927654ca6..534ded57eb85 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayWindowPropertiesRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayWindowPropertiesRepository.kt
@@ -16,11 +16,16 @@
 
 package com.android.systemui.display.data.repository
 
+import android.content.Context
+import android.view.Display
 import com.android.systemui.display.shared.model.DisplayWindowProperties
 import com.google.common.collect.HashBasedTable
+import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.mock
+import org.mockito.kotlin.spy
 
-class FakeDisplayWindowPropertiesRepository : DisplayWindowPropertiesRepository {
+class FakeDisplayWindowPropertiesRepository(private val context: Context) :
+    DisplayWindowPropertiesRepository {
 
     private val properties = HashBasedTable.create<Int, Int, DisplayWindowProperties>()
 
@@ -29,13 +34,26 @@ class FakeDisplayWindowPropertiesRepository : DisplayWindowPropertiesRepository
             ?: DisplayWindowProperties(
                     displayId = displayId,
                     windowType = windowType,
-                    context = mock(),
+                    context = contextWithDisplayId(context, displayId),
                     windowManager = mock(),
                     layoutInflater = mock(),
                 )
                 .also { properties.put(displayId, windowType, it) }
     }
 
+    private fun contextWithDisplayId(context: Context, displayId: Int): Context {
+        val newDisplay = displayWithId(context.display, displayId)
+        return spy(context) {
+            on { getDisplayId() } doReturn displayId
+            on { display } doReturn newDisplay
+            on { displayNoVerify } doReturn newDisplay
+        }
+    }
+
+    private fun displayWithId(display: Display, displayId: Int): Display {
+        return spy(display) { on { getDisplayId() } doReturn displayId }
+    }
+
     /** Sets an instance, just for testing purposes. */
     fun insert(instance: DisplayWindowProperties) {
         properties.put(instance.displayId, instance.windowType, instance)