From 17ee5c2ab8032506674b065b047add5bee82c835 Mon Sep 17 00:00:00 2001
From: Josh Hawkins <32435876+hawkeye217@users.noreply.github.com>
Date: Sun, 15 Dec 2024 12:24:03 -0600
Subject: [PATCH] backend

---
 frigate/comms/dispatcher.py |   5 ++
 frigate/comms/webpush.py    | 111 ++++++++++++++++++++++--------------
 frigate/const.py            |   1 +
 3 files changed, 74 insertions(+), 43 deletions(-)

diff --git a/frigate/comms/dispatcher.py b/frigate/comms/dispatcher.py
index 2bddc97a5..d42df289d 100644
--- a/frigate/comms/dispatcher.py
+++ b/frigate/comms/dispatcher.py
@@ -13,6 +13,7 @@ from frigate.const import (
     CLEAR_ONGOING_REVIEW_SEGMENTS,
     INSERT_MANY_RECORDINGS,
     INSERT_PREVIEW,
+    NOTIFICATION_TEST,
     REQUEST_REGION_GRID,
     UPDATE_CAMERA_ACTIVITY,
     UPDATE_EMBEDDINGS_REINDEX_PROGRESS,
@@ -191,6 +192,9 @@ class Dispatcher:
                 json.dumps(self.embeddings_reindex.copy()),
             )
 
+        def handle_notification_test():
+            self.publish("notification_test", "Test notification")
+
         # Dictionary mapping topic to handlers
         topic_handlers = {
             INSERT_MANY_RECORDINGS: handle_insert_many_recordings,
@@ -202,6 +206,7 @@ class Dispatcher:
             UPDATE_EVENT_DESCRIPTION: handle_update_event_description,
             UPDATE_MODEL_STATE: handle_update_model_state,
             UPDATE_EMBEDDINGS_REINDEX_PROGRESS: handle_update_embeddings_reindex_progress,
+            NOTIFICATION_TEST: handle_notification_test,
             "restart": handle_restart,
             "embeddingsReindexProgress": handle_embeddings_reindex_progress,
             "modelState": handle_model_state,
diff --git a/frigate/comms/webpush.py b/frigate/comms/webpush.py
index 602f8d11e..132e85ccc 100644
--- a/frigate/comms/webpush.py
+++ b/frigate/comms/webpush.py
@@ -116,17 +116,71 @@ class WebPushClient(Communicator):  # type: ignore[misc]
 
         if topic == "reviews":
             self.send_alert(json.loads(payload))
+        elif topic == "notification_test":
+            self.send_notification_test()
 
-    def send_alert(self, payload: dict[str, any]) -> None:
+    def send_push_notification(
+        self,
+        user: str,
+        payload: dict[str, Any],
+        title: str,
+        message: str,
+        direct_url: str = "",
+        image: str = "",
+        notification_type: str = "alert",
+        ttl: int = 0,
+    ) -> None:
+        for pusher in self.web_pushers[user]:
+            endpoint = pusher.subscription_info["endpoint"]
+            headers = self.claim_headers[endpoint[: endpoint.index("/", 10)]].copy()
+            headers["urgency"] = "high"
+
+            resp = pusher.send(
+                headers=headers,
+                ttl=ttl,
+                data=json.dumps(
+                    {
+                        "title": title,
+                        "message": message,
+                        "direct_url": direct_url,
+                        "image": image,
+                        "id": payload.get("after", {}).get("id", ""),
+                        "type": notification_type,
+                    }
+                ),
+            )
+
+            if resp.status_code in (404, 410):
+                self.expired_subs.setdefault(user, []).append(endpoint)
+            elif resp.status_code != 201:
+                logger.warning(
+                    f"Failed to send notification to {user} :: {resp.headers}"
+                )
+
+    def send_notification_test(self) -> None:
         if not self.config.notifications.email:
             return
 
         self.check_registrations()
 
-        # Only notify for alerts
-        if payload["after"]["severity"] != "alert":
+        for user in self.web_pushers:
+            self.send_push_notification(
+                user=user,
+                payload={},
+                title="Test Notification",
+                message="This is a test notification from Frigate.",
+                notification_type="test",
+            )
+
+    def send_alert(self, payload: dict[str, Any]) -> None:
+        if (
+            not self.config.notifications.email
+            or payload["after"]["severity"] != "alert"
+        ):
             return
 
+        self.check_registrations()
+
         state = payload["type"]
 
         # Don't notify if message is an update and important fields don't have an update
@@ -155,47 +209,18 @@ class WebPushClient(Communicator):  # type: ignore[misc]
 
         # if event is ongoing open to live view otherwise open to recordings view
         direct_url = f"/review?id={reviewId}" if state == "end" else f"/#{camera}"
+        ttl = 3600 if state == "end" else 0
 
-        for user, pushers in self.web_pushers.items():
-            for pusher in pushers:
-                endpoint = pusher.subscription_info["endpoint"]
-
-                # set headers for notification behavior
-                headers = self.claim_headers[
-                    endpoint[0 : endpoint.index("/", 10)]
-                ].copy()
-                headers["urgency"] = "high"
-                ttl = 3600 if state == "end" else 0
-
-                # send message
-                resp = pusher.send(
-                    headers=headers,
-                    ttl=ttl,
-                    data=json.dumps(
-                        {
-                            "title": title,
-                            "message": message,
-                            "direct_url": direct_url,
-                            "image": image,
-                            "id": reviewId,
-                            "type": "alert",
-                        }
-                    ),
-                )
-
-                if resp.status_code == 201:
-                    pass
-                elif resp.status_code == 404 or resp.status_code == 410:
-                    # subscription is not found or has been unsubscribed
-                    if not self.expired_subs.get(user):
-                        self.expired_subs[user] = []
-
-                    self.expired_subs[user].append(pusher.subscription_info["endpoint"])
-                    # the subscription no longer exists and should be removed
-                else:
-                    logger.warning(
-                        f"Failed to send notification to {user} :: {resp.headers}"
-                    )
+        for user in self.web_pushers:
+            self.send_push_notification(
+                user=user,
+                payload=payload,
+                title=title,
+                message=message,
+                direct_url=direct_url,
+                image=image,
+                ttl=ttl,
+            )
 
         self.cleanup_registrations()
 
diff --git a/frigate/const.py b/frigate/const.py
index 5976f47b1..71063f25a 100644
--- a/frigate/const.py
+++ b/frigate/const.py
@@ -102,6 +102,7 @@ UPDATE_CAMERA_ACTIVITY = "update_camera_activity"
 UPDATE_EVENT_DESCRIPTION = "update_event_description"
 UPDATE_MODEL_STATE = "update_model_state"
 UPDATE_EMBEDDINGS_REINDEX_PROGRESS = "handle_embeddings_reindex_progress"
+NOTIFICATION_TEST = "notification_test"
 
 # Stats Values