Persist snoozed notifications after restart

The state of snoozed notifications is lost after restart. To prevent
snoozed notifications from sending notifications, we intercept posted
notifications and re-snooze them if they were supposed to be snoozed
before restart/turning off. We store that state in the
notifications_polixy.xml file. (usually in /data/system)

Bug: 120617189
Test: atest SnoozeHelperTest NotificationManagerServiceTest
Change-Id: I3a19ad2f0ba98ab385a0ff12a5c7affba6e4f1e5

4 files changed, 387 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 6fe924e7432d..da5104a613ee 100644
--- a/services/core/java/com/android/server/notification/NotificationManagerService.java
+++ b/services/core/java/com/android/server/notification/NotificationManagerService.java
@@ -621,6 +621,8 @@ public class NotificationManagerService extends SystemService {
                 mConditionProviders.readXml(
                         parser, mAllowedManagedServicePackages, forRestore, userId);
                 migratedManagedServices = true;
+            } else if (mSnoozeHelper.XML_TAG_NAME.equals(parser.getName())) {
+                mSnoozeHelper.readXml(parser);
             }
             if (LOCKSCREEN_ALLOW_SECURE_NOTIFICATIONS_TAG.equals(parser.getName())) {
                 if (forRestore && userId != UserHandle.USER_SYSTEM) {
@@ -709,6 +711,7 @@ public class NotificationManagerService extends SystemService {
         mPreferencesHelper.writeXml(out, forBackup, userId);
         mListeners.writeXml(out, forBackup, userId);
         mAssistants.writeXml(out, forBackup, userId);
+        mSnoozeHelper.writeXml(out);
         mConditionProviders.writeXml(out, forBackup, userId);
         if (!forBackup || userId == UserHandle.USER_SYSTEM) {
             writeSecureNotificationsPolicy(out);
@@ -1753,6 +1756,7 @@ public class NotificationManagerService extends SystemService {
                 com.android.internal.R.integer.config_notificationWarnRemoteViewSizeBytes);
         mStripRemoteViewsSizeBytes = getContext().getResources().getInteger(
                 com.android.internal.R.integer.config_notificationStripRemoteViewSizeBytes);
+
     }
 
     @Override
@@ -5284,7 +5288,7 @@ public class NotificationManagerService extends SystemService {
             updateLightsLocked();
             if (mSnoozeCriterionId != null) {
                 mAssistants.notifyAssistantSnoozedLocked(r.sbn, mSnoozeCriterionId);
-                mSnoozeHelper.snooze(r);
+                mSnoozeHelper.snooze(r, mSnoozeCriterionId);
             } else {
                 mSnoozeHelper.snooze(r, mDuration);
             }
@@ -5387,6 +5391,27 @@ public class NotificationManagerService extends SystemService {
         @Override
         public void run() {
             synchronized (mNotificationLock) {
+                final Long snoozeAt =
+                        mSnoozeHelper.getSnoozeTimeForUnpostedNotification(
+                                r.getUser().getIdentifier(),
+                                r.sbn.getPackageName(), r.sbn.getKey());
+                final long currentTime = System.currentTimeMillis();
+                if (snoozeAt.longValue() > currentTime) {
+                    (new SnoozeNotificationRunnable(r.sbn.getKey(),
+                            snoozeAt.longValue() - currentTime, null)).snoozeLocked(r);
+                    return;
+                }
+
+                final String contextId =
+                        mSnoozeHelper.getSnoozeContextForUnpostedNotification(
+                                r.getUser().getIdentifier(),
+                                r.sbn.getPackageName(), r.sbn.getKey());
+                if (contextId != null) {
+                    (new SnoozeNotificationRunnable(r.sbn.getKey(),
+                            0, contextId)).snoozeLocked(r);
+                    return;
+                }
+
                 mEnqueuedNotifications.add(r);
                 scheduleTimeoutLocked(r);
 
@@ -6937,6 +6962,7 @@ public class NotificationManagerService extends SystemService {
         if (DBG) {
             Slog.d(TAG, String.format("unsnooze event(%s, %s)", key, listenerName));
         }
+        mSnoozeHelper.cleanupPersistedContext(key);
         mSnoozeHelper.repost(key);
         handleSavePolicyFile();
     }
diff --git a/services/core/java/com/android/server/notification/SnoozeHelper.java b/services/core/java/com/android/server/notification/SnoozeHelper.java
index 91f497cf9607..8125d0d653ad 100644
--- a/services/core/java/com/android/server/notification/SnoozeHelper.java
+++ b/services/core/java/com/android/server/notification/SnoozeHelper.java
@@ -55,6 +55,21 @@ import java.util.Set;
  * NotificationManagerService helper for handling snoozed notifications.
  */
 public class SnoozeHelper {
+    public static final String XML_SNOOZED_NOTIFICATION_VERSION = "1";
+
+    protected static final String XML_TAG_NAME = "snoozed-notifications";
+
+    private static final String XML_SNOOZED_NOTIFICATION = "notification";
+    private static final String XML_SNOOZED_NOTIFICATION_CONTEXT = "context";
+    private static final String XML_SNOOZED_NOTIFICATION_PKG = "pkg";
+    private static final String XML_SNOOZED_NOTIFICATION_USER_ID = "user-id";
+    private static final String XML_SNOOZED_NOTIFICATION_KEY = "key";
+    //the time the snoozed notification should be reposted
+    private static final String XML_SNOOZED_NOTIFICATION_TIME = "time";
+    private static final String XML_SNOOZED_NOTIFICATION_CONTEXT_ID = "id";
+    private static final String XML_SNOOZED_NOTIFICATION_VERSION_LABEL = "version";
+
+
     private static final String TAG = "SnoozeHelper";
     private static final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG);
     private static final String INDENT = "    ";
@@ -72,6 +87,17 @@ public class SnoozeHelper {
     // User id : package name : notification key : record.
     private ArrayMap<Integer, ArrayMap<String, ArrayMap<String, NotificationRecord>>>
             mSnoozedNotifications = new ArrayMap<>();
+    // User id : package name : notification key : time-milliseconds .
+    // This member stores persisted snoozed notification trigger times. it persists through reboots
+    // It should have the notifications that haven't expired or re-posted yet
+    private ArrayMap<Integer, ArrayMap<String, ArrayMap<String, Long>>>
+            mPersistedSnoozedNotifications = new ArrayMap<>();
+    // User id : package name : notification key : creation ID .
+    // This member stores persisted snoozed notification trigger context for the assistant
+    // it persists through reboots.
+    // It should have the notifications that haven't expired or re-posted yet
+    private ArrayMap<Integer, ArrayMap<String, ArrayMap<String, String>>>
+            mPersistedSnoozedNotificationsWithContext = new ArrayMap<>();
     // notification key : package.
     private ArrayMap<String, String> mPackages = new ArrayMap<>();
     // key : userId
@@ -89,6 +115,34 @@ public class SnoozeHelper {
         mUserProfiles = userProfiles;
     }
 
+    void cleanupPersistedContext(String key){
+        int userId = mUsers.get(key);
+        String pkg = mPackages.get(key);
+        synchronized (mPersistedSnoozedNotificationsWithContext) {
+            removeRecord(pkg, key, userId, mPersistedSnoozedNotificationsWithContext);
+        }
+    }
+
+    //This function has a side effect of removing the time from the list of persisted notifications.
+    //IT IS NOT IDEMPOTENT!
+    @NonNull
+    protected Long getSnoozeTimeForUnpostedNotification(int userId, String pkg, String key) {
+        Long time;
+        synchronized (mPersistedSnoozedNotifications) {
+            time = removeRecord(pkg, key, userId, mPersistedSnoozedNotifications);
+        }
+        if (time == null) {
+            return 0L;
+        }
+        return time;
+    }
+
+    protected String getSnoozeContextForUnpostedNotification(int userId, String pkg, String key) {
+        synchronized (mPersistedSnoozedNotificationsWithContext) {
+            return removeRecord(pkg, key, userId, mPersistedSnoozedNotificationsWithContext);
+        }
+    }
+
     protected boolean isSnoozed(int userId, String pkg, String key) {
         return mSnoozedNotifications.containsKey(userId)
                 && mSnoozedNotifications.get(userId).containsKey(pkg)
@@ -169,32 +223,82 @@ public class SnoozeHelper {
      * Snoozes a notification and schedules an alarm to repost at that time.
      */
     protected void snooze(NotificationRecord record, long duration) {
+        String pkg = record.sbn.getPackageName();
+        String key = record.getKey();
+        int userId = record.getUser().getIdentifier();
+
         snooze(record);
-        scheduleRepost(record.sbn.getPackageName(), record.getKey(), record.getUserId(), duration);
+        scheduleRepost(pkg, key, userId, duration);
+        Long activateAt = System.currentTimeMillis() + duration;
+        synchronized (mPersistedSnoozedNotifications) {
+            storeRecord(pkg, key, userId, mPersistedSnoozedNotifications, activateAt);
+        }
     }
 
     /**
      * Records a snoozed notification.
      */
-    protected void snooze(NotificationRecord record) {
+    protected void snooze(NotificationRecord record, String contextId) {
+        int userId = record.getUser().getIdentifier();
+        if (contextId != null) {
+            synchronized (mPersistedSnoozedNotificationsWithContext) {
+                storeRecord(record.sbn.getPackageName(), record.getKey(),
+                        userId, mPersistedSnoozedNotificationsWithContext, contextId);
+            }
+        }
+        snooze(record);
+    }
+
+    private void snooze(NotificationRecord record) {
         int userId = record.getUser().getIdentifier();
         if (DEBUG) {
             Slog.d(TAG, "Snoozing " + record.getKey());
         }
-        ArrayMap<String, ArrayMap<String, NotificationRecord>> records =
-                mSnoozedNotifications.get(userId);
+        storeRecord(record.sbn.getPackageName(), record.getKey(),
+                userId, mSnoozedNotifications, record);
+        mPackages.put(record.getKey(), record.sbn.getPackageName());
+        mUsers.put(record.getKey(), userId);
+    }
+
+    private <T> void storeRecord(String pkg, String key, Integer userId,
+            ArrayMap<Integer, ArrayMap<String, ArrayMap<String, T>>> targets, T object) {
+
+        ArrayMap<String, ArrayMap<String, T>> records =
+                targets.get(userId);
         if (records == null) {
             records = new ArrayMap<>();
         }
-        ArrayMap<String, NotificationRecord> pkgRecords = records.get(record.sbn.getPackageName());
+        ArrayMap<String, T> pkgRecords = records.get(pkg);
         if (pkgRecords == null) {
             pkgRecords = new ArrayMap<>();
         }
-        pkgRecords.put(record.getKey(), record);
-        records.put(record.sbn.getPackageName(), pkgRecords);
-        mSnoozedNotifications.put(userId, records);
-        mPackages.put(record.getKey(), record.sbn.getPackageName());
-        mUsers.put(record.getKey(), userId);
+        pkgRecords.put(key, object);
+        records.put(pkg, pkgRecords);
+        targets.put(userId, records);
+
+    }
+
+    private <T> T removeRecord(String pkg, String key, Integer userId,
+            ArrayMap<Integer, ArrayMap<String, ArrayMap<String, T>>> targets) {
+        T object = null;
+
+        ArrayMap<String, ArrayMap<String, T>> records =
+                targets.get(userId);
+        if (records == null) {
+            return null;
+        }
+        ArrayMap<String, T> pkgRecords = records.get(pkg);
+        if (pkgRecords == null) {
+            return null;
+        }
+        object = pkgRecords.remove(key);
+        if (pkgRecords.size() == 0) {
+            records.remove(pkg);
+        }
+        if (records.size() == 0) {
+            targets.remove(userId);
+        }
+        return object;
     }
 
     protected boolean cancel(int userId, String pkg, String tag, int id) {
@@ -414,13 +518,121 @@ public class SnoozeHelper {
         }
     }
 
-    protected void writeXml(XmlSerializer out, boolean forBackup) throws IOException {
+    protected void writeXml(XmlSerializer out) throws IOException {
+        final long currentTime = System.currentTimeMillis();
+        out.startTag(null, XML_TAG_NAME);
+        writeXml(out, mPersistedSnoozedNotifications, XML_SNOOZED_NOTIFICATION,
+                value -> {
+                    if (value < currentTime) {
+                        return;
+                    }
+                    out.attribute(null, XML_SNOOZED_NOTIFICATION_TIME,
+                            value.toString());
+                });
+        writeXml(out, mPersistedSnoozedNotificationsWithContext, XML_SNOOZED_NOTIFICATION_CONTEXT,
+                value -> {
+                    out.attribute(null, XML_SNOOZED_NOTIFICATION_CONTEXT_ID,
+                            value);
+                });
+        out.endTag(null, XML_TAG_NAME);
+    }
+
+    private interface Inserter<T> {
+        void insert(T t) throws IOException;
+    }
+    private <T> void writeXml(XmlSerializer out,
+            ArrayMap<Integer, ArrayMap<String, ArrayMap<String, T>>> targets, String tag,
+            Inserter<T> attributeInserter)
+            throws IOException {
+        synchronized (targets) {
+            final int M = targets.size();
+            for (int i = 0; i < M; i++) {
+                final ArrayMap<String, ArrayMap<String, T>> packages =
+                        targets.valueAt(i);
+                if (packages == null) {
+                    continue;
+                }
+                final int N = packages.size();
+                for (int j = 0; j < N; j++) {
+                    final ArrayMap<String, T> keyToValue = packages.valueAt(j);
+                    if (keyToValue == null) {
+                        continue;
+                    }
+                    final int O = keyToValue.size();
+                    for (int k = 0; k < O; k++) {
+                        T value = keyToValue.valueAt(k);
+
+                        out.startTag(null, tag);
 
+                        attributeInserter.insert(value);
+
+                        out.attribute(null, XML_SNOOZED_NOTIFICATION_VERSION_LABEL,
+                                XML_SNOOZED_NOTIFICATION_VERSION);
+                        out.attribute(null, XML_SNOOZED_NOTIFICATION_KEY, keyToValue.keyAt(k));
+                        out.attribute(null, XML_SNOOZED_NOTIFICATION_PKG, packages.keyAt(j));
+                        out.attribute(null, XML_SNOOZED_NOTIFICATION_USER_ID,
+                                targets.keyAt(i).toString());
+
+                        out.endTag(null, tag);
+
+                    }
+                }
+            }
+        }
     }
 
-    public void readXml(XmlPullParser parser, boolean forRestore)
+    protected void readXml(XmlPullParser parser)
             throws XmlPullParserException, IOException {
+        int type;
+        while ((type = parser.next()) != XmlPullParser.END_DOCUMENT) {
+            String tag = parser.getName();
+            if (type == XmlPullParser.END_TAG
+                    && XML_TAG_NAME.equals(tag)) {
+                break;
+            }
+            if (type == XmlPullParser.START_TAG
+                    && (XML_SNOOZED_NOTIFICATION.equals(tag)
+                        || tag.equals(XML_SNOOZED_NOTIFICATION_CONTEXT))
+                    && parser.getAttributeValue(null, XML_SNOOZED_NOTIFICATION_VERSION_LABEL)
+                        .equals(XML_SNOOZED_NOTIFICATION_VERSION)) {
+                try {
+                    final String key = parser.getAttributeValue(null, XML_SNOOZED_NOTIFICATION_KEY);
+                    final String pkg = parser.getAttributeValue(null, XML_SNOOZED_NOTIFICATION_PKG);
+                    final int userId = Integer.parseInt(
+                            parser.getAttributeValue(null, XML_SNOOZED_NOTIFICATION_USER_ID));
+                    if (tag.equals(XML_SNOOZED_NOTIFICATION)) {
+                        final Long time = Long.parseLong(
+                                parser.getAttributeValue(null, XML_SNOOZED_NOTIFICATION_TIME));
+                        if (time > System.currentTimeMillis()) { //only read new stuff
+                            synchronized (mPersistedSnoozedNotifications) {
+                                storeRecord(pkg, key, userId, mPersistedSnoozedNotifications, time);
+                            }
+                            scheduleRepost(pkg, key, userId, time - System.currentTimeMillis());
+                        }
+                        continue;
+                    }
+                    if (tag.equals(XML_SNOOZED_NOTIFICATION_CONTEXT)) {
+                        final String creationId = parser.getAttributeValue(
+                                null, XML_SNOOZED_NOTIFICATION_CONTEXT_ID);
+                        synchronized (mPersistedSnoozedNotificationsWithContext) {
+                            storeRecord(pkg, key, userId, mPersistedSnoozedNotificationsWithContext,
+                                    creationId);
+                        }
+                        continue;
+                    }
+
 
+                } catch (Exception e) {
+                    //we dont cre if it is a number format exception or a null pointer exception.
+                    //we just want to debug it and continue with our lives
+                    if (DEBUG) {
+                        Slog.d(TAG,
+                                "Exception in reading snooze data from policy xml: "
+                                        + e.getMessage());
+                    }
+                }
+            }
+        }
     }
 
     @VisibleForTesting
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 c1c0a308e48a..57caa1db423c 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
@@ -2766,6 +2766,18 @@ public class NotificationManagerServiceTest extends UiServiceTestCase {
     }
 
     @Test
+    public void testReadPolicyXml_readSnoozedNotificationsFromXml() throws Exception {
+        final String upgradeXml = "<notification-policy version=\"1\">"
+                + "<snoozed-notifications>></snoozed-notifications>"
+                + "</notification-policy>";
+        mService.readPolicyXml(
+                new BufferedInputStream(new ByteArrayInputStream(upgradeXml.getBytes())),
+                false,
+                UserHandle.USER_ALL);
+        verify(mSnoozeHelper, times(1)).readXml(any(XmlPullParser.class));
+    }
+
+    @Test
     public void testReadPolicyXml_readApprovedServicesFromSettings() throws Exception {
         final String preupgradeXml = "<notification-policy version=\"1\">"
                 + "<ranking></ranking>"
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 2e7277f5af01..36175a93f667 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/SnoozeHelperTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/SnoozeHelperTest.java
@@ -18,6 +18,7 @@ package com.android.server.notification;
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertTrue;
+import static junit.framework.Assert.assertNull;
 
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
@@ -37,9 +38,11 @@ import android.os.UserHandle;
 import android.service.notification.StatusBarNotification;
 import android.test.suitebuilder.annotation.SmallTest;
 import android.util.IntArray;
+import android.util.Xml;
 
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.internal.util.FastXmlSerializer;
 import com.android.server.UiServiceTestCase;
 
 import org.junit.Before;
@@ -48,6 +51,15 @@ import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+import org.xmlpull.v1.XmlSerializer;
+
+import java.io.BufferedInputStream;
+import java.io.BufferedOutputStream;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
@@ -69,6 +81,117 @@ public class SnoozeHelperTest extends UiServiceTestCase {
     }
 
     @Test
+    public void testWriteXMLformattedCorrectly_testReadingCorrectTime()
+            throws XmlPullParserException, IOException {
+        final String max_time_str = Long.toString(Long.MAX_VALUE);
+        final String xml_string = "<snoozed-notifications>"
+                + "<notification version=\"1\" user-id=\"0\" notification=\"notification\" "
+                + "pkg=\"pkg\" key=\"key\" time=\"" + max_time_str + "\"/>"
+                + "<notification version=\"1\" user-id=\"0\" notification=\"notification\" "
+                + "pkg=\"pkg\" key=\"key2\" time=\"" + max_time_str + "\"/>"
+                + "</snoozed-notifications>";
+        XmlPullParser parser = Xml.newPullParser();
+        parser.setInput(new BufferedInputStream(
+                new ByteArrayInputStream(xml_string.getBytes())), null);
+        mSnoozeHelper.readXml(parser);
+        assertTrue("Should read the notification time from xml and it should be more than zero",
+                0 < mSnoozeHelper.getSnoozeTimeForUnpostedNotification(
+                        0, "pkg", "key").doubleValue());
+    }
+
+    @Test
+    public void testWriteXMLformattedCorrectly_testCorrectContextURI()
+            throws XmlPullParserException, IOException {
+        final String max_time_str = Long.toString(Long.MAX_VALUE);
+        final String xml_string = "<snoozed-notifications>"
+                + "<context version=\"1\" user-id=\"0\" notification=\"notification\" "
+                + "pkg=\"pkg\" key=\"key\" id=\"uri\"/>"
+                + "<context version=\"1\" user-id=\"0\" notification=\"notification\" "
+                + "pkg=\"pkg\" key=\"key2\" id=\"uri\"/>"
+                + "</snoozed-notifications>";
+        XmlPullParser parser = Xml.newPullParser();
+        parser.setInput(new BufferedInputStream(
+                new ByteArrayInputStream(xml_string.getBytes())), null);
+        mSnoozeHelper.readXml(parser);
+        assertEquals("Should read the notification context from xml and it should be `uri",
+                "uri", mSnoozeHelper.getSnoozeContextForUnpostedNotification(
+                        0, "pkg", "key"));
+    }
+
+    @Test
+    public void testReadValidSnoozedFromCorrectly_timeDeadline()
+            throws XmlPullParserException, IOException {
+        NotificationRecord r = getNotificationRecord("pkg", 1, "one", UserHandle.SYSTEM);
+        mSnoozeHelper.snooze(r, 999999999);
+        XmlSerializer serializer = new FastXmlSerializer();
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        serializer.setOutput(new BufferedOutputStream(baos), "utf-8");
+        serializer.startDocument(null, true);
+        mSnoozeHelper.writeXml(serializer);
+        serializer.endDocument();
+        serializer.flush();
+
+        XmlPullParser parser = Xml.newPullParser();
+        parser.setInput(new BufferedInputStream(
+                new ByteArrayInputStream(baos.toByteArray())), "utf-8");
+        mSnoozeHelper.readXml(parser);
+        assertTrue("Should read the notification time from xml and it should be more than zero",
+                0 < mSnoozeHelper.getSnoozeTimeForUnpostedNotification(
+                        0, "pkg", r.getKey()).doubleValue());
+    }
+
+
+    @Test
+    public void testReadExpiredSnoozedNotification() throws
+            XmlPullParserException, IOException, InterruptedException {
+        NotificationRecord r = getNotificationRecord("pkg", 1, "one", UserHandle.SYSTEM);
+        mSnoozeHelper.snooze(r, 0);
+       // Thread.sleep(100);
+        XmlSerializer serializer = new FastXmlSerializer();
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        serializer.setOutput(new BufferedOutputStream(baos), "utf-8");
+        serializer.startDocument(null, true);
+        mSnoozeHelper.writeXml(serializer);
+        serializer.endDocument();
+        serializer.flush();
+        Thread.sleep(10);
+        XmlPullParser parser = Xml.newPullParser();
+        parser.setInput(new BufferedInputStream(
+                new ByteArrayInputStream(baos.toByteArray())), "utf-8");
+        mSnoozeHelper.readXml(parser);
+        int systemUser = UserHandle.SYSTEM.getIdentifier();
+        assertTrue("Should see a past time returned",
+                System.currentTimeMillis() >  mSnoozeHelper.getSnoozeTimeForUnpostedNotification(
+                        systemUser, "pkg", r.getKey()).longValue());
+    }
+
+    @Test
+    public void testCleanupContextShouldRemovePersistedRecord() {
+        NotificationRecord r = getNotificationRecord("pkg", 1, "one", UserHandle.SYSTEM);
+        mSnoozeHelper.snooze(r, "context");
+        mSnoozeHelper.cleanupPersistedContext(r.sbn.getKey());
+        assertNull(mSnoozeHelper.getSnoozeContextForUnpostedNotification(
+                r.getUser().getIdentifier(),
+                r.sbn.getPackageName(),
+                r.sbn.getKey()
+        ));
+    }
+
+    @Test
+    public void testReadNoneSnoozedNotification() throws XmlPullParserException,
+            IOException, InterruptedException {
+        NotificationRecord r = getNotificationRecord(
+                "pkg", 1, "one", UserHandle.SYSTEM);
+        mSnoozeHelper.snooze(r, 0);
+
+        assertEquals("should see a zero value for unsnoozed notification",
+                0L,
+                mSnoozeHelper.getSnoozeTimeForUnpostedNotification(
+                        UserHandle.SYSTEM.getIdentifier(),
+                        "not_my_package", r.getKey()).longValue());
+    }
+
+    @Test
     public void testSnoozeForTime() throws Exception {
         NotificationRecord r = getNotificationRecord("pkg", 1, "one", UserHandle.SYSTEM);
         mSnoozeHelper.snooze(r, 1000);
@@ -84,7 +207,7 @@ public class SnoozeHelperTest extends UiServiceTestCase {
     @Test
     public void testSnooze() throws Exception {
         NotificationRecord r = getNotificationRecord("pkg", 1, "one", UserHandle.SYSTEM);
-        mSnoozeHelper.snooze(r);
+        mSnoozeHelper.snooze(r, (String) null);
         verify(mAm, never()).setExactAndAllowWhileIdle(
                 anyInt(), anyLong(), any(PendingIntent.class));
         assertTrue(mSnoozeHelper.isSnoozed(