| import asyncio,unittest |
| from unittest.mock import AsyncMock |
| from contextlib import asynccontextmanager |
| from shared_turn import admitted_turn |
|
|
| class Client: |
| def __init__(self,busy=0,statuses=(200,)): |
| self.busy=busy;self.statuses=list(statuses);self.posts=0 |
| async def get(self,path,**kwargs): |
| busy=self.busy>0;self.busy-=1 |
| class State: |
| status_code=200 |
| def json(self):return {'busy':busy,'ready':True} |
| return State() |
| @asynccontextmanager |
| async def stream(self,*args,**kwargs): |
| self.posts+=1;code=self.statuses.pop(0) |
| class Response: |
| status_code=code |
| def raise_for_status(self): |
| if code>=400:raise ValueError('Non-retryable status') |
| yield Response() |
|
|
| class AdmissionTests(unittest.IsolatedAsyncioTestCase): |
| async def test_busy_and_admission_race_retry_only_before_success(self): |
| c=Client(busy=3,statuses=(409,503,200));notice=AsyncMock() |
| async with admitted_turn(c,{},notice,timeout=.2,poll=.001) as response:self.assertEqual(response.status_code,200) |
| self.assertEqual(c.posts,3);notice.assert_awaited() |
| async def test_accepted_stream_failure_is_never_retried(self): |
| c=Client(statuses=(200,200)) |
| with self.assertRaises(RuntimeError): |
| async with admitted_turn(c,{},AsyncMock(),poll=.001):raise RuntimeError('Accepted stream disconnected') |
| self.assertEqual(c.posts,1) |
| async def test_wait_is_bounded_and_does_not_send_busy_generation(self): |
| c=Client(busy=100) |
| with self.assertRaises(TimeoutError): |
| async with admitted_turn(c,{},AsyncMock(),timeout=.01,poll=.001):pass |
| self.assertEqual(c.posts,0) |
| async def test_wait_is_immediately_cancellable(self): |
| c=Client(busy=100) |
| async def wait(): |
| async with admitted_turn(c,{},AsyncMock(),poll=.001):pass |
| task=asyncio.create_task(wait());await asyncio.sleep(.005);task.cancel() |
| with self.assertRaises(asyncio.CancelledError):await task |
| self.assertEqual(c.posts,0) |
|
|
