feat: add WebSocket announce manager
Made-with: Cursor
This commit is contained in:
37
src/ntr_fetcher/websocket.py
Normal file
37
src/ntr_fetcher/websocket.py
Normal 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
45
tests/test_websocket.py
Normal 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
|
||||
Reference in New Issue
Block a user