sse sub resource

openapi_first/templates/vet_app/sse.py

@@ -1,44 +1,44 @@
-"""
-SSE broadcast for the animal-sounds worker.
-
-Not part of the openapi_first library API surface.
-"""
-
 import asyncio
 import random
 import json
 
-
-_sounds = ["woof", "meow", "coo"]
-_subscribers: list[asyncio.Queue] = []
-_worker_task: asyncio.Task | None = None
+_sounds_by_species = {"dog": ["woof"], "cat": ["meow"], "bird": ["coo"]}
+_subscribers: dict[int, list[asyncio.Queue]] = {}
+_worker_tasks: dict[int, asyncio.Task] = {}
 
 
-async def _sound_worker():
+async def _sound_worker(pet_id: int, species: str):
+    sounds = _sounds_by_species.get(species, ["woof"])
     while True:
-        sound = random.choice(_sounds)
+        sound = random.choice(sounds)
         data = json.dumps({"sound": sound})
-        for q in _subscribers:
+        queues = _subscribers.get(pet_id, [])
+        for q in queues:
             await q.put(data)
         await asyncio.sleep(random.uniform(1, 5))
 
 
-def start_worker():
-    global _worker_task
-    _worker_task = asyncio.create_task(_sound_worker())
+def _ensure_worker(pet_id: int, species: str):
+    if pet_id not in _worker_tasks or _worker_tasks[pet_id].done():
+        _subscribers.setdefault(pet_id, [])
+        _worker_tasks[pet_id] = asyncio.create_task(
+            _sound_worker(pet_id, species)
+        )
 
 
-def stop_worker():
-    if _worker_task is not None:
-        _worker_task.cancel()
-
-
-async def subscribe() -> asyncio.Queue:
+async def subscribe(pet_id: int, species: str) -> asyncio.Queue:
     q: asyncio.Queue = asyncio.Queue()
-    _subscribers.append(q)
+    _subscribers.setdefault(pet_id, []).append(q)
+    _ensure_worker(pet_id, species)
     return q
 
 
-def unsubscribe(q: asyncio.Queue):
-    if q in _subscribers:
-        _subscribers.remove(q)
+def unsubscribe(pet_id: int, q: asyncio.Queue):
+    queues = _subscribers.get(pet_id, [])
+    if q in queues:
+        queues.remove(q)
+    if not queues:
+        task = _worker_tasks.pop(pet_id, None)
+        if task and not task.done():
+            task.cancel()
+        _subscribers.pop(pet_id, None)