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)