def test_context_process(self):
client = yield from connect(AMQP_URL, loop=self.loop)
queue_name = self.get_random_name("test_connection")
routing_key = self.get_random_name()
channel = yield from client.channel()
exchange = yield from channel.declare_exchange('direct', auto_delete=True)
queue = yield from channel.declare_queue(queue_name, auto_delete=True)
yield from queue.bind(exchange, routing_key)
body = bytes(shortuuid.uuid(), 'utf-8')
yield from exchange.publish(
Message(
body, content_type='text/plain',
headers={'foo': 'bar'}
),
routing_key
)
incoming_message = yield from queue.get(timeout=5)
with self.assertRaises(AssertionError):
with incoming_message.process(requeue=True):
raise AssertionError
self.assertEqual(incoming_message.locked, True)
incoming_message = yield from queue.get(timeout=5)
with incoming_message.process():
pass
self.assertEqual(incoming_message.body, body)
yield from exchange.publish(
Message(
body, content_type='text/plain',
headers={'foo': 'bar'}
),
routing_key
)
incoming_message = yield from queue.get(timeout=5)
with self.assertRaises(AssertionError):
with incoming_message.process(requeue=True, reject_on_redelivered=True):
raise AssertionError
incoming_message = yield from queue.get(timeout=5)
with self.assertRaises(AssertionError):
with incoming_message.process(requeue=True, reject_on_redelivered=True):
raise AssertionError
self.assertEqual(incoming_message.locked, True)
yield from queue.unbind(exchange, routing_key)
yield from queue.delete()
yield from wait((client.close(), client.closing), loop=self.loop)
def test_simple(kwargs, expected):
loop = asyncio.get_event_loop()
# noinspection PyTypeChecker
conn = loop.run_until_complete(
connect(connection_class=FakeConnection, **kwargs), )
assert conn.url == URL(expected)
def test_consuming_not_coroutine(self):
client = yield from connect(AMQP_URL, loop=self.loop)
queue_name = self.get_random_name("tc2")
routing_key = self.get_random_name()
channel = yield from client.channel()
exchange = yield from channel.declare_exchange('direct',
auto_delete=True)
queue = yield from channel.declare_queue(queue_name, auto_delete=True)
yield from queue.bind(exchange, routing_key)
body = bytes(shortuuid.uuid(), 'utf-8')
f = asyncio.Future(loop=self.loop)
def handle(message):
message.ack()
self.assertEqual(message.body, body)
self.assertEqual(message.routing_key, routing_key)
f.set_result(True)
queue.consume(handle)
yield from exchange.publish(
Message(body, content_type='text/plain', headers={'foo': 'bar'}),
routing_key)
if not f.done():
yield from f
yield from queue.unbind(exchange, routing_key)
yield from exchange.delete()
yield from wait((client.close(), client.closing), loop=self.loop)
def test_purge_queue(self):
client = yield from connect(AMQP_URL, loop=self.loop)
queue_name = self.get_random_name("test_connection4")
routing_key = self.get_random_name()
channel = yield from client.channel()
exchange = yield from channel.declare_exchange('direct',
auto_delete=True)
queue = yield from channel.declare_queue(queue_name, auto_delete=True)
yield from queue.bind(exchange, routing_key)
try:
body = bytes(shortuuid.uuid(), 'utf-8')
yield from exchange.publish(
Message(body,
content_type='text/plain',
headers={'foo': 'bar'}), routing_key)
yield from queue.purge()
with self.assertRaises(TimeoutError):
yield from queue.get(timeout=1)
except:
yield from queue.unbind(exchange, routing_key)
yield from queue.delete()
yield from wait((client.close(), client.closing), loop=self.loop)
def test_exchange_delete(self):
client = yield from connect(AMQP_URL, loop=self.loop)
channel = yield from client.channel()
exchange = yield from channel.declare_exchange("test",
auto_delete=True)
yield from exchange.delete()
yield from client.close()
def test_internal_exchange(self):
client = yield from connect(AMQP_URL, loop=self.loop)
routing_key = self.get_random_name()
exchange_name = self.get_random_name("internal", "exchange")
channel = yield from client.channel()
exchange = yield from channel.declare_exchange(exchange_name,
auto_delete=True,
internal=True)
queue = yield from channel.declare_queue(auto_delete=True)
yield from queue.bind(exchange, routing_key)
body = bytes(shortuuid.uuid(), 'utf-8')
with self.assertRaises(ValueError):
f = exchange.publish(
Message(body,
content_type='text/plain',
headers={'foo': 'bar'}), routing_key)
yield from f
yield from queue.unbind(exchange, routing_key)
yield from queue.delete()
yield from wait((client.close(), client.closing), loop=self.loop)
def test_simple_publish_and_receive(self):
client = yield from connect(AMQP_URL, loop=self.loop)
queue_name = self.get_random_name("test_connection")
routing_key = self.get_random_name()
channel = yield from client.channel()
exchange = yield from channel.declare_exchange('direct',
auto_delete=True)
queue = yield from channel.declare_queue(queue_name, auto_delete=True)
yield from queue.bind(exchange, routing_key)
body = bytes(shortuuid.uuid(), 'utf-8')
yield from exchange.publish(
Message(body, content_type='text/plain', headers={'foo': 'bar'}),
routing_key)
incoming_message = yield from queue.get(timeout=5)
incoming_message.ack()
self.assertEqual(incoming_message.body, body)
yield from queue.unbind(exchange, routing_key)
yield from queue.delete()
yield from wait((client.close(), client.closing), loop=self.loop)
def test_dlx(self):
client = yield from connect(AMQP_URL, loop=self.loop)
direct_queue_name = self.get_random_name("test_dlx", "direct")
dlx_queue_name = self.get_random_name("test_dlx", "dlx")
routing_key = self.get_random_name()
channel = yield from client.channel()
direct_exchange = yield from channel.declare_exchange('direct', auto_delete=True)
direct_queue = yield from channel.declare_queue(
direct_queue_name,
auto_delete=True,
arguments={
'x-message-ttl': 300,
'x-dead-letter-exchange': 'dlx',
'x-dead-letter-routing-key': routing_key
}
)
direct_queue.bind(direct_exchange, routing_key)
@asyncio.coroutine
def dlx_handle(message):
message.ack()
self.assertEqual(message.body, body)
self.assertEqual(message.routing_key, routing_key)
f.set_result(True)
dlx_exchange = yield from channel.declare_exchange('dlx', ExchangeType.DIRECT, auto_delete=True)
dlx_queue = yield from channel.declare_queue(dlx_queue_name, auto_delete=True)
dlx_queue.consume(dlx_handle)
yield from dlx_queue.bind(dlx_exchange, routing_key)
body = bytes(shortuuid.uuid(), 'utf-8')
try:
f = asyncio.Future(loop=self.loop)
yield from direct_exchange.publish(
Message(
body,
content_type='text/plain',
headers={
'x-message-ttl': 100,
'x-dead-letter-exchange': 'dlx',
}
),
routing_key
)
if not f.done():
yield from f
finally:
yield from dlx_queue.unbind(dlx_exchange, routing_key)
yield from direct_queue.unbind(direct_exchange, routing_key)
yield from direct_queue.delete()
yield from direct_exchange.delete()
yield from dlx_exchange.delete()
yield from client.close()
def create_connection(self,
cleanup=True) -> Generator[Any, None, Connection]:
client = yield from connect(AMQP_URL, loop=self.loop)
if cleanup:
self.addCleanup(client.close)
return client
def test_wrong_credentials(self):
amqp_url = AMQP_URL.with_user(uuid.uuid4().hex).with_password(uuid.uuid4().hex)
with self.assertRaises(ProbableAuthenticationError):
yield from connect(
amqp_url,
loop=self.loop
)
def test_add_close_callback(self):
client = yield from connect(AMQP_URL, loop=self.loop)
f = create_future(loop=self.loop)
client.add_close_callback(f.set_result)
yield from client.close()
self.assertTrue(f.done())
def test_basic_return(self):
client = yield from connect(AMQP_URL, loop=self.loop)
channel = yield from client.channel() # type: aio_pika.Channel
f = asyncio.Future(loop=self.loop)
channel.add_on_return_callback(f.set_result)
body = bytes(shortuuid.uuid(), 'utf-8')
yield from channel.default_exchange.publish(
Message(
body,
content_type='text/plain',
headers={'foo': 'bar'}
),
self.get_random_name("test_basic_return")
)
returned = yield from f
self.assertEqual(returned.body, body)
# handler with exception
f = asyncio.Future(loop=self.loop)
yield from channel.close()
channel = yield from client.channel() # type: aio_pika.Channel
def bad_handler(message):
try:
raise ValueError
finally:
f.set_result(message)
channel.add_on_return_callback(bad_handler)
body = bytes(shortuuid.uuid(), 'utf-8')
yield from channel.default_exchange.publish(
Message(
body,
content_type='text/plain',
headers={'foo': 'bar'}
),
self.get_random_name("test_basic_return")
)
returned = yield from f
self.assertEqual(returned.body, body)
yield from wait((client.close(), client.closing), loop=self.loop)
def test_delete_queue_and_exchange(self):
client = yield from connect(AMQP_URL, loop=self.loop)
queue_name = self.get_random_name("test_connection")
exchange = self.get_random_name()
channel = yield from client.channel()
yield from channel.declare_exchange(exchange, auto_delete=True)
yield from channel.declare_queue(queue_name, auto_delete=True)
yield from channel.queue_delete(queue_name)
yield from channel.exchange_delete(exchange)
yield from wait((client.close(), client.closing), loop=self.loop)
def test_expiration(self):
client = yield from connect(AMQP_URL, loop=self.loop)
channel = yield from client.channel() # type: aio_pika.Channel
dlx_queue = yield from channel.declare_queue(
self.get_random_name("test_dlx")
) # type: aio_pika.Queue
dlx_exchange = yield from channel.declare_exchange(
self.get_random_name("dlx"),
) # type: aio_pika.Exchange
yield from dlx_queue.bind(dlx_exchange, routing_key=dlx_queue.name)
queue = yield from channel.declare_queue(
self.get_random_name("test_expiration"),
arguments={
"x-message-ttl": 10000,
"x-dead-letter-exchange": dlx_exchange.name,
"x-dead-letter-routing-key": dlx_queue.name,
}
) # type: aio_pika.Queue
body = bytes(shortuuid.uuid(), 'utf-8')
yield from channel.default_exchange.publish(
Message(
body,
content_type='text/plain',
headers={'foo': 'bar'},
expiration=0.5
),
queue.name
)
f = asyncio.Future(loop=self.loop)
dlx_queue.consume(f.set_result, no_ack=True)
message = yield from f
self.assertEqual(message.body, body)
self.assertEqual(message.headers['x-death'][0]['original-expiration'], '500')
yield from wait((client.close(), client.closing), loop=self.loop)
def test_temporary_queue(self):
client = yield from connect(AMQP_URL, loop=self.loop)
channel = yield from client.channel()
queue = yield from channel.declare_queue(auto_delete=True)
self.assertNotEqual(queue.name, '')
body = os.urandom(32)
yield from channel.default_exchange.publish(Message(body=body), routing_key=queue.name)
message = yield from queue.get()
self.assertEqual(message.body, body)
yield from channel.queue_delete(queue.name)
yield from wait((client.close(), client.closing), loop=self.loop)
def test_context_process_redelivery(self):
client = yield from connect(AMQP_URL, loop=self.loop)
queue_name = self.get_random_name("test_connection")
routing_key = self.get_random_name()
channel = yield from client.channel()
exchange = yield from channel.declare_exchange('direct',
auto_delete=True)
queue = yield from channel.declare_queue(queue_name, auto_delete=True)
yield from queue.bind(exchange, routing_key)
body = bytes(shortuuid.uuid(), 'utf-8')
yield from exchange.publish(
Message(body, content_type='text/plain', headers={'foo': 'bar'}),
routing_key)
incoming_message = yield from queue.get(timeout=5)
with self.assertRaises(AssertionError):
with incoming_message.process(requeue=True,
reject_on_redelivered=True):
raise AssertionError
incoming_message = yield from queue.get(timeout=5)
with mock.patch('aio_pika.message.log') as message_logger:
with self.assertRaises(Exception):
with incoming_message.process(requeue=True,
reject_on_redelivered=True):
raise Exception
self.assertTrue(message_logger.info.called)
self.assertEqual(message_logger.info.mock_calls[0][1][1].body,
incoming_message.body)
self.assertEqual(incoming_message.body, body)
yield from queue.unbind(exchange, routing_key)
yield from queue.delete()
yield from wait((client.close(), client.closing), loop=self.loop)
def test_connection_close(self):
client = yield from connect(AMQP_URL, loop=self.loop) # type: Connection
routing_key = self.get_random_name()
channel = yield from client.channel() # type: Channel
exchange = yield from channel.declare_exchange('direct', auto_delete=True)
try:
with self.assertRaises(aio_pika.exceptions.ChannelClosed):
msg = Message(bytes(shortuuid.uuid(), 'utf-8'))
msg.delivery_mode = 8
yield from exchange.publish(msg, routing_key)
channel = yield from client.channel()
exchange = yield from channel.declare_exchange('direct', auto_delete=True)
finally:
yield from exchange.delete()
yield from wait((client.close(), client.closing), loop=self.loop)
def test_no_ack_redelivery(self):
client = yield from connect(AMQP_URL, loop=self.loop)
queue_name = self.get_random_name("test_connection")
routing_key = self.get_random_name()
channel = yield from client.channel()
exchange = yield from channel.declare_exchange('direct',
auto_delete=True)
queue = yield from channel.declare_queue(queue_name, auto_delete=False)
yield from queue.bind(exchange, routing_key)
# publish 2 messages
for _ in range(2):
body = bytes(shortuuid.uuid(), 'utf-8')
msg = Message(body)
yield from exchange.publish(msg, routing_key)
# ack 1 message out of 2
first_message = yield from queue.get(timeout=5)
last_message = yield from queue.get(timeout=5)
last_message.ack()
# close channel, not acked message should be redelivered
yield from channel.close()
channel = yield from client.channel()
exchange = yield from channel.declare_exchange('direct',
auto_delete=True)
queue = yield from channel.declare_queue(queue_name, auto_delete=False)
# receive not acked message
message = yield from queue.get(timeout=5)
self.assertEqual(message.body, first_message.body)
message.ack()
yield from queue.unbind(exchange, routing_key)
yield from queue.delete()
yield from wait((client.close(), client.closing), loop=self.loop)
async def send_message(text, source):
routing = source
try:
connection = await (aio_pika.connect(
"amqp://{user}:{user}@{host}/".format(
user=os.getenv('MQ_USER', 'guest'),
host=os.getenv('MQ_HOST', 'mq'))))
except Exception as exp:
raise RabbitConnectError(str(exp))
try:
channel = await connection.channel()
exchange = await channel.declare_exchange('messages',
aio_pika.ExchangeType.DIRECT)
await exchange.publish(
aio_pika.Message(body=text.encode(),
delivery_mode=aio_pika.DeliveryMode.PERSISTENT),
routing_key=routing)
await connection.close()
except Exception as exp:
raise RabbitWorkingError(f'{str(exp)}')
def test_ack_multiple(self):
client = yield from connect(AMQP_URL, loop=self.loop)
queue_name = self.get_random_name("test_connection")
routing_key = self.get_random_name()
channel = yield from client.channel()
exchange = yield from channel.declare_exchange('direct',
auto_delete=True)
queue = yield from channel.declare_queue(queue_name, auto_delete=False)
yield from queue.bind(exchange, routing_key)
# publish 2 messages
for _ in range(2):
body = bytes(shortuuid.uuid(), 'utf-8')
msg = Message(body)
yield from exchange.publish(msg, routing_key)
# ack only last mesage with multiple flag, first message should be acked too
yield from queue.get(timeout=5)
last_message = yield from queue.get(timeout=5)
last_message.ack(multiple=True)
# close channel, no messages should be redelivered
yield from channel.close()
channel = yield from client.channel()
exchange = yield from channel.declare_exchange('direct',
auto_delete=True)
queue = yield from channel.declare_queue(queue_name, auto_delete=False)
with self.assertRaises(TimeoutError):
yield from queue.get(timeout=1)
yield from queue.unbind(exchange, routing_key)
yield from queue.delete()
yield from wait((client.close(), client.closing), loop=self.loop)
def test_channel_close(self):
client = yield from connect(AMQP_URL, loop=self.loop)
self.get_random_name("test_connection")
self.get_random_name()
self.__closed = False
def on_close(ch):
log.info("Close called")
self.__closed = True
channel = yield from client.channel()
channel.add_close_callback(on_close)
yield from channel.close()
yield from asyncio.sleep(1, loop=self.loop)
self.assertTrue(self.__closed)
with self.assertRaises(RuntimeError):
yield from channel.initialize()
yield from wait((client.close(), client.closing), loop=self.loop)
def create_connection(self):
client = yield from connect(AMQP_URL, loop=self.loop)
self.addCleanup(client.close)
return client
def test_set_qos(self):
client = yield from connect(AMQP_URL, loop=self.loop)
channel = yield from client.channel()
yield from channel.set_qos(prefetch_count=1, all_channels=True)
yield from wait((client.close(), client.closing), loop=self.loop)
def test_connection_refused(self):
with self.assertRaises(ConnectionRefusedError):
yield from connect('amqp://*****:*****@localhost:9999',
loop=self.loop)
storage = RedisStorage2(host=REDIS_HOST,
port=REDIS_PORT,
password=REDIS_PASS if REDIS_PASS else None)
bot = Bot(token=TG_TOKEN, parse_mode="HTML", loop=loop)
dp = Dispatcher(bot, storage=storage)
logging.basicConfig(level=logging.DEBUG)
dp.middleware.setup(LoggingMiddleware())
dp.middleware.setup(GetUserMiddleware())
i18n = setup_middleware(dp)
i18n.reload()
_ = i18n.gettext
loop.run_until_complete(Tortoise.init(config=TORTOISE_ORM))
# await Tortoise.generate_schemas()
rabbit_connection = loop.run_until_complete(connect(
host=RABBIT_HOST,
port=int(RABBIT_PORT),
virtualhost=RABBIT_VIRTUAL_HOST,
login=RABBIT_USER,
password=RABBIT_PASSWORD,
loop=loop
))
rabbit_channel = loop.run_until_complete(rabbit_connection.channel())
rabbit_donation_queue = loop.run_until_complete(rabbit_channel.declare_queue(RABBIT_QUEUE, passive=True))
async def _connect(self):
self.conn = await asyncio.wait_for(aio_pika.connect(**self.config),
timeout=RabbitMQ.timeout)
self.channel = await self.conn.channel()
def test_incoming_message_info(self):
client = yield from connect(AMQP_URL, loop=self.loop)
queue_name = self.get_random_name("test_connection")
routing_key = self.get_random_name()
channel = yield from client.channel()
exchange = yield from channel.declare_exchange('direct',
auto_delete=True)
queue = yield from channel.declare_queue(queue_name, auto_delete=True)
yield from queue.bind(exchange, routing_key)
body = bytes(shortuuid.uuid(), 'utf-8')
info = {
'headers': {
"foo": "bar"
},
'content_type': "application/json",
'content_encoding': "text",
'delivery_mode': DeliveryMode.PERSISTENT.value,
'priority': 0,
'correlation_id': b'1',
'reply_to': 'test',
'expiration': 1.5,
'message_id': shortuuid.uuid(),
'timestamp': int(time.time()),
'type': '0',
'user_id': 'guest',
'app_id': 'test',
'body_size': len(body)
}
msg = Message(body=body,
headers={'foo': 'bar'},
content_type='application/json',
content_encoding='text',
delivery_mode=DeliveryMode.PERSISTENT,
priority=0,
correlation_id=1,
reply_to='test',
expiration=1.5,
message_id=info['message_id'],
timestamp=info['timestamp'],
type='0',
user_id='guest',
app_id='test')
yield from exchange.publish(msg, routing_key)
incoming_message = yield from queue.get(timeout=5)
incoming_message.ack()
info['synchronous'] = incoming_message.synchronous
info['routing_key'] = incoming_message.routing_key
info['redelivered'] = incoming_message.redelivered
info['exchange'] = incoming_message.exchange
info['delivery_tag'] = incoming_message.delivery_tag
info['consumer_tag'] = incoming_message.consumer_tag
info['cluster_id'] = incoming_message.cluster_id
self.assertEqual(incoming_message.body, body)
self.assertDictEqual(incoming_message.info(), info)
yield from queue.unbind(exchange, routing_key)
yield from queue.delete()
yield from wait((client.close(), client.closing), loop=self.loop)
