Update app/services/alert_hub.py

app/services/alert_hub.py

@@ -1,48 +1,51 @@
-from collections import defaultdict
-from typing import Any
-
-from fastapi import WebSocket
-
-
-class AlertHub:
-    def __init__(self) -> None:
-        self.admin_clients: set[WebSocket] = set()
-        self.device_clients: dict[str, set[WebSocket]] = defaultdict(set)
-
-    async def connect_admin(self, websocket: WebSocket) -> None:
-        await websocket.accept()
-        self.admin_clients.add(websocket)
-
-    async def connect_device(self, device_id: str, websocket: WebSocket) -> None:
-        await websocket.accept()
-        self.device_clients[device_id].add(websocket)
-
-    def disconnect_admin(self, websocket: WebSocket) -> None:
-        self.admin_clients.discard(websocket)
-
-    def disconnect_device(self, device_id: str, websocket: WebSocket) -> None:
-        if device_id in self.device_clients:
-            self.device_clients[device_id].discard(websocket)
-
-    async def broadcast_admin(self, payload: dict[str, Any]) -> None:
-        stale: list[WebSocket] = []
-        for client in self.admin_clients:
-            try:
-                await client.send_json(payload)
-            except Exception:
-                stale.append(client)
-        for client in stale:
-            self.disconnect_admin(client)
-
-    async def send_device(self, device_id: str, payload: dict[str, Any]) -> None:
-        stale: list[WebSocket] = []
-        for client in self.device_clients.get(device_id, set()):
-            try:
-                await client.send_json(payload)
-            except Exception:
-                stale.append(client)
-        for client in stale:
-            self.disconnect_device(device_id, client)
-
-
-alert_hub = AlertHub()
+from collections import defaultdict
+from typing import Any
+
+from fastapi import WebSocket
+
+
+class AlertHub:
+    def __init__(self) -> None:
+        self.admin_clients: set[WebSocket] = set()
+        self.device_clients: dict[str, set[WebSocket]] = defaultdict(set)
+
+    async def connect_admin(self, websocket: WebSocket) -> None:
+        await websocket.accept()
+        self.admin_clients.add(websocket)
+
+    async def connect_device(self, device_id: str, websocket: WebSocket) -> None:
+        await websocket.accept()
+        self.device_clients[device_id].add(websocket)
+
+    def disconnect_admin(self, websocket: WebSocket) -> None:
+        self.admin_clients.discard(websocket)
+
+    def disconnect_device(self, device_id: str, websocket: WebSocket) -> None:
+        if device_id in self.device_clients:
+            self.device_clients[device_id].discard(websocket)
+
+    async def broadcast_admin(self, payload: dict[str, Any]) -> None:
+        stale: list[WebSocket] = []
+        for client in self.admin_clients:
+            try:
+                await client.send_json(payload)
+            except Exception:
+                stale.append(client)
+        for client in stale:
+            self.disconnect_admin(client)
+
+    async def send_device(self, device_id: str, payload: dict[str, Any]) -> None:
+        stale: list[WebSocket] = []
+        for client in self.device_clients.get(device_id, set()):
+            try:
+                await client.send_json(payload)
+            except Exception:
+                stale.append(client)
+        for client in stale:
+            self.disconnect_device(device_id, client)
+
+    def device_connection_count(self, device_id: str) -> int:
+        return len(self.device_clients.get(device_id, set()))
+
+
+alert_hub = AlertHub()