From 788225b3b6ecb002fa629ae8223b846535c142f5 Mon Sep 17 00:00:00 2001
From: cottongin <cottongin@users.noreply.github.com>
Date: Thu, 12 Mar 2026 07:14:04 -0400
Subject: [PATCH] feat: add WebSocket announce manager

Made-with: Cursor
---
 src/ntr_fetcher/websocket.py | 37 +++++++++++++++++++++++++++++
 tests/test_websocket.py      | 45 ++++++++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+)
 create mode 100644 src/ntr_fetcher/websocket.py
 create mode 100644 tests/test_websocket.py

diff --git a/src/ntr_fetcher/websocket.py b/src/ntr_fetcher/websocket.py
new file mode 100644
index 0000000..23df44e
--- /dev/null
+++ b/src/ntr_fetcher/websocket.py
@@ -0,0 +1,37 @@
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class AnnounceManager:
+    def __init__(self):
+        self._subscribers: list = []
+
+    @property
+    def subscriber_count(self) -> int:
+        return len(self._subscribers)
+
+    def add_subscriber(self, websocket) -> None:
+        self._subscribers.append(websocket)
+        logger.info("Subscriber connected (%d total)", self.subscriber_count)
+
+    def remove_subscriber(self, websocket) -> None:
+        self._subscribers = [ws for ws in self._subscribers if ws is not websocket]
+        logger.info("Subscriber disconnected (%d total)", self.subscriber_count)
+
+    async def broadcast(self, message: dict) -> None:
+        dead = []
+        for ws in self._subscribers:
+            try:
+                await ws.send_json(message)
+            except Exception:
+                dead.append(ws)
+                logger.warning("Removing dead subscriber")
+        for ws in dead:
+            self.remove_subscriber(ws)
+
+    async def broadcast_status(self) -> None:
+        await self.broadcast({
+            "type": "status",
+            "subscribers": self.subscriber_count,
+        })
diff --git a/tests/test_websocket.py b/tests/test_websocket.py
new file mode 100644
index 0000000..aa86d08
--- /dev/null
+++ b/tests/test_websocket.py
@@ -0,0 +1,45 @@
+import pytest
+from ntr_fetcher.websocket import AnnounceManager
+
+
+@pytest.fixture
+def manager():
+    return AnnounceManager()
+
+
+def test_no_subscribers_initially(manager):
+    assert manager.subscriber_count == 0
+
+
+@pytest.mark.asyncio
+async def test_subscribe_and_broadcast(manager):
+    received = []
+
+    class FakeWS:
+        async def send_json(self, data):
+            received.append(data)
+
+    ws = FakeWS()
+    manager.add_subscriber(ws)
+    assert manager.subscriber_count == 1
+
+    await manager.broadcast({"type": "announce", "message": "Now Playing: Song #1"})
+    assert len(received) == 1
+    assert received[0]["message"] == "Now Playing: Song #1"
+
+    manager.remove_subscriber(ws)
+    assert manager.subscriber_count == 0
+
+
+@pytest.mark.asyncio
+async def test_broadcast_skips_dead_connections(manager):
+    class DeadWS:
+        async def send_json(self, data):
+            raise Exception("connection closed")
+
+    ws = DeadWS()
+    manager.add_subscriber(ws)
+    assert manager.subscriber_count == 1
+
+    await manager.broadcast({"type": "announce", "message": "test"})
+    assert manager.subscriber_count == 0