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)