Update notification group mapping when group key changes

Tested by running:
"runtest -x cts/tests/app/src/android/app/cts/NotificationManagerTest.java
-m testAutogrouping_autogroupStaysUntilAllNotificationsCanceled"

The test fails before this change, and passes afterwards.

Also manually verified that the original bug that introduced this
regression (b/65100024) remains fixed.

Bug: 68215987
Change-Id: I4af198fd5257ad1dfab5e052c11f58ea84b5095a
diff --git a/src/com/android/launcher3/notification/NotificationListener.java b/src/com/android/launcher3/notification/NotificationListener.java
index 114b2b8..cbdabf3 100644
--- a/src/com/android/launcher3/notification/NotificationListener.java
+++ b/src/com/android/launcher3/notification/NotificationListener.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.notification;
 
+import static com.android.launcher3.SettingsActivity.NOTIFICATION_BADGING;
+
 import android.annotation.TargetApi;
 import android.app.Notification;
 import android.app.NotificationChannel;
@@ -43,8 +45,6 @@
 import java.util.Map;
 import java.util.Set;
 
-import static com.android.launcher3.SettingsActivity.NOTIFICATION_BADGING;
-
 /**
  * A {@link NotificationListenerService} that sends updates to its
  * {@link NotificationsChangedListener} when notifications are posted or canceled,
@@ -71,6 +71,8 @@
     private final Ranking mTempRanking = new Ranking();
     /** Maps groupKey's to the corresponding group of notifications. */
     private final Map<String, NotificationGroup> mNotificationGroupMap = new HashMap<>();
+    /** Maps keys to their corresponding current group key */
+    private final Map<String, String> mNotificationGroupKeyMap = new HashMap<>();
 
     private SettingsObserver mNotificationBadgingObserver;
 
@@ -258,6 +260,48 @@
         }
     }
 
+    @Override
+    public void onNotificationRankingUpdate(RankingMap rankingMap) {
+        super.onNotificationRankingUpdate(rankingMap);
+        String[] keys = rankingMap.getOrderedKeys();
+        for (StatusBarNotification sbn : getActiveNotifications(keys)) {
+            updateGroupKeyIfNecessary(sbn);
+        }
+    }
+
+    private void updateGroupKeyIfNecessary(StatusBarNotification sbn) {
+        String childKey = sbn.getKey();
+        String oldGroupKey = mNotificationGroupKeyMap.get(childKey);
+        String newGroupKey = sbn.getGroupKey();
+        if (oldGroupKey == null || !oldGroupKey.equals(newGroupKey)) {
+            // The group key has changed.
+            mNotificationGroupKeyMap.put(childKey, newGroupKey);
+            if (oldGroupKey != null && mNotificationGroupMap.containsKey(oldGroupKey)) {
+                // Remove the child key from the old group.
+                NotificationGroup oldGroup = mNotificationGroupMap.get(oldGroupKey);
+                oldGroup.removeChildKey(childKey);
+                if (oldGroup.isEmpty()) {
+                    mNotificationGroupMap.remove(oldGroupKey);
+                }
+            }
+        }
+        if (sbn.isGroup() && newGroupKey != null) {
+            // Maintain group info so we can cancel the summary when the last child is canceled.
+            NotificationGroup notificationGroup = mNotificationGroupMap.get(newGroupKey);
+            if (notificationGroup == null) {
+                notificationGroup = new NotificationGroup();
+                mNotificationGroupMap.put(newGroupKey, notificationGroup);
+            }
+            boolean isGroupSummary = (sbn.getNotification().flags
+                    & Notification.FLAG_GROUP_SUMMARY) != 0;
+            if (isGroupSummary) {
+                notificationGroup.setGroupSummaryKey(childKey);
+            } else {
+                notificationGroup.addChildKey(childKey);
+            }
+        }
+    }
+
     /** This makes a potentially expensive binder call and should be run on a background thread. */
     public List<StatusBarNotification> getNotificationsForKeys(List<NotificationKeyData> keys) {
         StatusBarNotification[] notifications = NotificationListener.this
@@ -295,20 +339,7 @@
     private boolean shouldBeFilteredOut(StatusBarNotification sbn) {
         Notification notification = sbn.getNotification();
 
-        boolean isGroupHeader = (notification.flags & Notification.FLAG_GROUP_SUMMARY) != 0;
-        if (sbn.isGroup()) {
-            // Maintain group info so we can cancel the summary when the last child is canceled.
-            NotificationGroup notificationGroup = mNotificationGroupMap.get(sbn.getGroupKey());
-            if (notificationGroup == null) {
-                notificationGroup = new NotificationGroup();
-                mNotificationGroupMap.put(sbn.getGroupKey(), notificationGroup);
-            }
-            if (isGroupHeader) {
-                notificationGroup.setGroupSummaryKey(sbn.getKey());
-            } else {
-                notificationGroup.addChildKey(sbn.getKey());
-            }
-        }
+        updateGroupKeyIfNecessary(sbn);
 
         getCurrentRanking().getRanking(sbn.getKey(), mTempRanking);
         if (!mTempRanking.canShowBadge()) {
@@ -324,6 +355,7 @@
         CharSequence title = notification.extras.getCharSequence(Notification.EXTRA_TITLE);
         CharSequence text = notification.extras.getCharSequence(Notification.EXTRA_TEXT);
         boolean missingTitleAndText = TextUtils.isEmpty(title) && TextUtils.isEmpty(text);
+        boolean isGroupHeader = (notification.flags & Notification.FLAG_GROUP_SUMMARY) != 0;
         return (isGroupHeader || missingTitleAndText);
     }