File size: 2,097 Bytes
0efd0f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | 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)
|