From 3cf79b35863eeb4c12e31f74afe8f4880a551e6b Mon Sep 17 00:00:00 2001
From: Hassan Malik <hbmalik88@gmail.com>
Date: Fri, 18 Oct 2024 18:30:20 -0400
Subject: [PATCH] add assertion to deleteNotificationById and add more tests

---
 .../NotificationServicesController.test.ts    | 47 +++++++++++++++++++
 .../NotificationServicesController.ts         | 12 +++++
 2 files changed, 59 insertions(+)

diff --git a/packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.test.ts b/packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.test.ts
index 99849c9730..d6ddc0d668 100644
--- a/packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.test.ts
+++ b/packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.test.ts
@@ -43,6 +43,7 @@ import * as OnChainNotifications from './services/onchain-notifications';
 import type { INotification } from './types';
 import type { UserStorage } from './types/user-storage/user-storage';
 import * as Utils from './utils/utils';
+import { processFeatureAnnouncement } from './processors';
 
 // Mock type used for testing purposes
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -606,6 +607,52 @@ describe('metamask-notifications - deleteNotificationById', () => {
 
     expect(controller.state.metamaskNotificationsList).toHaveLength(0);
   });
+
+  it('will throw if a notification is not found', async () => {
+    const { messenger } = mockNotificationMessenger();
+    const processedSnapNotification = processSnapNotification(
+      createMockSnapNotification(),
+    );
+    const controller = new NotificationServicesController({
+      messenger,
+      env: { featureAnnouncements: featureAnnouncementsEnv },
+      state: { metamaskNotificationsList: [processedSnapNotification] },
+    });
+
+    await expect(controller.deleteNotificationById('foo')).rejects.toThrow(
+      'The notification to be deleted does not exist.',
+    );
+
+    expect(controller.state.metamaskNotificationsList).toHaveLength(1);
+  });
+
+  it('will throw if the notification to be deleted is not locally persisted', async () => {
+    const { messenger } = mockNotificationMessenger();
+    const processedSnapNotification = processSnapNotification(
+      createMockSnapNotification(),
+    );
+    const processedFeatureAnnouncement = processFeatureAnnouncement(
+      createMockFeatureAnnouncementRaw(),
+    );
+    const controller = new NotificationServicesController({
+      messenger,
+      env: { featureAnnouncements: featureAnnouncementsEnv },
+      state: {
+        metamaskNotificationsList: [
+          processedFeatureAnnouncement,
+          processedSnapNotification,
+        ],
+      },
+    });
+
+    await expect(
+      controller.deleteNotificationById(processedFeatureAnnouncement.id),
+    ).rejects.toThrow(
+      'The notification type of "features_announcement" is not locally persisted, only the following types can use this function: snap.',
+    );
+
+    expect(controller.state.metamaskNotificationsList).toHaveLength(2);
+  });
 });
 
 describe('metamask-notifications - markMetamaskNotificationsAsRead()', () => {
diff --git a/packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.ts b/packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.ts
index 86a9c3345b..53f6a76e8b 100644
--- a/packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.ts
+++ b/packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.ts
@@ -174,6 +174,8 @@ export const defaultState: NotificationServicesControllerState = {
   isCheckingAccountsPresence: false,
 };
 
+const locallyPersistedNotificationTypes = new Set<string>([TRIGGER_TYPES.SNAP]);
+
 export type NotificationServicesControllerGetStateAction =
   ControllerGetStateAction<
     typeof controllerName,
@@ -1194,6 +1196,16 @@ export default class NotificationServicesController extends BaseController<
       'The notification to be deleted does not exist.',
     );
 
+    assert(
+      locallyPersistedNotificationTypes.has(fetchedNotification.type),
+      `The notification type of "${
+        // notifications are guaranteed to have type properties which equate to strings
+        fetchedNotification.type as string
+      }" is not locally persisted, only the following types can use this function: ${[
+        ...locallyPersistedNotificationTypes,
+      ].join(', ')}.`,
+    );
+
     const newList = this.state.metamaskNotificationsList.filter(
       (notification) => notification.id !== id,
     );