diff --git a/src/nsct/providers/priority_queue.py b/src/nsct/providers/priority_queue.py index 9540d95..0d99d55 100644 --- a/src/nsct/providers/priority_queue.py +++ b/src/nsct/providers/priority_queue.py @@ -116,8 +116,8 @@ class PriorityLimiter: self._lock = asyncio.Lock() self._semaphore = asyncio.Semaphore(limit) - def __aenter__(self) -> "PriorityLimiter": + async def __aenter__(self) -> "PriorityLimiter": return self - def __aexit__(self, *exc: Any) -> None: + async def __aexit__(self, *exc: Any) -> None: pass \ No newline at end of file diff --git a/tests/test_priority_limiter.py b/tests/test_priority_limiter.py new file mode 100644 index 0000000..e52a412 --- /dev/null +++ b/tests/test_priority_limiter.py @@ -0,0 +1,93 @@ +"""Tests for PriorityLimiter async context manager (Stage 19 async_fix). + +PriorityLimiter.__aenter__/__aexit__ must be coroutines for use with +`async with` in Python 3.14+. +""" + +import asyncio + +from src.nsct.providers.priority_queue import ( + LlmPriorityRejectError, + Priority, + PriorityLimiter, +) + + +async def test_aenter_aexit_are_awaitable(): + """__aenter__ and __aexit__ must be coroutines (awaitable).""" + lim = PriorityLimiter(max_concurrency=1) + enter = lim.__aenter__() + exit_ = lim.__aexit__(None, None, None) + assert asyncio.iscoroutine(enter), "__aenter__ must return a coroutine" + assert asyncio.iscoroutine(exit_), "__aexit__ must return a coroutine" + # clean up the coroutines + await enter + await exit_ + + +async def test_async_with_roundtrip(): + """`async with PriorityLimiter(...)` works end-to-end.""" + lim = PriorityLimiter(max_concurrency=1) + async with lim: + assert isinstance(lim, PriorityLimiter) + # no exception after exit + + +async def test_async_with_concurrent(): + """Multiple concurrent `async with` blocks work without error.""" + + async def worker(lim: PriorityLimiter, n: int) -> int: + async with lim: + return n + + lim = PriorityLimiter(max_concurrency=3) + results = await asyncio.gather(*(worker(lim, i) for i in range(10))) + assert sorted(results) == list(range(10)) + + +async def test_aexit_ignores_exception(): + """__aexit__ returns None (does not swallow exceptions).""" + lim = PriorityLimiter(max_concurrency=1) + result = await lim.__aexit__(ValueError, ValueError("boom"), None) + assert result is None + + +async def test_high_priority_bypasses_limit(): + """HIGH requests acquire immediately even when NORMAL slots are full.""" + lim = PriorityLimiter(max_concurrency=1) + await lim.acquire(Priority.NORMAL) + # NORMAL slots exhausted; HIGH must still succeed immediately + await asyncio.wait_for(lim.acquire(Priority.HIGH), timeout=0.5) + lim.release(Priority.NORMAL) + + +async def test_low_priority_rejected_when_full(): + """LOW requests are rejected when all slots are occupied.""" + lim = PriorityLimiter(max_concurrency=1) + await lim.acquire(Priority.NORMAL) + try: + await lim.acquire(Priority.LOW) + assert False, "expected LlmPriorityRejectError" + except LlmPriorityRejectError: + pass + finally: + lim.release(Priority.NORMAL) + + +async def test_low_priority_succeeds_when_free(): + """LOW requests succeed when slots are available.""" + lim = PriorityLimiter(max_concurrency=2) + await lim.acquire(Priority.LOW) + lim.release(Priority.LOW) + + +async def test_active_count_tracks_slots(): + """active_count() reflects the number of held slots.""" + lim = PriorityLimiter(max_concurrency=3) + assert lim.active_count() == 0 + await lim.acquire(Priority.NORMAL) + await lim.acquire(Priority.NORMAL) + assert lim.active_count() == 2 + lim.release(Priority.NORMAL) + lim.release(Priority.NORMAL) + assert lim.active_count() == 0