fix: check if daemon task is still alive while waiting, otherwise we …

saq/multiplexer.py

@@ -46,12 +46,13 @@ async def listen(self, *channels: str, timeout: float | None = None) -> t.AsyncG
         queue = await self.subscribe(*channels)
 
         try:
-            while True:
-                if timeout:
-                    yield await asyncio.wait_for(queue.get(), timeout)
-                else:
-                    yield await queue.get()
-                queue.task_done()
+            while self._daemon_task and not self._daemon_task.done():
+                try:
+                    yield await asyncio.wait_for(queue.get(), timeout if timeout else 1.0)
+                    queue.task_done()
+                except asyncio.TimeoutError:
+                    if timeout:
+                        raise
         finally:
             await self.unsubscribe(queue)
 

tests/test_queue.py

@@ -787,3 +787,13 @@ async def test_priority(self) -> None:
         assert await self.enqueue("test", priority=-1)
         self.assertEqual(await self.count("queued"), 1)
         assert not await self.queue.dequeue(0.01)
+
+    async def test_listen_cancel(self) -> None:
+        job = await self.queue.enqueue("test")
+        task = asyncio.create_task(job.refresh(0))
+        await asyncio.sleep(0.1)
+        await self.queue._listener.close()
+        await task
+
+        with self.assertRaises(asyncio.TimeoutError):
+            await job.refresh(0.1)