Fix forced-grouping on notification channel updates

 Fixes some edge-cases:
    - updating a channel moves a notification to a new, empty section => no autogrouping
    - updating a channel moves a notification to a section with ungrouped notifications,
        and auto-grouping minimum count is reached => trigger autogrouping

Flag: android.service.notification.notification_force_grouping

Test: atest GroupHelperTest
Bug: 336488844
Change-Id: Ic9a07fc4dea36751cead076d8928fea55958adee
diff --git a/services/core/java/com/android/server/notification/GroupHelper.java b/services/core/java/com/android/server/notification/GroupHelper.java
index 008746c..e5abb44 100644
--- a/services/core/java/com/android/server/notification/GroupHelper.java
+++ b/services/core/java/com/android/server/notification/GroupHelper.java
@@ -830,61 +830,19 @@
                 }
             }
 
+            // The list of notification operations required after the channel update
             final ArrayList<NotificationMoveOp> notificationsToMove = new ArrayList<>();
 
-            final Set<FullyQualifiedGroupKey> oldGroups =
-                    new HashSet<>(mAggregatedNotifications.keySet());
-            for (FullyQualifiedGroupKey oldFullAggKey : oldGroups) {
-                // Only check aggregate groups that match the same userId & packageName
-                if (pkgName.equals(oldFullAggKey.pkg) && userId == oldFullAggKey.userId) {
-                    final ArrayMap<String, NotificationAttributes> notificationsInAggGroup =
-                            mAggregatedNotifications.get(oldFullAggKey);
-                    if (notificationsInAggGroup == null) {
-                        continue;
-                    }
+            // Check any already auto-grouped notifications that may need to be re-grouped
+            // after the channel update
+            notificationsToMove.addAll(
+                    getAutogroupedNotificationsMoveOps(userId, pkgName,
+                        notificationsToCheck));
 
-                    FullyQualifiedGroupKey newFullAggregateGroupKey = null;
-                    for (String key : notificationsInAggGroup.keySet()) {
-                        if (notificationsToCheck.get(key) != null) {
-                            // check if section changes
-                            NotificationSectioner sectioner = getSection(
-                                    notificationsToCheck.get(key));
-                            if (sectioner == null) {
-                                continue;
-                            }
-                            newFullAggregateGroupKey = new FullyQualifiedGroupKey(userId, pkgName,
-                                    sectioner);
-                            if (!oldFullAggKey.equals(newFullAggregateGroupKey)) {
-                                if (DEBUG) {
-                                    Log.i(TAG, "Change section on channel update: " + key);
-                                }
-                                notificationsToMove.add(
-                                        new NotificationMoveOp(notificationsToCheck.get(key),
-                                            oldFullAggKey, newFullAggregateGroupKey));
-                            }
-                        }
-                    }
-
-                    if (newFullAggregateGroupKey != null) {
-                        // Add any notifications left ungrouped to the new section
-                        ArrayMap<String, NotificationAttributes> ungrouped =
-                            mUngroupedAbuseNotifications.get(newFullAggregateGroupKey);
-                        if (ungrouped != null) {
-                            for (NotificationRecord r : notificationList) {
-                                if (ungrouped.containsKey(r.getKey())) {
-                                    if (DEBUG) {
-                                        Log.i(TAG, "Add previously ungrouped: " + r);
-                                    }
-                                    notificationsToMove.add(
-                                        new NotificationMoveOp(r, null, newFullAggregateGroupKey));
-                                }
-                            }
-                            //Cleanup mUngroupedAbuseNotifications
-                            mUngroupedAbuseNotifications.remove(newFullAggregateGroupKey);
-                        }
-                    }
-                }
-            }
+            // Check any ungrouped notifications that may need to be auto-grouped
+            // after the channel update
+            notificationsToMove.addAll(
+                    getUngroupedNotificationsMoveOps(userId, pkgName, notificationsToCheck));
 
             // Batch move to new section
             if (!notificationsToMove.isEmpty()) {
@@ -894,10 +852,103 @@
     }
 
     @GuardedBy("mAggregatedNotifications")
+    private List<NotificationMoveOp> getAutogroupedNotificationsMoveOps(int userId, String pkgName,
+            ArrayMap<String, NotificationRecord> notificationsToCheck) {
+        final ArrayList<NotificationMoveOp> notificationsToMove = new ArrayList<>();
+        final Set<FullyQualifiedGroupKey> oldGroups =
+                new HashSet<>(mAggregatedNotifications.keySet());
+        // Move auto-grouped updated notifications from the old groups to the new groups (section)
+        for (FullyQualifiedGroupKey oldFullAggKey : oldGroups) {
+            // Only check aggregate groups that match the same userId & packageName
+            if (pkgName.equals(oldFullAggKey.pkg) && userId == oldFullAggKey.userId) {
+                final ArrayMap<String, NotificationAttributes> notificationsInAggGroup =
+                        mAggregatedNotifications.get(oldFullAggKey);
+                if (notificationsInAggGroup == null) {
+                    continue;
+                }
+
+                FullyQualifiedGroupKey newFullAggregateGroupKey = null;
+                for (String key : notificationsInAggGroup.keySet()) {
+                    if (notificationsToCheck.get(key) != null) {
+                        // check if section changes
+                        NotificationSectioner sectioner = getSection(notificationsToCheck.get(key));
+                        if (sectioner == null) {
+                            continue;
+                        }
+                        newFullAggregateGroupKey = new FullyQualifiedGroupKey(userId, pkgName,
+                                sectioner);
+                        if (!oldFullAggKey.equals(newFullAggregateGroupKey)) {
+                            if (DEBUG) {
+                                Log.i(TAG, "Change section on channel update: " + key);
+                            }
+                            notificationsToMove.add(
+                                    new NotificationMoveOp(notificationsToCheck.get(key),
+                                        oldFullAggKey, newFullAggregateGroupKey));
+                            notificationsToCheck.remove(key);
+                        }
+                    }
+                }
+            }
+        }
+        return notificationsToMove;
+    }
+
+    @GuardedBy("mAggregatedNotifications")
+    private List<NotificationMoveOp> getUngroupedNotificationsMoveOps(int userId, String pkgName,
+            final ArrayMap<String, NotificationRecord> notificationsToCheck) {
+        final ArrayList<NotificationMoveOp> notificationsToMove = new ArrayList<>();
+        // Move any remaining ungrouped updated notifications from the old ungrouped list
+        // to the new ungrouped section list, if necessary
+        if (!notificationsToCheck.isEmpty()) {
+            final Set<FullyQualifiedGroupKey> oldUngroupedSectionKeys =
+                    new HashSet<>(mUngroupedAbuseNotifications.keySet());
+            for (FullyQualifiedGroupKey oldFullAggKey : oldUngroupedSectionKeys) {
+                // Only check aggregate groups that match the same userId & packageName
+                if (pkgName.equals(oldFullAggKey.pkg) && userId == oldFullAggKey.userId) {
+                    final ArrayMap<String, NotificationAttributes> ungroupedOld =
+                            mUngroupedAbuseNotifications.get(oldFullAggKey);
+                    if (ungroupedOld == null) {
+                        continue;
+                    }
+
+                    FullyQualifiedGroupKey newFullAggregateGroupKey = null;
+                    final Set<String> ungroupedKeys = new HashSet<>(ungroupedOld.keySet());
+                    for (String key : ungroupedKeys) {
+                        NotificationRecord record = notificationsToCheck.get(key);
+                        if (record != null) {
+                            // check if section changes
+                            NotificationSectioner sectioner = getSection(record);
+                            if (sectioner == null) {
+                                continue;
+                            }
+                            newFullAggregateGroupKey = new FullyQualifiedGroupKey(userId, pkgName,
+                                    sectioner);
+                            if (!oldFullAggKey.equals(newFullAggregateGroupKey)) {
+                                if (DEBUG) {
+                                    Log.i(TAG, "Change ungrouped section: " + key);
+                                }
+                                notificationsToMove.add(
+                                        new NotificationMoveOp(record, oldFullAggKey,
+                                            newFullAggregateGroupKey));
+                                notificationsToCheck.remove(key);
+                                //Remove from previous ungrouped list
+                                ungroupedOld.remove(key);
+                            }
+                        }
+                    }
+                    mUngroupedAbuseNotifications.put(oldFullAggKey, ungroupedOld);
+                }
+            }
+        }
+        return notificationsToMove;
+    }
+
+    @GuardedBy("mAggregatedNotifications")
     private void moveNotificationsToNewSection(final int userId, final String pkgName,
             final List<NotificationMoveOp> notificationsToMove) {
         record GroupUpdateOp(FullyQualifiedGroupKey groupKey, NotificationRecord record,
                              boolean hasSummary) { }
+        // Bundled operations to apply to groups affected by the channel update
         ArrayMap<FullyQualifiedGroupKey, GroupUpdateOp> groupsToUpdate = new ArrayMap<>();
 
         for (NotificationMoveOp moveOp: notificationsToMove) {
@@ -923,35 +974,36 @@
                 // Only add once, for triggering notification
                 if (!groupsToUpdate.containsKey(oldFullAggregateGroupKey)) {
                     groupsToUpdate.put(oldFullAggregateGroupKey,
-                            new GroupUpdateOp(oldFullAggregateGroupKey, record, true));
+                        new GroupUpdateOp(oldFullAggregateGroupKey, record, true));
                 }
             }
 
-            // Add/update aggregate summary for new group
+            // Add moved notifications to the ungrouped list for new group and do grouping
+            // after all notifications have been handled
             if (newFullAggregateGroupKey != null) {
                 final ArrayMap<String, NotificationAttributes> newAggregatedNotificationsAttrs =
                         mAggregatedNotifications.getOrDefault(newFullAggregateGroupKey,
                             new ArrayMap<>());
-                boolean newGroupExists = !newAggregatedNotificationsAttrs.isEmpty();
-                newAggregatedNotificationsAttrs.put(record.getKey(),
-                        new NotificationAttributes(record.getFlags(),
-                            record.getNotification().getSmallIcon(),
-                            record.getNotification().color,
-                            record.getNotification().visibility,
-                            record.getNotification().getGroupAlertBehavior(),
-                            record.getChannel().getId()));
-                mAggregatedNotifications.put(newFullAggregateGroupKey,
-                        newAggregatedNotificationsAttrs);
+                boolean hasSummary = !newAggregatedNotificationsAttrs.isEmpty();
+                ArrayMap<String, NotificationAttributes> ungrouped =
+                        mUngroupedAbuseNotifications.getOrDefault(newFullAggregateGroupKey,
+                            new ArrayMap<>());
+                ungrouped.put(record.getKey(), new NotificationAttributes(
+                        record.getFlags(),
+                        record.getNotification().getSmallIcon(),
+                        record.getNotification().color,
+                        record.getNotification().visibility,
+                        record.getNotification().getGroupAlertBehavior(),
+                        record.getChannel().getId()));
+                mUngroupedAbuseNotifications.put(newFullAggregateGroupKey, ungrouped);
+
+                record.setOverrideGroupKey(null);
 
                 // Only add once, for triggering notification
                 if (!groupsToUpdate.containsKey(newFullAggregateGroupKey)) {
                     groupsToUpdate.put(newFullAggregateGroupKey,
-                            new GroupUpdateOp(newFullAggregateGroupKey, record, newGroupExists));
+                        new GroupUpdateOp(newFullAggregateGroupKey, record, hasSummary));
                 }
-
-                // Add notification to new group. do not request resort
-                record.setOverrideGroupKey(null);
-                mCallback.addAutoGroup(record.getKey(), newFullAggregateGroupKey.toString(), false);
             }
         }
 
@@ -959,18 +1011,26 @@
         for (FullyQualifiedGroupKey groupKey : groupsToUpdate.keySet()) {
             final ArrayMap<String, NotificationAttributes> aggregatedNotificationsAttrs =
                     mAggregatedNotifications.getOrDefault(groupKey, new ArrayMap<>());
-            if (aggregatedNotificationsAttrs.isEmpty()) {
-                mCallback.removeAutoGroupSummary(userId, pkgName, groupKey.toString());
-                mAggregatedNotifications.remove(groupKey);
-            } else {
-                NotificationRecord triggeringNotification = groupsToUpdate.get(groupKey).record;
-                boolean hasSummary = groupsToUpdate.get(groupKey).hasSummary;
+            final ArrayMap<String, NotificationAttributes> ungrouped =
+                    mUngroupedAbuseNotifications.getOrDefault(groupKey, new ArrayMap<>());
+
+            NotificationRecord triggeringNotification = groupsToUpdate.get(groupKey).record;
+            boolean hasSummary = groupsToUpdate.get(groupKey).hasSummary;
+            //Group needs to be created/updated
+            if (ungrouped.size() >= mAutoGroupAtCount
+                    || (hasSummary && !aggregatedNotificationsAttrs.isEmpty())) {
                 NotificationSectioner sectioner = getSection(triggeringNotification);
                 if (sectioner == null) {
                     continue;
                 }
-                updateAggregateAppGroup(groupKey, triggeringNotification.getKey(), hasSummary,
-                        sectioner.mSummaryId);
+                aggregateUngroupedNotifications(groupKey, triggeringNotification.getKey(),
+                        ungrouped, hasSummary, sectioner.mSummaryId);
+            } else {
+                // Remove empty groups
+                if (aggregatedNotificationsAttrs.isEmpty() && hasSummary) {
+                    mCallback.removeAutoGroupSummary(userId, pkgName, groupKey.toString());
+                    mAggregatedNotifications.remove(groupKey);
+                }
             }
         }
     }
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/GroupHelperTest.java b/services/tests/uiservicestests/src/com/android/server/notification/GroupHelperTest.java
index 51f64ba..3fc28f8 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/GroupHelperTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/GroupHelperTest.java
@@ -2204,7 +2204,7 @@
         verify(mCallback, times(1)).addAutoGroupSummary(anyInt(), eq(pkg), anyString(),
                 eq(expectedGroupKey_silent), anyInt(), eq(getNotificationAttributes(BASE_FLAGS)));
         verify(mCallback, times(AUTOGROUP_AT_COUNT)).addAutoGroup(anyString(),
-                eq(expectedGroupKey_silent), eq(false));
+                eq(expectedGroupKey_silent), eq(true));
 
         // Check that the alerting section group is removed
         verify(mCallback, times(1)).removeAutoGroupSummary(anyInt(), eq(pkg),
@@ -2264,13 +2264,15 @@
                 notificationList);
 
         // Check that channel1's notifications are moved to the silent section group
-        expectedSummaryAttr = new NotificationAttributes(BASE_FLAGS,
-                mSmallIcon, COLOR_DEFAULT, DEFAULT_VISIBILITY, DEFAULT_GROUP_ALERT,
-                "TEST_CHANNEL_ID1");
-        verify(mCallback, times(1)).addAutoGroupSummary(anyInt(), eq(pkg), anyString(),
-                eq(expectedGroupKey_silent), anyInt(), eq(expectedSummaryAttr));
-        verify(mCallback, times(AUTOGROUP_AT_COUNT/2 + 1)).addAutoGroup(anyString(),
-                eq(expectedGroupKey_silent), eq(false));
+        // But not enough to auto-group => remove override group key
+        verify(mCallback, never()).addAutoGroupSummary(anyInt(), eq(pkg), anyString(),
+                anyString(), anyInt(), any());
+        verify(mCallback, never()).addAutoGroup(anyString(), anyString(), anyBoolean());
+        for (NotificationRecord record: notificationList) {
+            if (record.getChannel().getId().equals(channel1.getId())) {
+                assertThat(record.getSbn().getOverrideGroupKey()).isNull();
+            }
+        }
 
         // Check that the alerting section group is not removed, only updated
         expectedSummaryAttr = new NotificationAttributes(BASE_FLAGS,
@@ -2343,7 +2345,7 @@
         verify(mCallback, times(1)).addAutoGroupSummary(anyInt(), eq(pkg), anyString(),
                 eq(expectedGroupKey_silent), anyInt(), eq(getNotificationAttributes(BASE_FLAGS)));
         verify(mCallback, times(numSilentGroupNotifications)).addAutoGroup(anyString(),
-                eq(expectedGroupKey_silent), eq(false));
+                eq(expectedGroupKey_silent), eq(true));
 
         // Check that the alerting section group is removed
         verify(mCallback, times(1)).removeAutoGroupSummary(anyInt(), eq(pkg),
@@ -2353,6 +2355,60 @@
     }
 
     @Test
+    @EnableFlags(FLAG_NOTIFICATION_FORCE_GROUPING)
+    public void testAutogroup_updateChannel_reachedMinAutogroupCount() {
+        final String pkg = "package";
+        final NotificationChannel channel1 = new NotificationChannel("TEST_CHANNEL_ID1",
+                "TEST_CHANNEL_ID1", IMPORTANCE_DEFAULT);
+        final NotificationChannel channel2 = new NotificationChannel("TEST_CHANNEL_ID2",
+                "TEST_CHANNEL_ID2", IMPORTANCE_LOW);
+        final List<NotificationRecord> notificationList = new ArrayList<>();
+        // Post notifications with different channels that would autogroup in different sections
+        NotificationRecord r;
+        // Not enough notifications to autogroup initially
+        for (int i = 0; i < AUTOGROUP_AT_COUNT; i++) {
+            if (i % 2 == 0) {
+                r = getNotificationRecord(pkg, i, String.valueOf(i),
+                    UserHandle.SYSTEM, null, false, channel1);
+            } else {
+                r = getNotificationRecord(pkg, i, String.valueOf(i),
+                    UserHandle.SYSTEM, null, false, channel2);
+            }
+            notificationList.add(r);
+            mGroupHelper.onNotificationPosted(r, false);
+        }
+        verify(mCallback, never()).addAutoGroupSummary(anyInt(), anyString(), anyString(),
+                anyString(), anyInt(), any());
+        verify(mCallback, never()).addAutoGroup(anyString(), anyString(), anyBoolean());
+        verify(mCallback, never()).removeAutoGroup(anyString());
+        verify(mCallback, never()).removeAutoGroupSummary(anyInt(), anyString(), anyString());
+        verify(mCallback, never()).updateAutogroupSummary(anyInt(), anyString(), anyString(),
+                any());
+        Mockito.reset(mCallback);
+
+        // Update channel1's importance
+        final String expectedGroupKey_silent = GroupHelper.getFullAggregateGroupKey(pkg,
+                AGGREGATE_GROUP_KEY + "SilentSection", UserHandle.SYSTEM.getIdentifier());
+        channel1.setImportance(IMPORTANCE_LOW);
+        for (NotificationRecord record: notificationList) {
+            if (record.getChannel().getId().equals(channel1.getId())) {
+                record.updateNotificationChannel(channel1);
+            }
+        }
+        mGroupHelper.onChannelUpdated(UserHandle.SYSTEM.getIdentifier(), pkg, channel1,
+                notificationList);
+
+        // Check that channel1's notifications are moved to the silent section & autogroup all
+        NotificationAttributes expectedSummaryAttr = new NotificationAttributes(BASE_FLAGS,
+                mSmallIcon, COLOR_DEFAULT, DEFAULT_VISIBILITY, DEFAULT_GROUP_ALERT,
+                "TEST_CHANNEL_ID1");
+        verify(mCallback, times(AUTOGROUP_AT_COUNT)).addAutoGroup(anyString(),
+                eq(expectedGroupKey_silent), eq(true));
+        verify(mCallback, times(1)).addAutoGroupSummary(anyInt(), eq(pkg), anyString(),
+                eq(expectedGroupKey_silent), anyInt(), eq(expectedSummaryAttr));
+    }
+
+    @Test
     @EnableFlags({FLAG_NOTIFICATION_FORCE_GROUPING,
             Flags.FLAG_NOTIFICATION_FORCE_GROUP_SINGLETONS})
     public void testNoGroup_singletonGroup_underLimit() {