feat: add WebSocket announce manager

Made-with: Cursor
This commit is contained in:
cottongin
2026-03-12 07:14:04 -04:00
parent 7ed7ace578
commit 788225b3b6
2 changed files with 82 additions and 0 deletions

View File

@@ -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,
})

45
tests/test_websocket.py Normal file
View File

@@ -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