async def test_consumer_gracefull_shutdown(ack_deadline_cache, message): queue = asyncio.Queue() ack_queue = asyncio.Queue() nack_queue = asyncio.Queue() mock = MagicMock() event = asyncio.Event() async def f(*args): mock(*args) await event.wait() consumer_task = asyncio.ensure_future( consumer(queue, f, ack_queue, ack_deadline_cache, 1, nack_queue, MagicMock())) await queue.put((message, 0.0)) await asyncio.sleep(0.1) mock.assert_called_once() consumer_task.cancel() await asyncio.sleep(0.1) assert not consumer_task.done() event.set() await asyncio.sleep(0) assert ack_queue.qsize() == 1 await ack_queue.get() ack_queue.task_done() await asyncio.sleep(0.1) assert consumer_task.done()
async def test_consumer_handles_callback_exception_nack( ack_deadline_cache, message ): queue = asyncio.Queue() ack_queue = asyncio.Queue() nack_queue = asyncio.Queue() mock = MagicMock() async def f(*args): mock(*args) raise RuntimeError consumer_task = asyncio.ensure_future( consumer(queue, f, ack_queue, ack_deadline_cache, 1, nack_queue, MagicMock())) await queue.put((message, 0.0)) await asyncio.sleep(0.1) consumer_task.cancel() mock.assert_called_once() assert ack_queue.qsize() == 0 assert nack_queue.qsize() == 1 assert queue.qsize() == 0 # cleanup nack_queue.get_nowait() nack_queue.task_done() with pytest.raises(asyncio.CancelledError): await asyncio.wait_for(consumer_task, 1)
async def test_consumer_calls_none_means_ack(ack_deadline_cache, message, application_callback): queue = asyncio.Queue() ack_queue = asyncio.Queue() nack_queue = asyncio.Queue() consumer_task = asyncio.ensure_future( consumer(queue, application_callback, ack_queue, ack_deadline_cache, 1, nack_queue, MagicMock())) await queue.put((message, 0.0)) await asyncio.sleep(0) consumer_task.cancel() result = await asyncio.wait_for(ack_queue.get(), 1) assert result == 'ack_id' ack_queue.task_done() application_callback.assert_called_once() assert queue.qsize() == 0 assert nack_queue.qsize() == 0 # verify that the consumer shuts down gracefully with pytest.raises(asyncio.CancelledError): await asyncio.wait_for(consumer_task, 1)
async def test_consumer_tasks_limited_by_pool_size(ack_deadline_cache): queue = asyncio.Queue() ack_queue = asyncio.Queue() async def callback(mock): mock() await asyncio.sleep(10) mock1 = MagicMock() mock1.ack_id = 'ack_id' mock2 = MagicMock() mock2.ack_id = 'ack_id' mock3 = MagicMock() mock3.ack_id = 'ack_id' mock4 = MagicMock() mock4.ack_id = 'ack_id' asyncio.ensure_future( consumer(queue, callback, ack_queue, ack_deadline_cache, 2, None, MagicMock())) for m in [mock1, mock2, mock3, mock4]: await queue.put((m, 0.0)) await asyncio.sleep(0.1) mock1.assert_called_once() mock2.assert_called_once() mock3.assert_not_called() assert queue.qsize() == 1 assert ack_queue.qsize() == 0
async def test_consumer_handles_callback_exception_no_nack( ack_deadline_cache, message ): queue = asyncio.Queue() ack_queue = asyncio.Queue() mock = MagicMock() async def f(*args): mock(*args) raise RuntimeError consumer_task = asyncio.ensure_future( consumer( queue, f, ack_queue, ack_deadline_cache, 1, None, MagicMock() ) ) await queue.put((message, 0.0)) await asyncio.sleep(0.1) consumer_task.cancel() mock.assert_called_once() assert ack_queue.qsize() == 0 assert queue.qsize() == 0
async def test_consumer_drops_expired_messages(ack_deadline_cache, message, application_callback): f = asyncio.Future() f.set_result(0.0) ack_deadline_cache.get = MagicMock(return_value=f) queue = asyncio.Queue() ack_queue = asyncio.Queue() nack_queue = asyncio.Queue() consumer_task = asyncio.ensure_future( consumer( queue, application_callback, ack_queue, ack_deadline_cache, 1, nack_queue, MagicMock() ) ) await queue.put((message, 0.0)) await asyncio.sleep(0) consumer_task.cancel() application_callback.assert_not_called() assert ack_queue.qsize() == 0 assert nack_queue.qsize() == 0 assert queue.qsize() == 0
async def test_consumer_calls_none_means_ack(ack_deadline_cache, message, application_callback): queue = asyncio.Queue() ack_queue = asyncio.Queue() nack_queue = asyncio.Queue() consumer_task = asyncio.ensure_future( consumer( queue, application_callback, ack_queue, ack_deadline_cache, 1, nack_queue, MagicMock() ) ) await queue.put((message, 0.0)) await asyncio.sleep(0) consumer_task.cancel() result = await asyncio.wait_for(ack_queue.get(), 1) assert result == 'ack_id' application_callback.assert_called_once() assert queue.qsize() == 0 assert nack_queue.qsize() == 0
async def test_consumer_gracefull_shutdown_without_pending_tasks( ack_deadline_cache): queue = asyncio.Queue() ack_queue = asyncio.Queue() nack_queue = asyncio.Queue() consumer_task = asyncio.ensure_future( consumer(queue, lambda _x: None, ack_queue, ack_deadline_cache, 1, nack_queue, MagicMock())) await asyncio.sleep(0.1) consumer_task.cancel() await asyncio.sleep(0.1) assert consumer_task.done()
async def test_consumer_tasks_limited_by_pool_size(ack_deadline_cache): queue = asyncio.Queue() ack_queue = asyncio.Queue() pause = asyncio.Event() async def callback(mock): mock() await pause.wait() mock1 = MagicMock() mock1.ack_id = 'ack_id' mock2 = MagicMock() mock2.ack_id = 'ack_id' mock3 = MagicMock() mock3.ack_id = 'ack_id' mock4 = MagicMock() mock4.ack_id = 'ack_id' consumer_task = asyncio.ensure_future( consumer(queue, callback, ack_queue, ack_deadline_cache, 2, None, MagicMock())) for m in [mock1, mock2, mock3, mock4]: await queue.put((m, 0.0)) await asyncio.sleep(0.1) mock1.assert_called_once() mock2.assert_called_once() mock3.assert_not_called() mock4.assert_not_called() assert queue.qsize() == 1 # two running, one dequeued assert ack_queue.qsize() == 0 # clean up pause.set() await queue.join() while not ack_queue.empty(): await ack_queue.get() ack_queue.task_done() consumer_task.cancel() with pytest.raises(asyncio.CancelledError): await asyncio.wait_for(consumer_task, 1)