async def test__iter_messages(mocker, event_loop):
# arrange
queue = mocker.Mock(spec=asynqp.Queue)
queue.consume.return_value = future()
consumer = get_consumer(callback=simple_callback, prefetch_count=0)
mocker.patch.object(consumer, '_queue', new=queue)
queue = mocker.patch('asynqp_consumer.consumer.asyncio.Queue').return_value
queue.get.side_effect = iter([
future(Message(mock.Mock(spec=asynqp.IncomingMessage))),
future(Message(mock.Mock(spec=asynqp.IncomingMessage))),
future(exception=SomeException),
])
# act
result = []
with pytest.raises(SomeException):
messages_iterator = await consumer._get_messages_iterator(
loop=event_loop)
async for message in messages_iterator:
result.append(message)
# assert
assert len(result) == 2
assert isinstance(result[0], Message)
assert isinstance(result[1], Message)
async def test_run__func_raised_exception__reject_all_messages(
self, mocker):
# arrange
message1 = Message(
body='message1',
channel=mocker.sentinel.channel,
envelope=mocker.sentinel.envelope,
properties=mocker.sentinel.properties,
)
mocker.patch.object(message1, 'reject', return_value=future())
message2 = Message(
body='message2',
channel=mocker.sentinel.channel,
envelope=mocker.sentinel.envelope,
properties=mocker.sentinel.properties,
)
mocker.patch.object(
message2,
'reject',
return_value=future(exception=MessageAlreadyResolved()))
middleware = ProcessBulk(
lambda messages: future(exception=Exception()))
inp = make_iterator([[message1, message2]])
# act
out = middleware(inp)
# assert
assert await collect_iterator(out) == [None]
message1.reject.assert_called_once_with()
message2.reject.assert_called_once_with()
async def test_connect_and_open_channel(mocker):
# arrange
transport = mocker.Mock(spec=asyncio.Transport)
channel = mocker.Mock(spec=Channel)
protocol = mocker.Mock(spec=aioamqp.AmqpProtocol)
protocol.channel.return_value = future(channel)
mocker.patch.object(aioamqp,
'connect',
return_value=future((transport, protocol)),
auto_spec=True)
connection_params = ConnectionParams()
# act
result = await connect_and_open_channel(
connection_params=connection_params,
on_error=mocker.sentinel.on_error,
)
# assert
assert result == (transport, protocol, channel)
aioamqp.connect.assert_called_once_with(
host=connection_params.host,
port=connection_params.port,
login=connection_params.username,
password=connection_params.password,
virtualhost=connection_params.virtual_host,
on_error=mocker.sentinel.on_error,
loop=asyncio.get_event_loop(),
)
protocol.channel.assert_called_once_with()
async def test_queue_to_iterator(mocker):
# arrange
queue = mocker.Mock(spec=asyncio.Queue)
queue.get.side_effect = [future(1), future(2), future(exception=Cancelled())]
# act
result = []
with pytest.raises(Cancelled):
async for item in queue_to_iterator(queue):
result.append(item)
# assert
assert result == [1, 2]
async def test__connect(self, mocker, event_loop):
# arrange
consumer = Consumer(
middleware=Process(lambda _: future(None)),
queue=mocker.sentinel.queue,
connection_params=[mocker.sentinel.connection_params],
prefetch_count=mocker.sentinel.prefetch_count,
default_reconnect_timeout=3.0,
max_reconnect_timeout=5.0,
)
channel = mocker.Mock(spec=Channel)
channel.basic_qos.return_value = future()
connect_and_open_channel = mocker.patch(
'aioamqp_consumer_best.consumer.connect_and_open_channel',
return_value=future((
mocker.sentinel.transport,
mocker.sentinel.protocol,
channel
)),
)
declare_queue = mocker.patch('aioamqp_consumer_best.consumer.declare_queue', return_value=future())
connection_closed_future = asyncio.Future()
# act
await consumer._connect(
connection_closed_future=connection_closed_future,
loop=event_loop,
)
# assert
on_error_arg = Arg()
connect_and_open_channel.assert_called_once_with(
connection_params=mocker.sentinel.connection_params,
on_error=on_error_arg,
loop=event_loop,
)
exc = Exception()
await on_error_arg.value(exc)
await on_error_arg.value(exc)
assert connection_closed_future.done()
assert connection_closed_future.exception() is exc
channel.basic_qos.assert_called_once_with(prefetch_count=mocker.sentinel.prefetch_count)
declare_queue.assert_called_once_with(channel=channel, queue=mocker.sentinel.queue)
async def test_connect__integer_heartbeat_interval__passed_to_aioamqp(mocker):
# arrange
transport = mocker.Mock(spec=asyncio.BaseTransport)
protocol = mocker.Mock(spec=aioamqp.AmqpProtocol)
protocol.wait_closed.return_value = future()
mocker.patch.object(aioamqp,
'connect',
return_value=future((transport, protocol)))
# act
async with connect(ConnectionParams(), heartbeat_interval=30):
pass
# assert
assert aioamqp.connect.call_args[1]['heartbeat'] == 30
async def test_load_json(mocker, event_loop):
# arrange
message1 = Message(
body='{"hello": "world"}',
channel=mocker.sentinel.channel,
envelope=mocker.sentinel.envelope,
properties=mocker.sentinel.properties,
)
message2 = Message(
body='invalid json',
channel=mocker.sentinel.channel,
envelope=mocker.sentinel.envelope,
properties=mocker.sentinel.properties,
)
mocker.patch.object(message2, 'reject', return_value=future())
inp = make_iterator([message1, message2])
# act
out = load_json.func(inp, event_loop)
# assert
result = await collect_iterator(out)
message_arg = Arg()
assert result == [message_arg]
assert message_arg.value.body == {'hello': 'world'}
message2.reject.assert_called_once_with(requeue=False)
async def test__disconnect_ok(mocker):
# arrange
consumer = get_consumer(callback=simple_callback)
connection = mocker.patch.object(consumer, '_connection', autospec=True)
connection.close.return_value = future()
channel = mocker.patch.object(consumer, '_channel', autospec=True)
channel.close.return_value = future()
# act
await consumer._disconnect()
# assert
consumer._connection.close.assert_called_once_with()
consumer._channel.close.assert_called_once_with()
async def test__process_queue__ack_message_when_json_is_invalid(
mocker, event_loop, capsys):
# arrange
consumer = get_consumer(callback=simple_callback,
prefetch_count=1,
reject_invalid_json=False)
message = asynqp.IncomingMessage(
body="Invalid JSON",
sender=None,
delivery_tag=None,
exchange_name=None,
routing_key=None,
)
mocker.patch.object(message, 'ack', autospec=True)
mocker.patch.object(consumer,
'_get_messages_iterator',
return_value=future(AsyncIter([message])))
logger_exception = mocker.patch(
'asynqp_consumer.consumer.logger.exception')
# act
await consumer._process_queue(loop=event_loop)
# assert
assert consumer._messages == []
message.ack.assert_called_once_with()
logger_exception.assert_called_once_with(
'Failed to parse message body: %s', b'Invalid JSON')
async def test_connect_and_open_channel(mocker, event_loop):
# arrange
connection = mocker.Mock(spec=asynqp.Connection)
channel = mocker.Mock(spec=asynqp.Channel)
asynqp_connect_and_open_channel = mocker.patch(
'asynqp_consumer.connect.asynqp.connect_and_open_channel',
autospec=True,
return_value=future((connection, channel)))
# act
result = await connect_and_open_channel(ConnectionParams(), event_loop)
# assert
assert result == (connection, channel)
asynqp_connect_and_open_channel.assert_called_once_with(
host='localhost',
port=5672,
username='******',
password='******',
virtual_host='/',
loop=event_loop,
)
class TestProcess:
@pytest.mark.parametrize('ack_result', [
future(None),
future(exception=MessageAlreadyResolved()),
])
async def test_run__success__ack_message(self, mocker, ack_result):
# arrange
message = Message(
body='message',
channel=mocker.sentinel.channel,
envelope=mocker.sentinel.envelope,
properties=mocker.sentinel.properties,
)
mocker.patch.object(message, 'ack', return_value=ack_result)
inp = make_iterator([message])
middleware = Process(lambda _: future(None))
# act
out = middleware(inp)
# assert
assert await collect_iterator(out) == [None]
message.ack.assert_called_once_with()
@pytest.mark.parametrize('reject_result', [
future(None),
future(exception=MessageAlreadyResolved()),
])
async def test_run__callback_raised_exception__reject_message(
self, mocker, reject_result):
# arrange
message = Message(
body='message',
channel=mocker.sentinel.channel,
envelope=mocker.sentinel.envelope,
properties=mocker.sentinel.properties,
)
mocker.patch.object(message, 'reject', return_value=reject_result)
inp = make_iterator([message])
middleware = Process(lambda _: future(exception=Exception()))
# act
out = middleware(inp)
# assert
assert await collect_iterator(out) == [None]
message.reject.assert_called_once_with()
async def test_declare_queue(mocker):
# arrange
asynqp_queue = mocker.Mock(spec=asynqp.Queue)
asynqp_queue.bind.return_value = future()
asynqp_exchange = mocker.Mock(spec=asynqp.Exchange)
channel = mocker.Mock(spec=asynqp.Channel)
channel.declare_queue.return_value = future(asynqp_queue)
channel.declare_exchange.return_value = future(asynqp_exchange)
# act
result = await declare_queue(
channel,
Queue(name='test-queue',
bindings=[
QueueBinding(
exchange=Exchange('test-exchange'),
routing_key='test-routing-key',
)
]))
# assert
assert result is asynqp_queue
channel.declare_queue.assert_called_once_with(
name='test-queue',
durable=True,
exclusive=False,
auto_delete=False,
arguments=None,
)
channel.declare_exchange.assert_called_once_with(
name='test-exchange',
type='topic',
durable=True,
auto_delete=False,
arguments=None,
)
asynqp_queue.bind.assert_called_once_with(
exchange=asynqp_exchange,
routing_key='test-routing-key',
arguments=None,
)
async def test__process_queue__when_prefetch_count_is_not_0(mocker, event_loop):
# arrange
consumer = get_consumer(callback=simple_callback, prefetch_count=1)
mocker.patch.object(consumer, '_get_messages_iterator', return_value=future(AsyncIter([
mock.Mock(spec=asynqp.IncomingMessage),
mock.Mock(spec=asynqp.IncomingMessage),
])))
mocker.patch.object(consumer, '_process_bulk', side_effect=iter([future(), future()]))
# act
await consumer._process_queue(loop=event_loop)
# assert
assert len(consumer._messages) == 2
assert isinstance(consumer._messages[0], Message)
assert isinstance(consumer._messages[1], Message)
consumer._process_bulk.mock_calls == [mocker.call(), mocker.call()]
async def test__connect__ok(mocker, event_loop):
# arrange
connection = mocker.Mock(spec=asynqp.Connection)
channel = mocker.Mock(spec=asynqp.Channel)
channel.set_qos.return_value = future()
connect_and_open_channel = mocker.patch(
'asynqp_consumer.consumer.connect_and_open_channel', autospec=True)
connect_and_open_channel.return_value = future((connection, channel))
asynqp_queue = mocker.Mock(spec=asynqp.Queue)
declare_queue = mocker.patch('asynqp_consumer.consumer.declare_queue',
autospec=True)
declare_queue.return_value = future(asynqp_queue)
consumer = get_consumer(callback=simple_callback)
# act
await consumer._connect(loop=event_loop)
# assert
connect_and_open_channel.assert_called_once_with(
ConnectionParams(
host='test_host',
port=1234,
username='******',
password='******',
virtual_host='test_virtual_host',
), event_loop)
declare_queue.assert_called_once_with(
channel,
Queue(name='test_queue',
bindings=[
QueueBinding(exchange=Exchange('test_exchange'),
routing_key='test_routing_key')
]))
channel.set_qos.assert_called_once_with(prefetch_count=0)
assert consumer._connection is connection
assert consumer._channel is channel
assert consumer._queue is asynqp_queue
async def test_start__two_attempts(mocker, event_loop):
# arrange
consumer = get_consumer(callback=simple_callback)
mocker.patch.object(consumer, '_connect', side_effect=iter([OSError, future()]))
mocker.patch.object(consumer, '_disconnect', return_value=future())
mocker.patch.object(consumer, '_process_queue', return_value=future())
mocker.patch.object(consumer, '_check_bulk', return_value=future())
consumer._connection = mocker.Mock(spec=asynqp.Connection)
consumer._connection.closed = asyncio.Future(loop=event_loop)
mocker.patch('asynqp_consumer.consumer.gather', autospec=True, return_value=future())
Future = mocker.patch('asynqp_consumer.consumer.asyncio.Future', autospec=True)
Future.return_value.done.side_effect = iter([False, False, True])
Future.return_value._loop = event_loop
consumer._connection.closed.set_exception(ConsumerCloseException)
sleep = mocker.patch('asynqp_consumer.consumer.asyncio.sleep', return_value=future())
# act
await consumer.start(loop=event_loop)
# assert
assert consumer._connect.mock_calls == [
mocker.call(loop=event_loop),
mocker.call(loop=event_loop),
]
consumer._disconnect.assert_called_once_with()
consumer._process_queue.assert_called_once_with(loop=event_loop)
consumer._check_bulk.assert_called_once_with(loop=event_loop)
sleep.assert_called_once_with(3, loop=event_loop)
async def test_declare_queue(mocker):
# arrange
channel = mocker.Mock(spec=Channel)
channel.queue_declare.return_value = future()
channel.exchange_declare.return_value = future()
channel.queue_bind.return_value = future()
exchange = Exchange('exchange')
queue_binding = QueueBinding(
exchange=exchange,
routing_key='routing-key',
)
queue = Queue(name='queue-name', bindings=[queue_binding])
# act
await declare_queue(channel=channel, queue=queue)
# assert
channel.queue_declare.assert_called_once_with(
queue_name=queue.name,
durable=queue.durable,
exclusive=queue.exclusive,
auto_delete=queue.auto_delete,
arguments=queue.arguments,
)
channel.exchange_declare.assert_called_once_with(
exchange_name=exchange.name,
type_name=exchange.type.value,
auto_delete=exchange.auto_delete,
durable=exchange.durable,
arguments=exchange.arguments,
)
channel.queue_bind.assert_called_once_with(
queue_name=queue.name,
exchange_name=exchange.name,
routing_key=queue_binding.routing_key,
arguments=exchange.arguments,
)
async def test_run__callback_cancelled__should_cancel(self, mocker):
# arrange
message = Message(
body='message',
channel=mocker.sentinel.channel,
envelope=mocker.sentinel.envelope,
properties=mocker.sentinel.properties,
)
inp = make_iterator([message])
middleware = Process(
lambda _: future(exception=asyncio.CancelledError()))
# act
with pytest.raises(asyncio.CancelledError):
await collect_iterator(middleware(inp))
async def test__consume_with_arguments(mocker, event_loop):
# arrange
queue = mocker.Mock(spec=asynqp.Queue)
queue.consume.return_value = future()
consume_arguments = {'x-priority': 100}
consumer = get_consumer(callback=simple_callback, prefetch_count=0, consume_arguments=consume_arguments)
mocker.patch.object(consumer, '_queue', new=queue)
asyncio_queue = mocker.patch('asynqp_consumer.consumer.asyncio.Queue').return_value
# act
await consumer._get_messages_iterator(loop=event_loop)
# assert
queue.consume.assert_called_once_with(callback=asyncio_queue.put_nowait, arguments=consume_arguments)
async def test__process_queue__when_message_is_invalid_json(mocker, event_loop):
# arrange
consumer = get_consumer(callback=simple_callback, prefetch_count=1)
message = mock.Mock(spec=asynqp.IncomingMessage)
message.json.side_effect = json.JSONDecodeError('message', '', 0)
message.body = 'Error json'
mocker.patch.object(consumer, '_get_messages_iterator', return_value=future(AsyncIter([message])))
# act
await consumer._process_queue(loop=event_loop)
# assert
assert consumer._messages == []
async def test__disconnect(self, mocker):
# arrange
consumer = Consumer(
middleware=Process(lambda _: future(None)),
queue=mocker.sentinel.queue,
connection_params=[mocker.sentinel.connection_params],
prefetch_count=mocker.sentinel.prefetch_count,
default_reconnect_timeout=3.0,
max_reconnect_timeout=5.0,
)
channel = mocker.Mock(spec=Channel)
channel.is_open = True
channel.close.return_value = future()
consumer._channel = channel
protocol = mocker.Mock(spec=Channel)
protocol.state = aioamqp.protocol.OPEN
protocol.close.return_value = future()
consumer._protocol = protocol
transport = mocker.Mock(spec=Channel)
consumer._transport = transport
# act
await consumer._disconnect()
# assert
assert consumer._channel is None
channel.close.assert_called_once_with()
assert consumer._protocol is None
protocol.close.assert_called_once_with()
assert consumer._transport is None
transport.close.assert_called_once_with()
async def test_start__expected_right_reconnect_timeouts_and_calls(self, mocker, event_loop):
# arrange
consumer = Consumer(
middleware=Process(lambda _: future(None)),
queue=mocker.sentinel.queue,
connection_params=[mocker.sentinel.connection_params],
prefetch_count=1,
default_reconnect_timeout=3.0,
max_reconnect_timeout=5.0,
)
mocker.patch.object(consumer, '_connect', side_effect=[
future(),
future(exception=aioamqp.AioamqpException()),
future(),
])
mocker.patch.object(consumer, '_disconnect', return_value=future())
mocker.patch.object(consumer, '_process_queue', side_effect=[
future(exception=aioamqp.AioamqpException()),
future(exception=_ConsumerCloseException()),
])
mocker.patch.object(asyncio, 'sleep', return_value=future())
# act
await consumer.start(event_loop)
# assert
connection_closed_future_arg = Arg()
assert consumer._connect.call_count == 3
consumer._connect.assert_called_with(
connection_closed_future=connection_closed_future_arg,
loop=event_loop,
)
assert isinstance(connection_closed_future_arg.value, asyncio.Future)
consumer._disconnect.assert_called_once_with()
assert consumer._process_queue.call_count == 2
consumer._process_queue.assert_called_with(loop=event_loop)
assert asyncio.sleep.call_args_list == [
mocker.call(3.0, loop=event_loop),
mocker.call(5.0, loop=event_loop),
]
async def test_logout__delete_user__deletes_user_and_logs_out(mocker):
# arrange
m_get_user = mocker.patch.object(auth,
'get_user_id_from_cookie',
return_value=1)
m_delete_user = mocker.patch.object(auth,
'_delete_user',
return_value=future())
m_response = mocker.Mock()
m_response.delete_cookie = mocker.Mock()
# act
await auth.logout(cookie='test', response=m_response, delete_user=True)
# assert
m_get_user.assert_called_once_with('test')
m_delete_user.assert_called_once_with(1)
m_response.delete_cookie.assert_called_once_with('fastapi_auth')
async def test_run__success__ack_message(self, mocker, ack_result):
# arrange
message = Message(
body='message',
channel=mocker.sentinel.channel,
envelope=mocker.sentinel.envelope,
properties=mocker.sentinel.properties,
)
mocker.patch.object(message, 'ack', return_value=ack_result)
inp = make_iterator([message])
middleware = Process(lambda _: future(None))
# act
out = middleware(inp)
# assert
assert await collect_iterator(out) == [None]
message.ack.assert_called_once_with()
async def test_run__callback_raised_exception__reject_message(
self, mocker, reject_result):
# arrange
message = Message(
body='message',
channel=mocker.sentinel.channel,
envelope=mocker.sentinel.envelope,
properties=mocker.sentinel.properties,
)
mocker.patch.object(message, 'reject', return_value=reject_result)
inp = make_iterator([message])
middleware = Process(lambda _: future(exception=Exception()))
# act
out = middleware(inp)
# assert
assert await collect_iterator(out) == [None]
message.reject.assert_called_once_with()
async def test_reject(self, mocker):
# arrange
channel = mocker.Mock(spec=Channel)
channel.basic_reject.return_value = future()
message = self._make_message(mocker, channel)
# act
await message.reject()
# assert
with pytest.raises(MessageAlreadyResolved):
await message.ack()
with pytest.raises(MessageAlreadyResolved):
await message.reject()
channel.basic_reject.assert_called_once_with(
delivery_tag=mocker.sentinel.delivery_tag, requeue=True)
assert not channel.basic_client_ack.called
async def test_close(self, mocker):
# arrange
consumer = Consumer(
middleware=Process(lambda _: future(None)),
queue=mocker.sentinel.queue,
connection_params=[mocker.sentinel.connection_params],
prefetch_count=1,
default_reconnect_timeout=3.0,
max_reconnect_timeout=5.0,
)
consumer._closed_future = asyncio.Future()
consumer._closed_ok = asyncio.Event()
consumer._closed_ok.set()
mocker.spy(consumer._closed_ok, 'wait')
# act
await consumer.close()
# assert
assert consumer._closed_future.done()
assert isinstance(consumer._closed_future.exception(), _ConsumerCloseException)
consumer._closed_ok.wait.assert_called_once_with()
async def test__process_queue(self, mocker, event_loop):
# arrange
consumer = Consumer(
middleware=Process(lambda _: future(None)),
queue=Queue('queue_name'),
connection_params=[mocker.sentinel.connection_params],
prefetch_count=mocker.sentinel.prefetch_count,
default_reconnect_timeout=3.0,
max_reconnect_timeout=5.0,
tag='tag',
)
consumer._middleware = mocker.Mock(spec=Middleware)
consumer._middleware.return_value = make_iterator([])
consumer._channel = mocker.Mock(spec=Channel)
consumer._channel.basic_consume.return_value = future()
queue_to_iterator_mock = mocker.patch(
'aioamqp_consumer_best.consumer.queue_to_iterator',
autospec=True,
return_value=mocker.sentinel.inp,
)
Message = mocker.patch('aioamqp_consumer_best.consumer.Message')
# act
await consumer._process_queue(loop=event_loop)
# assert
callback_arg = Arg()
consumer._channel.basic_consume.assert_called_once_with(
callback=callback_arg,
queue_name='queue_name',
consumer_tag='tag',
)
arg = Arg()
queue_to_iterator_mock.assert_called_once_with(arg)
input_queue = arg.value
assert isinstance(input_queue, asyncio.Queue)
consumer._middleware.assert_called_once_with(
inp=mocker.sentinel.inp,
loop=event_loop,
)
await callback_arg.value(
channel=mocker.sentinel.channel,
body=mocker.sentinel.body,
envelope=mocker.sentinel.envelope,
properties=mocker.sentinel.properties,
)
message = input_queue.get_nowait()
assert message is Message.return_value
Message.assert_called_once_with(
channel=mocker.sentinel.channel,
body=mocker.sentinel.body,
envelope=mocker.sentinel.envelope,
properties=mocker.sentinel.properties,
)
