83 lines
3.1 KiB
Python
83 lines
3.1 KiB
Python
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()
|