feat: improve library and audiobook workflow

server/audiobook/cmp/test_dynamic_queue.py

@@ -0,0 +1,82 @@
+import asyncio
+import json
+import sys
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'backend'))
+from services.cmp_multigpu import warm_chapters
+
+
+class DynamicQueueTests(unittest.IsolatedAsyncioTestCase):
+    async def test_free_gpu_continues_and_reports_before_slow_gpu_finishes(self):
+        slow = asyncio.Event()
+        dispatched = []
+        active = set()
+
+        class Response:
+            def raise_for_status(self): pass
+            def json(self): return {'duration_s': 1, 'cached': False}
+
+        class Client:
+            def __init__(self, **kwargs): pass
+            async def __aenter__(self): return self
+            async def __aexit__(self, *args): pass
+            async def post(self, url, json):
+                self_outer.assertNotIn(url, active)
+                active.add(url)
+                index = json['chapter_index']
+                dispatched.append((url, index))
+                try:
+                    if index == 0:
+                        await slow.wait()
+                    else:
+                        await asyncio.sleep(.001)
+                    return Response()
+                finally:
+                    active.remove(url)
+
+        self_outer = self
+        with patch('services.cmp_multigpu.httpx.AsyncClient', Client):
+            stream = warm_chapters('text', 6, ['gpu0', 'gpu1', 'gpu2'], {})
+            try:
+                first = await asyncio.wait_for(anext(stream), .3)
+                self.assertTrue(first.startswith('data:'))
+                self.assertFalse(slow.is_set())
+                while len(dispatched) < 6:
+                    await asyncio.wait_for(anext(stream), .3)
+                self.assertTrue(any(index >= 3 and url != 'gpu0/audiobook/preview' for url, index in dispatched))
+                slow.set()
+                remaining = [event async for event in stream]
+                self.assertTrue(any('"index": 0' in event for event in remaining))
+                self.assertEqual(list(range(6)), sorted(index for _, index in dispatched))
+            finally:
+                slow.set()
+                await stream.aclose()
+        self.assertFalse(active)
+
+    async def test_closing_stream_cancels_inflight_requests(self):
+        stopped = asyncio.Event()
+        class Response:
+            def raise_for_status(self): pass
+            def json(self): return {'duration_s': 1}
+        class Client:
+            def __init__(self, **kwargs): pass
+            async def __aenter__(self): return self
+            async def __aexit__(self, *args): pass
+            async def post(self, url, json):
+                if json['chapter_index'] == 0:
+                    try: await asyncio.Event().wait()
+                    finally: stopped.set()
+                return Response()
+        with patch('services.cmp_multigpu.httpx.AsyncClient', Client):
+            stream = warm_chapters('text', 2, ['gpu0', 'gpu1'], {})
+            try:
+                await asyncio.wait_for(anext(stream), .3)
+            finally:
+                await stream.aclose()
+            self.assertTrue(stopped.is_set())
+
+
+if __name__ == '__main__': unittest.main()