Merge "Preserve notif group stucture in repository" into main
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/data/repository/ActiveNotificationListRepository.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/data/repository/ActiveNotificationListRepository.kt
index 8064f04..12ee54d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/data/repository/ActiveNotificationListRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/data/repository/ActiveNotificationListRepository.kt
@@ -16,6 +16,9 @@
 package com.android.systemui.statusbar.notification.data.repository
 
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.statusbar.notification.data.repository.ActiveNotificationsStore.Key
+import com.android.systemui.statusbar.notification.shared.ActiveNotificationEntryModel
+import com.android.systemui.statusbar.notification.shared.ActiveNotificationGroupModel
 import com.android.systemui.statusbar.notification.shared.ActiveNotificationModel
 import javax.inject.Inject
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -29,13 +32,64 @@
  */
 @SysUISingleton
 class ActiveNotificationListRepository @Inject constructor() {
-    /**
-     * Notifications actively presented to the user in the notification stack.
-     *
-     * @see com.android.systemui.statusbar.notification.collection.listbuilder.OnAfterRenderListListener
-     */
-    val activeNotifications = MutableStateFlow(emptyMap<String, ActiveNotificationModel>())
+    /** Notifications actively presented to the user in the notification list. */
+    val activeNotifications = MutableStateFlow(ActiveNotificationsStore())
 
     /** Are any already-seen notifications currently filtered out of the active list? */
     val hasFilteredOutSeenNotifications = MutableStateFlow(false)
 }
+
+/** Represents the notification list, comprised of groups and individual notifications. */
+data class ActiveNotificationsStore(
+    /** Notification groups, stored by key. */
+    val groups: Map<String, ActiveNotificationGroupModel> = emptyMap(),
+    /** All individual notifications, including top-level and group children, stored by key. */
+    val individuals: Map<String, ActiveNotificationModel> = emptyMap(),
+    /**
+     * Ordered top-level list of entries in the notification list (either groups or individual),
+     * represented as [Key]s. The associated [ActiveNotificationEntryModel] can be retrieved by
+     * invoking [get].
+     */
+    val renderList: List<Key> = emptyList(),
+) {
+    operator fun get(key: Key): ActiveNotificationEntryModel? {
+        return when (key) {
+            is Key.Group -> groups[key.key]
+            is Key.Individual -> individuals[key.key]
+        }
+    }
+
+    /** Unique key identifying an [ActiveNotificationEntryModel] in the store. */
+    sealed class Key {
+        data class Individual(val key: String) : Key()
+        data class Group(val key: String) : Key()
+    }
+
+    /** Mutable builder for an [ActiveNotificationsStore]. */
+    class Builder {
+        private val groups = mutableMapOf<String, ActiveNotificationGroupModel>()
+        private val individuals = mutableMapOf<String, ActiveNotificationModel>()
+        private val renderList = mutableListOf<Key>()
+
+        fun build() = ActiveNotificationsStore(groups, individuals, renderList)
+
+        fun addEntry(entry: ActiveNotificationEntryModel) {
+            when (entry) {
+                is ActiveNotificationModel -> addIndividualNotif(entry)
+                is ActiveNotificationGroupModel -> addNotifGroup(entry)
+            }
+        }
+
+        fun addIndividualNotif(notif: ActiveNotificationModel) {
+            renderList.add(Key.Individual(notif.key))
+            individuals[notif.key] = notif
+        }
+
+        fun addNotifGroup(group: ActiveNotificationGroupModel) {
+            renderList.add(Key.Group(group.key))
+            groups[group.key] = group
+            individuals[group.summary.key] = group.summary
+            group.children.forEach { individuals[it.key] = it }
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/ActiveNotificationsInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/ActiveNotificationsInteractor.kt
index bfec60bc..85ba205 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/ActiveNotificationsInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/ActiveNotificationsInteractor.kt
@@ -16,6 +16,7 @@
 package com.android.systemui.statusbar.notification.domain.interactor
 
 import com.android.systemui.statusbar.notification.data.repository.ActiveNotificationListRepository
+import com.android.systemui.statusbar.notification.shared.ActiveNotificationGroupModel
 import com.android.systemui.statusbar.notification.shared.ActiveNotificationModel
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
@@ -27,6 +28,16 @@
     repository: ActiveNotificationListRepository,
 ) {
     /** Notifications actively presented to the user in the notification stack, in order. */
-    val notifications: Flow<Collection<ActiveNotificationModel>> =
-        repository.activeNotifications.map { it.values }
+    val topLevelRepresentativeNotifications: Flow<List<ActiveNotificationModel>> =
+        repository.activeNotifications.map { store ->
+            store.renderList.map { key ->
+                val entry =
+                    store[key]
+                        ?: error("Could not find notification with key $key in active notif store.")
+                when (entry) {
+                    is ActiveNotificationGroupModel -> entry.summary
+                    is ActiveNotificationModel -> entry
+                }
+            }
+        }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationListInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationListInteractor.kt
index 6a0abc9..6f4ed9d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationListInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationListInteractor.kt
@@ -16,16 +16,18 @@
 package com.android.systemui.statusbar.notification.domain.interactor
 
 import android.graphics.drawable.Icon
+import com.android.systemui.statusbar.notification.collection.GroupEntry
 import com.android.systemui.statusbar.notification.collection.ListEntry
 import com.android.systemui.statusbar.notification.collection.NotificationEntry
 import com.android.systemui.statusbar.notification.collection.provider.SectionStyleProvider
 import com.android.systemui.statusbar.notification.data.repository.ActiveNotificationListRepository
+import com.android.systemui.statusbar.notification.data.repository.ActiveNotificationsStore
+import com.android.systemui.statusbar.notification.shared.ActiveNotificationEntryModel
+import com.android.systemui.statusbar.notification.shared.ActiveNotificationGroupModel
 import com.android.systemui.statusbar.notification.shared.ActiveNotificationModel
 import javax.inject.Inject
 import kotlinx.coroutines.flow.update
 
-private typealias ModelStore = Map<String, ActiveNotificationModel>
-
 /**
  * Logic for passing information from the
  * [com.android.systemui.statusbar.notification.collection.NotifPipeline] to the presentation
@@ -38,24 +40,61 @@
     private val sectionStyleProvider: SectionStyleProvider,
 ) {
     /**
-     * Sets the current list of rendered notification entries as displayed in the notification
-     * stack.
-     *
-     * @see com.android.systemui.statusbar.notification.data.repository.ActiveNotificationListRepository.activeNotifications
+     * Sets the current list of rendered notification entries as displayed in the notification list.
      */
     fun setRenderedList(entries: List<ListEntry>) {
         repository.activeNotifications.update { existingModels ->
-            entries
-                .asSequence()
-                .mapNotNull { it.representativeEntry }
-                .associateBy(
-                    keySelector = { it.key },
-                    valueTransform = { it.toModel(existingModels) },
-                )
+            buildActiveNotificationsStore(existingModels, sectionStyleProvider) {
+                entries.forEach(::addListEntry)
+            }
+        }
+    }
+}
+
+private fun buildActiveNotificationsStore(
+    existingModels: ActiveNotificationsStore,
+    sectionStyleProvider: SectionStyleProvider,
+    block: ActiveNotificationsStoreBuilder.() -> Unit
+): ActiveNotificationsStore =
+    ActiveNotificationsStoreBuilder(existingModels, sectionStyleProvider).apply(block).build()
+
+private class ActiveNotificationsStoreBuilder(
+    private val existingModels: ActiveNotificationsStore,
+    private val sectionStyleProvider: SectionStyleProvider,
+) {
+    private val builder = ActiveNotificationsStore.Builder()
+
+    fun build(): ActiveNotificationsStore = builder.build()
+
+    /**
+     * Convert a [ListEntry] into [ActiveNotificationEntryModel]s, and add them to the
+     * [ActiveNotificationsStore]. Special care is taken to avoid re-allocating models if the result
+     * would be identical to an existing model (at the expense of additional computations).
+     */
+    fun addListEntry(entry: ListEntry) {
+        when (entry) {
+            is GroupEntry -> {
+                entry.summary?.let { summary ->
+                    val summaryModel = summary.toModel()
+                    val childModels = entry.children.map { it.toModel() }
+                    builder.addNotifGroup(
+                        existingModels.createOrReuse(
+                            key = entry.key,
+                            summary = summaryModel,
+                            children = childModels
+                        )
+                    )
+                }
+            }
+            else -> {
+                entry.representativeEntry?.let { notifEntry ->
+                    builder.addIndividualNotif(notifEntry.toModel())
+                }
+            }
         }
     }
 
-    private fun NotificationEntry.toModel(existingModels: ModelStore): ActiveNotificationModel =
+    private fun NotificationEntry.toModel(): ActiveNotificationModel =
         existingModels.createOrReuse(
             key = key,
             groupKey = sbn.groupKey,
@@ -69,76 +108,98 @@
             shelfIcon = icons.shelfIcon?.sourceIcon,
             statusBarIcon = icons.statusBarIcon?.sourceIcon,
         )
+}
 
-    private fun ModelStore.createOrReuse(
-        key: String,
-        groupKey: String?,
-        isAmbient: Boolean,
-        isRowDismissed: Boolean,
-        isSilent: Boolean,
-        isLastMessageFromReply: Boolean,
-        isSuppressedFromStatusBar: Boolean,
-        isPulsing: Boolean,
-        aodIcon: Icon?,
-        shelfIcon: Icon?,
-        statusBarIcon: Icon?
-    ): ActiveNotificationModel {
-        return this[key]?.takeIf {
-            it.isCurrent(
-                key = key,
-                groupKey = groupKey,
-                isAmbient = isAmbient,
-                isRowDismissed = isRowDismissed,
-                isSilent = isSilent,
-                isLastMessageFromReply = isLastMessageFromReply,
-                isSuppressedFromStatusBar = isSuppressedFromStatusBar,
-                isPulsing = isPulsing,
-                aodIcon = aodIcon,
-                shelfIcon = shelfIcon,
-                statusBarIcon = statusBarIcon
-            )
-        }
-            ?: ActiveNotificationModel(
-                key = key,
-                groupKey = groupKey,
-                isAmbient = isAmbient,
-                isRowDismissed = isRowDismissed,
-                isSilent = isSilent,
-                isLastMessageFromReply = isLastMessageFromReply,
-                isSuppressedFromStatusBar = isSuppressedFromStatusBar,
-                isPulsing = isPulsing,
-                aodIcon = aodIcon,
-                shelfIcon = shelfIcon,
-                statusBarIcon = statusBarIcon,
-            )
+private fun ActiveNotificationsStore.createOrReuse(
+    key: String,
+    groupKey: String?,
+    isAmbient: Boolean,
+    isRowDismissed: Boolean,
+    isSilent: Boolean,
+    isLastMessageFromReply: Boolean,
+    isSuppressedFromStatusBar: Boolean,
+    isPulsing: Boolean,
+    aodIcon: Icon?,
+    shelfIcon: Icon?,
+    statusBarIcon: Icon?
+): ActiveNotificationModel {
+    return individuals[key]?.takeIf {
+        it.isCurrent(
+            key = key,
+            groupKey = groupKey,
+            isAmbient = isAmbient,
+            isRowDismissed = isRowDismissed,
+            isSilent = isSilent,
+            isLastMessageFromReply = isLastMessageFromReply,
+            isSuppressedFromStatusBar = isSuppressedFromStatusBar,
+            isPulsing = isPulsing,
+            aodIcon = aodIcon,
+            shelfIcon = shelfIcon,
+            statusBarIcon = statusBarIcon
+        )
     }
+        ?: ActiveNotificationModel(
+            key = key,
+            groupKey = groupKey,
+            isAmbient = isAmbient,
+            isRowDismissed = isRowDismissed,
+            isSilent = isSilent,
+            isLastMessageFromReply = isLastMessageFromReply,
+            isSuppressedFromStatusBar = isSuppressedFromStatusBar,
+            isPulsing = isPulsing,
+            aodIcon = aodIcon,
+            shelfIcon = shelfIcon,
+            statusBarIcon = statusBarIcon,
+        )
+}
 
-    private fun ActiveNotificationModel.isCurrent(
-        key: String,
-        groupKey: String?,
-        isAmbient: Boolean,
-        isRowDismissed: Boolean,
-        isSilent: Boolean,
-        isLastMessageFromReply: Boolean,
-        isSuppressedFromStatusBar: Boolean,
-        isPulsing: Boolean,
-        aodIcon: Icon?,
-        shelfIcon: Icon?,
-        statusBarIcon: Icon?
-    ): Boolean {
-        return when {
-            key != this.key -> false
-            groupKey != this.groupKey -> false
-            isAmbient != this.isAmbient -> false
-            isRowDismissed != this.isRowDismissed -> false
-            isSilent != this.isSilent -> false
-            isLastMessageFromReply != this.isLastMessageFromReply -> false
-            isSuppressedFromStatusBar != this.isSuppressedFromStatusBar -> false
-            isPulsing != this.isPulsing -> false
-            aodIcon != this.aodIcon -> false
-            shelfIcon != this.shelfIcon -> false
-            statusBarIcon != this.statusBarIcon -> false
-            else -> true
-        }
+private fun ActiveNotificationModel.isCurrent(
+    key: String,
+    groupKey: String?,
+    isAmbient: Boolean,
+    isRowDismissed: Boolean,
+    isSilent: Boolean,
+    isLastMessageFromReply: Boolean,
+    isSuppressedFromStatusBar: Boolean,
+    isPulsing: Boolean,
+    aodIcon: Icon?,
+    shelfIcon: Icon?,
+    statusBarIcon: Icon?
+): Boolean {
+    return when {
+        key != this.key -> false
+        groupKey != this.groupKey -> false
+        isAmbient != this.isAmbient -> false
+        isRowDismissed != this.isRowDismissed -> false
+        isSilent != this.isSilent -> false
+        isLastMessageFromReply != this.isLastMessageFromReply -> false
+        isSuppressedFromStatusBar != this.isSuppressedFromStatusBar -> false
+        isPulsing != this.isPulsing -> false
+        aodIcon != this.aodIcon -> false
+        shelfIcon != this.shelfIcon -> false
+        statusBarIcon != this.statusBarIcon -> false
+        else -> true
+    }
+}
+
+private fun ActiveNotificationsStore.createOrReuse(
+    key: String,
+    summary: ActiveNotificationModel,
+    children: List<ActiveNotificationModel>,
+): ActiveNotificationGroupModel {
+    return groups[key]?.takeIf { it.isCurrent(key, summary, children) }
+        ?: ActiveNotificationGroupModel(key, summary, children)
+}
+
+private fun ActiveNotificationGroupModel.isCurrent(
+    key: String,
+    summary: ActiveNotificationModel,
+    children: List<ActiveNotificationModel>,
+): Boolean {
+    return when {
+        key != this.key -> false
+        summary != this.summary -> false
+        children != this.children -> false
+        else -> true
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractor.kt
index 00d873e..30e2f0e0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractor.kt
@@ -48,7 +48,7 @@
         showPulsing: Boolean = true,
     ): Flow<Set<ActiveNotificationModel>> {
         return combine(
-            activeNotificationsInteractor.notifications,
+            activeNotificationsInteractor.topLevelRepresentativeNotifications,
             keyguardViewStateRepository.areNotificationsFullyHidden,
         ) { notifications, notifsFullyHidden ->
             notifications
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModel.kt
index 53631e3..82626acc 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModel.kt
@@ -61,7 +61,7 @@
             darkIconInteractor.tintAreas,
             darkIconInteractor.tintColor,
             // Included so that tints are re-applied after entries are changed.
-            notificationsInteractor.notifications,
+            notificationsInteractor.topLevelRepresentativeNotifications,
         ) { areas, tint, _ ->
             NotificationIconColorLookup { viewBounds: Rect ->
                 if (DarkIconDispatcher.isInAreas(areas, viewBounds)) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/shared/ActiveNotificationModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/shared/ActiveNotificationModel.kt
index 78370ba..eb1c1ba 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/shared/ActiveNotificationModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/shared/ActiveNotificationModel.kt
@@ -17,9 +17,17 @@
 
 import android.graphics.drawable.Icon
 
-/** Model for entries in the notification stack. */
+/**
+ * Model for a top-level "entry" in the notification list, either an
+ * [individual notification][ActiveNotificationModel], or a [group][ActiveNotificationGroupModel].
+ */
+sealed class ActiveNotificationEntryModel
+
+/**
+ * Model for an individual notification in the notification list. These can appear as either an
+ * individual top-level notification, or as a child or summary of a [ActiveNotificationGroupModel].
+ */
 data class ActiveNotificationModel(
-    /** Notification key associated with this entry. */
     val key: String,
     /** Notification group key associated with this entry. */
     val groupKey: String?,
@@ -47,4 +55,11 @@
     val shelfIcon: Icon?,
     /** Icon to display in the status bar. */
     val statusBarIcon: Icon?,
-)
+) : ActiveNotificationEntryModel()
+
+/** Model for a group of notifications. */
+data class ActiveNotificationGroupModel(
+    val key: String,
+    val summary: ActiveNotificationModel,
+    val children: List<ActiveNotificationModel>,
+) : ActiveNotificationEntryModel()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationsListInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationsListInteractorTest.kt
index 8a5d93f..b86f841 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationsListInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationsListInteractorTest.kt
@@ -41,7 +41,7 @@
 
     @Test
     fun setRenderedList_preservesOrdering() = runTest {
-        val notifs by collectLastValue(notifsInteractor.notifications)
+        val notifs by collectLastValue(notifsInteractor.topLevelRepresentativeNotifications)
         val keys = (1..50).shuffled().map { "$it" }
         val entries =
             keys.map {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorTest.kt
index ec80e5f..f8252a7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorTest.kt
@@ -25,6 +25,7 @@
 import com.android.systemui.deviceentry.data.repository.FakeDeviceEntryRepository
 import com.android.systemui.statusbar.data.repository.NotificationListenerSettingsRepository
 import com.android.systemui.statusbar.notification.data.repository.ActiveNotificationListRepository
+import com.android.systemui.statusbar.notification.data.repository.ActiveNotificationsStore
 import com.android.systemui.statusbar.notification.data.repository.FakeNotificationsKeyguardViewStateRepository
 import com.android.systemui.statusbar.notification.shared.activeNotificationModel
 import com.android.systemui.statusbar.notification.shared.byIsAmbient
@@ -77,7 +78,9 @@
     fun setup() =
         with(testComponent) {
             activeNotificationListRepository.activeNotifications.value =
-                testIcons.associateBy { it.key }
+                ActiveNotificationsStore.Builder()
+                    .apply { testIcons.forEach(::addIndividualNotif) }
+                    .build()
         }
 
     @Test
@@ -196,7 +199,9 @@
     fun setup() =
         with(testComponent) {
             activeNotificationListRepository.activeNotifications.value =
-                testIcons.associateBy { it.key }
+                ActiveNotificationsStore.Builder()
+                    .apply { testIcons.forEach(::addIndividualNotif) }
+                    .build()
         }
 
     @Test
@@ -318,7 +323,9 @@
     fun setup() =
         with(testComponent) {
             activeNotificationListRepository.activeNotifications.value =
-                testIcons.associateBy { it.key }
+                ActiveNotificationsStore.Builder()
+                    .apply { testIcons.forEach(::addIndividualNotif) }
+                    .build()
         }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModelTest.kt
index ba68fbb..44acac8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModelTest.kt
@@ -42,6 +42,7 @@
 import com.android.systemui.power.shared.model.WakefulnessState
 import com.android.systemui.shade.data.repository.FakeShadeRepository
 import com.android.systemui.statusbar.notification.data.repository.ActiveNotificationListRepository
+import com.android.systemui.statusbar.notification.data.repository.ActiveNotificationsStore
 import com.android.systemui.statusbar.notification.data.repository.HeadsUpNotificationIconViewStateRepository
 import com.android.systemui.statusbar.notification.shared.activeNotificationModel
 import com.android.systemui.statusbar.phone.DozeParameters
@@ -342,14 +343,17 @@
                 val icon: Icon = mock()
                 shadeRepository.setLegacyShadeExpansion(0f)
                 activeNotificationsRepository.activeNotifications.value =
-                    listOf(
-                            activeNotificationModel(
-                                key = "notif1",
-                                groupKey = "group",
-                                statusBarIcon = icon
+                    ActiveNotificationsStore.Builder()
+                        .apply {
+                            addIndividualNotif(
+                                activeNotificationModel(
+                                    key = "notif1",
+                                    groupKey = "group",
+                                    statusBarIcon = icon
+                                )
                             )
-                        )
-                        .associateBy { it.key }
+                        }
+                        .build()
                 val isolatedIcon by collectLastValue(underTest.isolatedIcon)
                 runCurrent()
 
@@ -368,14 +372,17 @@
                 val icon: Icon = mock()
                 shadeRepository.setLegacyShadeExpansion(.5f)
                 activeNotificationsRepository.activeNotifications.value =
-                    listOf(
-                            activeNotificationModel(
-                                key = "notif1",
-                                groupKey = "group",
-                                statusBarIcon = icon
+                    ActiveNotificationsStore.Builder()
+                        .apply {
+                            addIndividualNotif(
+                                activeNotificationModel(
+                                    key = "notif1",
+                                    groupKey = "group",
+                                    statusBarIcon = icon
+                                )
                             )
-                        )
-                        .associateBy { it.key }
+                        }
+                        .build()
                 val isolatedIcon by collectLastValue(underTest.isolatedIcon)
                 runCurrent()