async def middle_stage(self, in_q, out_q, num, minsize):
async for batch in Stage.batches(in_q, minsize):
self.assertTrue(batch)
self.assertGreaterEqual(len(batch), min(minsize, num))
num -= len(batch)
for b in batch:
await out_q.put(b)
self.assertEqual(num, 0)
await out_q.put(None)
async def test_batch_and_single_none(self):
in_q = asyncio.Queue()
in_q.put_nowait(1)
in_q.put_nowait(2)
batch_it = Stage.batches(in_q)
self.assertEqual([1, 2], await batch_it.__anext__())
in_q.put_nowait(None)
with self.assertRaises(StopAsyncIteration):
await batch_it.__anext__()
async def test_two_batches(self):
in_q = asyncio.Queue()
in_q.put_nowait(1)
in_q.put_nowait(2)
batch_it = Stage.batches(in_q, minsize=1)
self.assertEqual([1, 2], await batch_it.__anext__())
in_q.put_nowait(3)
in_q.put_nowait(4)
in_q.put_nowait(None)
self.assertEqual([3, 4], await batch_it.__anext__())
with self.assertRaises(StopAsyncIteration):
await batch_it.__anext__()
class TestStage(asynctest.TestCase):
def setUp(self):
self.in_q = asyncio.Queue()
self.stage = Stage()
self.stage._connect(self.in_q, None)
async def test_none_only(self):
self.in_q.put_nowait(None)
batch_it = self.stage.batches(minsize=1)
with self.assertRaises(StopAsyncIteration):
await batch_it.__anext__()
async def test_single_batch_and_none(self):
c1 = mock.Mock(does_batch=True)
c2 = mock.Mock(does_batch=True)
self.in_q.put_nowait(c1)
self.in_q.put_nowait(c2)
self.in_q.put_nowait(None)
batch_it = self.stage.batches(minsize=1)
self.assertEqual([c1, c2], await batch_it.__anext__())
with self.assertRaises(StopAsyncIteration):
await batch_it.__anext__()
async def test_batch_and_single_none(self):
c1 = mock.Mock(does_batch=True)
c2 = mock.Mock(does_batch=True)
self.in_q.put_nowait(c1)
self.in_q.put_nowait(c2)
batch_it = self.stage.batches(minsize=1)
self.assertEqual([c1, c2], await batch_it.__anext__())
self.in_q.put_nowait(None)
with self.assertRaises(StopAsyncIteration):
await batch_it.__anext__()
async def test_two_batches(self):
c1 = mock.Mock(does_batch=True)
c2 = mock.Mock(does_batch=True)
c3 = mock.Mock(does_batch=True)
c4 = mock.Mock(does_batch=True)
self.in_q.put_nowait(c1)
self.in_q.put_nowait(c2)
batch_it = self.stage.batches(minsize=1)
self.assertEqual([c1, c2], await batch_it.__anext__())
self.in_q.put_nowait(c3)
self.in_q.put_nowait(c4)
self.in_q.put_nowait(None)
self.assertEqual([c3, c4], await batch_it.__anext__())
with self.assertRaises(StopAsyncIteration):
await batch_it.__anext__()
def setUp(self):
self.in_q = asyncio.Queue()
self.stage = Stage()
self.stage._connect(self.in_q, None)
class TestStage(asynctest.TestCase):
def setUp(self):
self.in_q = asyncio.Queue()
self.stage = Stage()
self.stage._connect(self.in_q, None)
async def test_none_only(self):
self.in_q.put_nowait(None)
batch_it = self.stage.batches(minsize=1)
with self.assertRaises(StopAsyncIteration):
await batch_it.__anext__()
async def test_single_batch_and_none(self):
c1 = mock.Mock()
c2 = mock.Mock()
self.in_q.put_nowait(c1)
self.in_q.put_nowait(c2)
self.in_q.put_nowait(None)
batch_it = self.stage.batches(minsize=1)
self.assertEqual([c1, c2], await batch_it.__anext__())
with self.assertRaises(StopAsyncIteration):
await batch_it.__anext__()
async def test_batch_and_single_none(self):
c1 = mock.Mock()
c2 = mock.Mock()
self.in_q.put_nowait(c1)
self.in_q.put_nowait(c2)
batch_it = self.stage.batches(minsize=1)
self.assertEqual([c1, c2], await batch_it.__anext__())
self.in_q.put_nowait(None)
with self.assertRaises(StopAsyncIteration):
await batch_it.__anext__()
async def test_two_batches(self):
c1 = mock.Mock()
c2 = mock.Mock()
c3 = mock.Mock()
c4 = mock.Mock()
self.in_q.put_nowait(c1)
self.in_q.put_nowait(c2)
batch_it = self.stage.batches(minsize=1)
self.assertEqual([c1, c2], await batch_it.__anext__())
self.in_q.put_nowait(c3)
self.in_q.put_nowait(c4)
self.in_q.put_nowait(None)
self.assertEqual([c3, c4], await batch_it.__anext__())
with self.assertRaises(StopAsyncIteration):
await batch_it.__anext__()
async def test_thaw_queue(self):
# Test that awaiting `resolved` on DeclarativeContent does not dead lock
c1 = DeclarativeContent(mock.Mock())
c2 = DeclarativeContent(mock.Mock())
c3 = DeclarativeContent(mock.Mock())
c4 = DeclarativeContent(mock.Mock())
batch_it = self.stage.batches(minsize=4)
fetch_task = asyncio.ensure_future(batch_it.__anext__())
fetch_task.add_done_callback(lambda _: c2.resolve())
self.in_q.put_nowait(c1)
self.in_q.put_nowait(c2)
self.in_q.put_nowait(c3)
wait_task = asyncio.ensure_future(c2.resolution())
wait_task.add_done_callback(lambda _: self.in_q.put_nowait(c4))
batch, c2_res = await asyncio.gather(fetch_task, wait_task)
self.assertEqual([c1, c2, c3], batch)
self.assertEqual(c2.content, c2_res)
self.in_q.put_nowait(None)
self.assertEqual([c4], await batch_it.__anext__())
with self.assertRaises(StopAsyncIteration):
await batch_it.__anext__()
async def test_none_only(self):
in_q = asyncio.Queue()
in_q.put_nowait(None)
batch_it = Stage.batches(in_q)
with self.assertRaises(StopAsyncIteration):
await batch_it.__anext__()
async def last_stage(self, in_q, num, minsize):
async for batch in Stage.batches(in_q, minsize):
self.assertTrue(batch)
self.assertGreaterEqual(len(batch), min(minsize, num))
num -= len(batch)
self.assertEqual(num, 0)
