fix: priority_queue async context methods (async_fix for Python 3.14+). Add test_priority_limiter.py tests for async context manager

This commit is contained in:
NSCT Agent
2026-09-05 13:36:49 +00:00
parent 6e2e7386ad
commit f9b761ced7
2 changed files with 95 additions and 2 deletions

View File

@@ -116,8 +116,8 @@ class PriorityLimiter:
self._lock = asyncio.Lock() self._lock = asyncio.Lock()
self._semaphore = asyncio.Semaphore(limit) self._semaphore = asyncio.Semaphore(limit)
def __aenter__(self) -> "PriorityLimiter": async def __aenter__(self) -> "PriorityLimiter":
return self return self
def __aexit__(self, *exc: Any) -> None: async def __aexit__(self, *exc: Any) -> None:
pass pass

View File

@@ -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