Notify listeners when syncing group expansion mgr with the pipeline.

Fix: 293887828
Bug: 282865576
Bug: 293434635
Test: atest GroupExpansionManagerTest & manually verified with the flag
on and off
Flag: NOTIFICATION_GROUP_EXPANSION_CHANGE

Change-Id: I105d710d745ba6e998cdf63ec7b14156f525c2a1

3 files changed, 108 insertions, 23 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/StackCoordinator.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/StackCoordinator.kt
index 5c2f9a8d28ec..62a0d138fd05 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/StackCoordinator.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/StackCoordinator.kt
@@ -39,10 +39,7 @@ class StackCoordinator @Inject internal constructor(
 
     override fun attach(pipeline: NotifPipeline) {
         pipeline.addOnAfterRenderListListener(::onAfterRenderList)
-        // TODO(b/282865576): This has an issue where it makes changes to some groups without
-        // notifying listeners. To be fixed in QPR, but for now let's comment it out to avoid the
-        // group expansion bug.
-        // groupExpansionManagerImpl.attach(pipeline)
+        groupExpansionManagerImpl.attach(pipeline)
     }
 
     fun onAfterRenderList(entries: List<ListEntry>, controller: NotifStackController) =
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerImpl.java
index 46af03a438f5..5d33804ab6a7 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerImpl.java
@@ -67,18 +67,29 @@ public class GroupExpansionManagerImpl implements GroupExpansionManager, Dumpabl
      * Cleanup entries from mExpandedGroups that no longer exist in the pipeline.
      */
     private final OnBeforeRenderListListener mNotifTracker = (entries) -> {
+        if (mExpandedGroups.isEmpty()) {
+            return; // nothing to do
+        }
+
         final Set<NotificationEntry> renderingSummaries = new HashSet<>();
         for (ListEntry entry : entries) {
             if (entry instanceof GroupEntry) {
                 renderingSummaries.add(entry.getRepresentativeEntry());
             }
         }
-        mExpandedGroups.removeIf(expandedGroup -> !renderingSummaries.contains(expandedGroup));
+
+        // Create a copy of mExpandedGroups so we can modify it in a thread-safe way.
+        final var currentExpandedGroups = new HashSet<>(mExpandedGroups);
+        for (NotificationEntry entry : currentExpandedGroups) {
+            setExpanded(entry, renderingSummaries.contains(entry));
+        }
     };
 
     public void attach(NotifPipeline pipeline) {
-        mDumpManager.registerDumpable(this);
-        pipeline.addOnBeforeRenderListListener(mNotifTracker);
+        if (mFeatureFlags.isEnabled(Flags.NOTIFICATION_GROUP_EXPANSION_CHANGE)) {
+            mDumpManager.registerDumpable(this);
+            pipeline.addOnBeforeRenderListListener(mNotifTracker);
+        }
     }
 
     @Override
@@ -94,11 +105,24 @@ public class GroupExpansionManagerImpl implements GroupExpansionManager, Dumpabl
     @Override
     public void setGroupExpanded(NotificationEntry entry, boolean expanded) {
         final NotificationEntry groupSummary = mGroupMembershipManager.getGroupSummary(entry);
+        setExpanded(groupSummary, expanded);
+    }
+
+    /**
+     * Add or remove {@code entry} to/from {@code mExpandedGroups} and notify listeners if
+     * something changed. This assumes that {@code entry} is a group summary.
+     * <p>
+     * TODO(b/293434635): Currently, in spite of its docs,
+     * {@code mGroupMembershipManager.getGroupSummary(entry)} returns null if {@code entry} is
+     * already a summary. Instead of needing this helper method to bypass that, we probably want to
+     * move this code back to {@code setGroupExpanded} and use that everywhere.
+     */
+    private void setExpanded(NotificationEntry entry, boolean expanded) {
         boolean changed;
         if (expanded) {
-            changed = mExpandedGroups.add(groupSummary);
+            changed = mExpandedGroups.add(entry);
         } else {
-            changed = mExpandedGroups.remove(groupSummary);
+            changed = mExpandedGroups.remove(entry);
         }
 
         // Only notify listeners if something changed.
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerTest.kt
index 4a94dc819a9e..38a8f414b0fb 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerTest.kt
@@ -21,11 +21,21 @@ import com.android.systemui.SysuiTestCase
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.flags.FakeFeatureFlags
 import com.android.systemui.flags.Flags
+import com.android.systemui.statusbar.notification.collection.GroupEntryBuilder
+import com.android.systemui.statusbar.notification.collection.ListEntry
+import com.android.systemui.statusbar.notification.collection.NotifPipeline
 import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder
+import com.android.systemui.statusbar.notification.collection.listbuilder.OnBeforeRenderListListener
+import com.android.systemui.statusbar.notification.collection.render.GroupExpansionManager.OnGroupExpansionChangeListener
+import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.withArgCaptor
 import org.junit.Assert
 import org.junit.Before
 import org.junit.Test
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
+import org.mockito.Mockito.verifyNoMoreInteractions
 import org.mockito.Mockito.`when` as whenever
 
 @SmallTest
@@ -36,13 +46,43 @@ class GroupExpansionManagerTest : SysuiTestCase() {
     private val groupMembershipManager: GroupMembershipManager = mock()
     private val featureFlags = FakeFeatureFlags()
 
-    private val entry1 = NotificationEntryBuilder().build()
-    private val entry2 = NotificationEntryBuilder().build()
+    private val pipeline: NotifPipeline = mock()
+    private lateinit var beforeRenderListListener: OnBeforeRenderListListener
+
+    private val summary1 = notificationEntry("foo", 1)
+    private val summary2 = notificationEntry("bar", 1)
+    private val entries =
+        listOf<ListEntry>(
+            GroupEntryBuilder()
+                .setSummary(summary1)
+                .setChildren(
+                    listOf(
+                        notificationEntry("foo", 2),
+                        notificationEntry("foo", 3),
+                        notificationEntry("foo", 4)
+                    )
+                )
+                .build(),
+            GroupEntryBuilder()
+                .setSummary(summary2)
+                .setChildren(
+                    listOf(
+                        notificationEntry("bar", 2),
+                        notificationEntry("bar", 3),
+                        notificationEntry("bar", 4)
+                    )
+                )
+                .build(),
+            notificationEntry("baz", 1)
+        )
+
+    private fun notificationEntry(pkg: String, id: Int) =
+        NotificationEntryBuilder().setPkg(pkg).setId(id).build().apply { row = mock() }
 
     @Before
     fun setUp() {
-        whenever(groupMembershipManager.getGroupSummary(entry1)).thenReturn(entry1)
-        whenever(groupMembershipManager.getGroupSummary(entry2)).thenReturn(entry2)
+        whenever(groupMembershipManager.getGroupSummary(summary1)).thenReturn(summary1)
+        whenever(groupMembershipManager.getGroupSummary(summary2)).thenReturn(summary2)
 
         gem = GroupExpansionManagerImpl(dumpManager, groupMembershipManager, featureFlags)
     }
@@ -54,15 +94,15 @@ class GroupExpansionManagerTest : SysuiTestCase() {
         var listenerCalledCount = 0
         gem.registerGroupExpansionChangeListener { _, _ -> listenerCalledCount++ }
 
-        gem.setGroupExpanded(entry1, false)
+        gem.setGroupExpanded(summary1, false)
         Assert.assertEquals(0, listenerCalledCount)
-        gem.setGroupExpanded(entry1, true)
+        gem.setGroupExpanded(summary1, true)
         Assert.assertEquals(1, listenerCalledCount)
-        gem.setGroupExpanded(entry2, true)
+        gem.setGroupExpanded(summary2, true)
         Assert.assertEquals(2, listenerCalledCount)
-        gem.setGroupExpanded(entry1, true)
+        gem.setGroupExpanded(summary1, true)
         Assert.assertEquals(2, listenerCalledCount)
-        gem.setGroupExpanded(entry2, false)
+        gem.setGroupExpanded(summary2, false)
         Assert.assertEquals(3, listenerCalledCount)
     }
 
@@ -73,15 +113,39 @@ class GroupExpansionManagerTest : SysuiTestCase() {
         var listenerCalledCount = 0
         gem.registerGroupExpansionChangeListener { _, _ -> listenerCalledCount++ }
 
-        gem.setGroupExpanded(entry1, false)
+        gem.setGroupExpanded(summary1, false)
         Assert.assertEquals(1, listenerCalledCount)
-        gem.setGroupExpanded(entry1, true)
+        gem.setGroupExpanded(summary1, true)
         Assert.assertEquals(2, listenerCalledCount)
-        gem.setGroupExpanded(entry2, true)
+        gem.setGroupExpanded(summary2, true)
         Assert.assertEquals(3, listenerCalledCount)
-        gem.setGroupExpanded(entry1, true)
+        gem.setGroupExpanded(summary1, true)
         Assert.assertEquals(4, listenerCalledCount)
-        gem.setGroupExpanded(entry2, false)
+        gem.setGroupExpanded(summary2, false)
         Assert.assertEquals(5, listenerCalledCount)
     }
+
+    @Test
+    fun testSyncWithPipeline() {
+        featureFlags.set(Flags.NOTIFICATION_GROUP_EXPANSION_CHANGE, true)
+        gem.attach(pipeline)
+        beforeRenderListListener = withArgCaptor {
+            verify(pipeline).addOnBeforeRenderListListener(capture())
+        }
+
+        val listener: OnGroupExpansionChangeListener = mock()
+        gem.registerGroupExpansionChangeListener(listener)
+
+        beforeRenderListListener.onBeforeRenderList(entries)
+        verify(listener, never()).onGroupExpansionChange(any(), any())
+
+        // Expand one of the groups.
+        gem.setGroupExpanded(summary1, true)
+        verify(listener).onGroupExpansionChange(summary1.row, true)
+
+        // Empty the pipeline list and verify that the group is no longer expanded.
+        beforeRenderListListener.onBeforeRenderList(emptyList())
+        verify(listener).onGroupExpansionChange(summary1.row, false)
+        verifyNoMoreInteractions(listener)
+    }
 }