class RabbitMQ(BaseExtension):
key = "rabbitmq"
def __init__(
self,
vk,
queue_name: str,
rabbitmq_url: str = "amqp://*****:*****@127.0.0.1/",
max_connections: int = 2,
max_channels: int = 15,
):
if aio_pika:
self._vk = vk
self._queue_name = queue_name
self._url = rabbitmq_url
self._conn_pool = Pool(
self.get_connection, max_size=max_connections, loop=vk.loop
)
self._chann_pool = Pool(
self.get_channel, max_size=max_channels, loop=vk.loop
)
else:
raise RuntimeWarning(
"Please install aio_pika (pip install aio_pika) for use this extension"
)
async def get_events(self) -> None:
pass
async def get_connection(self):
return await aio_pika.connect_robust(self._url)
async def get_channel(self) -> "aio_pika.Channel":
async with self._conn_pool.acquire() as connection:
return await connection.channel()
async def run(self, dp):
logger.info("RabbitMQ consumer started!")
async with self._chann_pool.acquire() as channel: # type: aio_pika.Channel
await channel.set_qos(10)
queue = await channel.declare_queue(
self._queue_name, durable=False, auto_delete=False
)
async with queue.iterator() as queue_iter:
async for message in queue_iter:
event = JSON_LIBRARY.loads(message.body.decode())
await dp._process_events([event])
await message.ack()
class TestCaseItemReuse(BaseTestCase):
max_size = 10
call_count = max_size * 5
def setUp(self):
super().setUp()
self.counter = set()
self.pool = Pool(self.create_instance,
max_size=self.max_size,
loop=self.loop)
async def create_instance(self):
obj = object()
self.counter.add(obj)
return obj
async def test_simple(self):
counter = Counter()
async def getter():
nonlocal counter
async with self.pool.acquire() as instance:
await asyncio.sleep(0.05, loop=self.loop)
counter[instance] += 1
await asyncio.gather(*[getter() for _ in range(self.call_count)],
loop=self.loop,
return_exceptions=True)
self.assertEqual(sum(counter.values()), self.call_count)
self.assertEqual(self.counter, set(counter))
self.assertEqual(len(set(counter.values())), 1)
async def send_msg(pool: Pool, rk: str, msg: Message):
async with pool.acquire() as channel:
exchange = rk.split(".")[0]
topic_exchange = await channel.declare_exchange(
exchange, ExchangeType.TOPIC)
await topic_exchange.publish(msg, routing_key=rk)
