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