Re-snooze Snoozed Notification

When Snoozing an already snoozed notification, the procedure is a NooP.
To fix this, we check if the notification we are trying to snooze is in
the snoozed notification list. We only do that if we cannot find the
notification in the displayed list (which was the only place where the
code looked for the notification). When we find the notification, we
update the snooze time.
Updated The code to look for snoozed group notifications
Bug: 70917141
Test: atest SnoozeHelperTest NotificationManagerServiceTest
Change-Id: I5ae17dcbc34ebccdcca6a64bb487fdef21703dd7

4 files changed, 196 insertions, 14 deletions

diff --git a/services/core/java/com/android/server/notification/NotificationManagerService.java b/services/core/java/com/android/server/notification/NotificationManagerService.java
index 217c0bdceede..6fe924e7432d 100644
--- a/services/core/java/com/android/server/notification/NotificationManagerService.java
+++ b/services/core/java/com/android/server/notification/NotificationManagerService.java
@@ -5229,7 +5229,7 @@ public class NotificationManagerService extends SystemService {
         @Override
         public void run() {
             synchronized (mNotificationLock) {
-                final NotificationRecord r = findNotificationByKeyLocked(mKey);
+                final NotificationRecord r = findInCurrentAndSnoozedNotificationByKeyLocked(mKey);
                 if (r != null) {
                     snoozeLocked(r);
                 }
@@ -5239,33 +5239,34 @@ public class NotificationManagerService extends SystemService {
         @GuardedBy("mNotificationLock")
         void snoozeLocked(NotificationRecord r) {
             if (r.sbn.isGroup()) {
-                final List<NotificationRecord> groupNotifications = findGroupNotificationsLocked(
+                final List<NotificationRecord> groupNotifications =
+                        findCurrentAndSnoozedGroupNotificationsLocked(
                         r.sbn.getPackageName(), r.sbn.getGroupKey(), r.sbn.getUserId());
                 if (r.getNotification().isGroupSummary()) {
-                    // snooze summary and all children
+                    // snooze all children
                     for (int i = 0; i < groupNotifications.size(); i++) {
-                        snoozeNotificationLocked(groupNotifications.get(i));
+                        if (mKey != groupNotifications.get(i).getKey()) {
+                            snoozeNotificationLocked(groupNotifications.get(i));
+                        }
                     }
                 } else {
                     // if there is a valid summary for this group, and we are snoozing the only
                     // child, also snooze the summary
                     if (mSummaryByGroupKey.containsKey(r.sbn.getGroupKey())) {
-                        if (groupNotifications.size() != 2) {
-                            snoozeNotificationLocked(r);
-                        } else {
+                        if (groupNotifications.size() == 2) {
                             // snooze summary and the one child
                             for (int i = 0; i < groupNotifications.size(); i++) {
-                                snoozeNotificationLocked(groupNotifications.get(i));
+                                if (mKey != groupNotifications.get(i).getKey()) {
+                                    snoozeNotificationLocked(groupNotifications.get(i));
+                                }
                             }
                         }
-                    } else {
-                        snoozeNotificationLocked(r);
                     }
                 }
-            } else {
-                // just snooze the one notification
-                snoozeNotificationLocked(r);
             }
+            // snooze the notification
+            snoozeNotificationLocked(r);
+
         }
 
         @GuardedBy("mNotificationLock")
@@ -7050,6 +7051,15 @@ public class NotificationManagerService extends SystemService {
     }
 
     @GuardedBy("mNotificationLock")
+    @NonNull
+    List<NotificationRecord> findCurrentAndSnoozedGroupNotificationsLocked(String pkg,
+            String groupKey, int userId) {
+        List<NotificationRecord> records = mSnoozeHelper.getNotifications(pkg, groupKey, userId);
+        records.addAll(findGroupNotificationsLocked(pkg, groupKey, userId));
+        return records;
+    }
+
+    @GuardedBy("mNotificationLock")
     @NonNull List<NotificationRecord> findGroupNotificationsLocked(String pkg,
             String groupKey, int userId) {
         List<NotificationRecord> records = new ArrayList<>();
@@ -7059,6 +7069,15 @@ public class NotificationManagerService extends SystemService {
         return records;
     }
 
+    @GuardedBy("mNotificationLock")
+    private NotificationRecord findInCurrentAndSnoozedNotificationByKeyLocked(String key) {
+        NotificationRecord r = findNotificationByKeyLocked(key);
+        if (r == null) {
+            r = mSnoozeHelper.getNotification(key);
+        }
+        return r;
+
+    }
 
     @GuardedBy("mNotificationLock")
     private @NonNull List<NotificationRecord> findGroupNotificationByListLocked(
diff --git a/services/core/java/com/android/server/notification/SnoozeHelper.java b/services/core/java/com/android/server/notification/SnoozeHelper.java
index abc98412e126..91f497cf9607 100644
--- a/services/core/java/com/android/server/notification/SnoozeHelper.java
+++ b/services/core/java/com/android/server/notification/SnoozeHelper.java
@@ -17,7 +17,6 @@ package com.android.server.notification;
 
 import android.annotation.NonNull;
 import android.app.AlarmManager;
-import android.app.Notification;
 import android.app.PendingIntent;
 import android.content.BroadcastReceiver;
 import android.content.Context;
@@ -104,6 +103,46 @@ public class SnoozeHelper {
         return Collections.EMPTY_LIST;
     }
 
+    @NonNull
+    ArrayList<NotificationRecord> getNotifications(String pkg,
+            String groupKey, Integer userId) {
+        ArrayList<NotificationRecord> records =  new ArrayList<>();
+        if (mSnoozedNotifications.containsKey(userId)
+                && mSnoozedNotifications.get(userId).containsKey(pkg)) {
+            ArrayMap<String, NotificationRecord> packages =
+                    mSnoozedNotifications.get(userId).get(pkg);
+            for (int i = 0; i < packages.size(); i++) {
+                String currentGroupKey = packages.valueAt(i).sbn.getGroup();
+                if (currentGroupKey.equals(groupKey)) {
+                    records.add(packages.valueAt(i));
+                }
+            }
+        }
+        return records;
+    }
+
+    protected NotificationRecord getNotification(String key) {
+        List<NotificationRecord> snoozedForUser = new ArrayList<>();
+        IntArray userIds = mUserProfiles.getCurrentProfileIds();
+        if (userIds != null) {
+            final int userIdsSize = userIds.size();
+            for (int i = 0; i < userIdsSize; i++) {
+                final ArrayMap<String, ArrayMap<String, NotificationRecord>> snoozedPkgs =
+                        mSnoozedNotifications.get(userIds.get(i));
+                if (snoozedPkgs != null) {
+                    final int snoozedPkgsSize = snoozedPkgs.size();
+                    for (int j = 0; j < snoozedPkgsSize; j++) {
+                        final ArrayMap<String, NotificationRecord> records = snoozedPkgs.valueAt(j);
+                        if (records != null) {
+                            return records.get(key);
+                        }
+                    }
+                }
+            }
+        }
+        return null;
+    }
+
     protected @NonNull List<NotificationRecord> getSnoozed() {
         List<NotificationRecord> snoozedForUser = new ArrayList<>();
         IntArray userIds = mUserProfiles.getCurrentProfileIds();
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
index f14e8d216cab..c1c0a308e48a 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
@@ -2005,6 +2005,73 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
     }
 
     @Test
+    public void testSnoozeRunnable_reSnoozeASingleSnoozedNotification() throws Exception {
+        final NotificationRecord notification = generateNotificationRecord(
+                mTestNotificationChannel, 1, null, true);
+        mService.addNotification(notification);
+        when(mSnoozeHelper.getNotification(any())).thenReturn(notification);
+
+        NotificationManagerService.SnoozeNotificationRunnable snoozeNotificationRunnable =
+                mService.new SnoozeNotificationRunnable(
+                notification.getKey(), 100, null);
+        snoozeNotificationRunnable.run();
+        NotificationManagerService.SnoozeNotificationRunnable snoozeNotificationRunnable2 =
+                mService.new SnoozeNotificationRunnable(
+                notification.getKey(), 100, null);
+        snoozeNotificationRunnable.run();
+
+        // snooze twice
+        verify(mSnoozeHelper, times(2)).snooze(any(NotificationRecord.class), anyLong());
+    }
+
+    @Test
+    public void testSnoozeRunnable_reSnoozeASnoozedNotificationWithGroupKey() throws Exception {
+        final NotificationRecord notification = generateNotificationRecord(
+                mTestNotificationChannel, 1, "group", true);
+        mService.addNotification(notification);
+        when(mSnoozeHelper.getNotification(any())).thenReturn(notification);
+
+        NotificationManagerService.SnoozeNotificationRunnable snoozeNotificationRunnable =
+                mService.new SnoozeNotificationRunnable(
+                notification.getKey(), 100, null);
+        snoozeNotificationRunnable.run();
+        NotificationManagerService.SnoozeNotificationRunnable snoozeNotificationRunnable2 =
+                mService.new SnoozeNotificationRunnable(
+                notification.getKey(), 100, null);
+        snoozeNotificationRunnable.run();
+
+        // snooze twice
+        verify(mSnoozeHelper, times(2)).snooze(any(NotificationRecord.class), anyLong());
+    }
+
+    @Test
+    public void testSnoozeRunnable_reSnoozeMultipleNotificationsWithGroupKey() throws Exception {
+        final NotificationRecord notification = generateNotificationRecord(
+                mTestNotificationChannel, 1, "group", true);
+        final NotificationRecord notification2 = generateNotificationRecord(
+                mTestNotificationChannel, 2, "group", true);
+        mService.addNotification(notification);
+        mService.addNotification(notification2);
+        when(mSnoozeHelper.getNotification(any())).thenReturn(notification);
+        when(mSnoozeHelper.getNotifications(
+                anyString(), anyString(), anyInt())).thenReturn(new ArrayList<>());
+
+        NotificationManagerService.SnoozeNotificationRunnable snoozeNotificationRunnable =
+                mService.new SnoozeNotificationRunnable(
+                        notification.getKey(), 100, null);
+        snoozeNotificationRunnable.run();
+        when(mSnoozeHelper.getNotifications(anyString(), anyString(), anyInt()))
+                .thenReturn(new ArrayList<>(Arrays.asList(notification, notification2)));
+        NotificationManagerService.SnoozeNotificationRunnable snoozeNotificationRunnable2 =
+                mService.new SnoozeNotificationRunnable(
+                        notification.getKey(), 100, null);
+        snoozeNotificationRunnable.run();
+
+        // snooze twice
+        verify(mSnoozeHelper, times(4)).snooze(any(NotificationRecord.class), anyLong());
+    }
+
+    @Test
     public void testSnoozeRunnable_snoozeNonGrouped() throws Exception {
         final NotificationRecord nonGrouped = generateNotificationRecord(
                 mTestNotificationChannel, 1, null, false);
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/SnoozeHelperTest.java b/services/tests/uiservicestests/src/com/android/server/notification/SnoozeHelperTest.java
index 1e645436e7ea..2e7277f5af01 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/SnoozeHelperTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/SnoozeHelperTest.java
@@ -264,6 +264,63 @@ public class SnoozeHelperTest extends UiServiceTestCase {
     }
 
     @Test
+    public void testGetSnoozedGroupNotifications() throws Exception {
+        IntArray profileIds = new IntArray();
+        profileIds.add(UserHandle.USER_CURRENT);
+        when(mUserProfiles.getCurrentProfileIds()).thenReturn(profileIds);
+        NotificationRecord r = getNotificationRecord("pkg", 1, "tag",
+                UserHandle.CURRENT, "group", true);
+        NotificationRecord r2 = getNotificationRecord("pkg", 2, "tag",
+                UserHandle.CURRENT, "group", true);
+        NotificationRecord r3 = getNotificationRecord("pkg2", 3, "tag",
+                UserHandle.CURRENT, "group", true);
+        NotificationRecord r4 = getNotificationRecord("pkg2", 4, "tag",
+                UserHandle.CURRENT, "group", true);
+        mSnoozeHelper.snooze(r, 1000);
+        mSnoozeHelper.snooze(r2, 1000);
+        mSnoozeHelper.snooze(r3, 1000);
+        mSnoozeHelper.snooze(r4, 1000);
+
+        assertEquals(2,
+                mSnoozeHelper.getNotifications("pkg", "group", UserHandle.USER_CURRENT).size());
+    }
+
+    @Test
+    public void testGetSnoozedNotificationByKey() throws Exception {
+        IntArray profileIds = new IntArray();
+        profileIds.add(UserHandle.USER_CURRENT);
+        when(mUserProfiles.getCurrentProfileIds()).thenReturn(profileIds);
+        NotificationRecord r = getNotificationRecord("pkg", 1, "tag",
+                UserHandle.CURRENT, "group", true);
+        NotificationRecord r2 = getNotificationRecord("pkg", 2, "tag",
+                UserHandle.CURRENT, "group", true);
+        NotificationRecord r3 = getNotificationRecord("pkg2", 3, "tag",
+                UserHandle.CURRENT, "group", true);
+        NotificationRecord r4 = getNotificationRecord("pkg2", 4, "tag",
+                UserHandle.CURRENT, "group", true);
+        mSnoozeHelper.snooze(r, 1000);
+        mSnoozeHelper.snooze(r2, 1000);
+        mSnoozeHelper.snooze(r3, 1000);
+        mSnoozeHelper.snooze(r4, 1000);
+
+        assertEquals(r, mSnoozeHelper.getNotification(r.getKey()));
+    }
+
+    @Test
+    public void testGetUnSnoozedNotificationByKey() throws Exception {
+        IntArray profileIds = new IntArray();
+        profileIds.add(UserHandle.USER_CURRENT);
+        when(mUserProfiles.getCurrentProfileIds()).thenReturn(profileIds);
+        NotificationRecord r = getNotificationRecord("pkg", 1, "tag",
+                UserHandle.CURRENT, "group", true);
+        NotificationRecord r2 = getNotificationRecord("pkg", 2, "tag",
+                UserHandle.CURRENT, "group", true);
+        mSnoozeHelper.snooze(r2, 1000);
+
+        assertEquals(null, mSnoozeHelper.getNotification(r.getKey()));
+    }
+
+    @Test
     public void repostGroupSummary_onlyFellowGroupChildren() throws Exception {
         NotificationRecord r = getNotificationRecord(
                 "pkg", 1, "one", UserHandle.SYSTEM, "group1", false);