ISSUE #283

* Improve `RabbitMQBrokerPublisher` and `RabbitMQBrokerSubscriber`.

Author: Sergio García Prado

packages/plugins/minos-broker-rabbitmq/minos/plugins/rabbitmq/publisher.py

@@ -3,6 +3,9 @@
 )
 
 import logging
+from typing import (
+    Optional,
+)
 
 from aio_pika import (
     Message,
@@ -25,23 +28,30 @@
 class RabbitMQBrokerPublisher(BrokerPublisher):
     """RabbitMQ Broker Publisher class."""
 
-    def __init__(self, *args, host: str, port: int, **kwargs):
+    def __init__(self, *args, host: Optional[str] = None, port: Optional[int] = None, **kwargs):
         super().__init__(*args, **kwargs)
+
+        if host is None:
+            host = "localhost"
+        if port is None:
+            port = 5672
+
         self.host = host
         self.port = port
 
+        self.connection = None
+        self.channel = None
+
     async def _setup(self) -> None:
         await super()._setup()
         self.connection = await connect(f"amqp://guest:guest@{self.host}:{self.port}/")
+        self.channel = await self.connection.channel()
 
     async def _destroy(self) -> None:
         await self.connection.close()
 
     async def _send(self, message: BrokerMessage) -> None:
-        async with self.connection:
-            channel = await self.connection.channel()
-            queue = await channel.declare_queue(message.topic)
-            await channel.default_exchange.publish(Message(message.avro_bytes), routing_key=queue.name)
+        await self.channel.default_exchange.publish(Message(message.avro_bytes), routing_key=message.topic)
 
 
 class RabbitMQBrokerPublisherBuilder(BrokerPublisherBuilder[RabbitMQBrokerPublisher], RabbitMQBrokerBuilderMixin):

packages/plugins/minos-broker-rabbitmq/minos/plugins/rabbitmq/subscriber.py

@@ -3,19 +3,28 @@
 )
 
 import logging
+from asyncio import (
+    CancelledError,
+    Queue,
+    TimeoutError,
+    create_task,
+    gather,
+    wait_for,
+)
 from collections.abc import (
     Iterable,
 )
+from contextlib import (
+    suppress,
+)
 from typing import (
+    NoReturn,
     Optional,
 )
 
 from aio_pika import (
     connect,
 )
-from aio_pika.exceptions import (
-    QueueEmpty,
-)
 
 from minos.networks import (
     BrokerMessage,
@@ -36,36 +45,58 @@ class RabbitMQBrokerSubscriber(BrokerSubscriber):
     def __init__(
         self,
         topics: Iterable[str],
-        host: str,
-        port: int,
-        group_id: Optional[str] = None,
-        remove_topics_on_destroy: bool = False,
+        host: Optional[str] = None,
+        port: Optional[int] = None,
         **kwargs,
     ):
         super().__init__(topics, **kwargs)
+
+        if host is None:
+            host = "localhost"
+        if port is None:
+            port = 5672
+
         self.host = host
         self.port = port
-        self.group_id = group_id
 
-        self.remove_topics_on_destroy = remove_topics_on_destroy
+        self.connection = None
+
+        self._run_task = None
+        self._queue: Queue[bytes] = Queue(maxsize=1)
 
     async def _setup(self) -> None:
         await super()._setup()
         self.connection = await connect(f"amqp://guest:guest@{self.host}:{self.port}/")
+        await self._start_task()
 
     async def _destroy(self) -> None:
+        await self._stop_task()
         await self.connection.close()
 
+    async def _start_task(self):
+        if self._run_task is None:
+            self._run_task = create_task(self._run())
+
+    async def _stop_task(self):
+        if self._run_task is not None:
+            self._run_task.cancel()
+            with suppress(TimeoutError, CancelledError):
+                await wait_for(self._run_task, 0.5)
+            self._run_task = None
+
+    async def _run(self) -> NoReturn:
+        await gather(*(self._run_one(topic) for topic in self.topics))
+
+    async def _run_one(self, topic: str) -> None:
+        channel = await self.connection.channel()
+        queue = await channel.declare_queue(topic)
+        async for message in queue.iterator():
+            await self._queue.put(message.body)
+
     async def _receive(self) -> BrokerMessage:
-        for topic in self.topics:
-            async with self.connection:
-                channel = await self.connection.channel()
-                queue = await channel.declare_queue(topic)
-                try:
-                    message = await queue.get()
-                    return BrokerMessage.from_avro_bytes(message.body)
-                except QueueEmpty:
-                    pass
+        bytes_ = await self._queue.get()
+        message = BrokerMessage.from_avro_bytes(bytes_)
+        return message
 
 
 class RabbitMQBrokerSubscriberBuilder(BrokerSubscriberBuilder[RabbitMQBrokerSubscriber], RabbitMQBrokerBuilderMixin):

packages/plugins/minos-broker-rabbitmq/tests/test_rabbitmq/test_integration.py

@@ -1,3 +1,7 @@
+from asyncio import (
+    TimeoutError,
+    wait_for,
+)
 from unittest import (
     IsolatedAsyncioTestCase,
 )
@@ -29,4 +33,5 @@ async def test_one_topic(self):
 
     async def test_empty_topic(self):
         async with RabbitMQBrokerSubscriber.from_config(CONFIG_FILE_PATH, topics={"empty_topic"}) as subscriber:
-            await subscriber.receive()
+            with self.assertRaises(TimeoutError):
+                await wait_for(subscriber.receive(), 0.1)

packages/plugins/minos-broker-rabbitmq/tests/test_rabbitmq/test_publisher.py

@@ -23,6 +23,11 @@ class TestRabbitMQBrokerPublisher(unittest.IsolatedAsyncioTestCase):
     def test_is_subclass(self):
         self.assertTrue(issubclass(RabbitMQBrokerPublisher, BrokerPublisher))
 
+    def test_constructor(self):
+        publisher = RabbitMQBrokerPublisher()
+        self.assertEqual("localhost", publisher.host)
+        self.assertEqual(5672, publisher.port)
+
     def test_from_config(self):
         config = Config(CONFIG_FILE_PATH)
         broker_config = config.get_interface_by_name("broker")["common"]
@@ -33,21 +38,14 @@ def test_from_config(self):
         self.assertEqual(broker_config["host"], publisher.host)
         self.assertEqual(broker_config["port"], publisher.port)
 
-    @patch("minos.plugins.rabbitmq.publisher.connect")
-    async def test_send(self, connect_mock):
-        message = BrokerMessageV1("foo", BrokerMessageV1Payload("bar"))
-
-        async with RabbitMQBrokerPublisher.from_config(CONFIG_FILE_PATH) as publisher:
-            await publisher.send(message)
-
-            self.assertEqual(1, connect_mock.call_count)
-
-    @patch("minos.plugins.rabbitmq.publisher.connect")
-    async def test_destroy(self, destroy_mock):
+    async def test_send(self):
         async with RabbitMQBrokerPublisher.from_config(CONFIG_FILE_PATH) as publisher:
-            await publisher.destroy()
+            with patch("aio_pika.Exchange.publish") as mock:
+                await publisher.send(BrokerMessageV1("foo1", BrokerMessageV1Payload("bar")))
+                await publisher.send(BrokerMessageV1("foo2", BrokerMessageV1Payload("bar")))
+                await publisher.send(BrokerMessageV1("foo1", BrokerMessageV1Payload("bar")))
 
-            self.assertEqual(1, destroy_mock.call_count)
+        self.assertEqual(3, mock.call_count)
 
 
 if __name__ == "__main__":

packages/plugins/minos-broker-rabbitmq/tests/test_rabbitmq/test_subscriber.py

@@ -8,6 +8,7 @@
 
 from minos.common import (
     Config,
+    ModelType,
 )
 from minos.networks import (
     BrokerSubscriber,
@@ -18,31 +19,39 @@
 )
 from tests.utils import (
     CONFIG_FILE_PATH,
+    FakeAsyncIterator,
 )
 
-_ConsumerMessage = namedtuple("_ConsumerMessage", ["value"])
+_Foo = ModelType.build("_Foo", {"bar": str})
+_ConsumerMessage = namedtuple("_ConsumerMessage", ["body"])
 
 
 class TestRabbitMQBrokerSubscriber(unittest.IsolatedAsyncioTestCase):
     def test_is_subclass(self):
         self.assertTrue(issubclass(RabbitMQBrokerSubscriber, BrokerSubscriber))
 
+    def test_constructor(self):
+        subscriber = RabbitMQBrokerSubscriber({"foo", "bar"})
+        self.assertEqual({"foo", "bar"}, subscriber.topics)
+        self.assertEqual("localhost", subscriber.host)
+        self.assertEqual(5672, subscriber.port)
+
     async def test_from_config(self):
         config = Config(CONFIG_FILE_PATH)
         broker_config = config.get_interface_by_name("broker")["common"]
         async with RabbitMQBrokerSubscriber.from_config(config, topics={"foo", "bar"}) as subscriber:
             self.assertEqual(broker_config["host"], subscriber.host)
             self.assertEqual(broker_config["port"], subscriber.port)
-            self.assertEqual(False, subscriber.remove_topics_on_destroy)
             self.assertEqual({"foo", "bar"}, subscriber.topics)
 
-    @patch("minos.plugins.rabbitmq.subscriber.connect")
-    @patch("minos.networks.BrokerMessage.from_avro_bytes")
-    async def test_receive(self, connect_mock, mock_avro):
+    async def test_receive(self):
+        message = _ConsumerMessage(_Foo("foobar").avro_bytes)
+
         async with RabbitMQBrokerSubscriber.from_config(CONFIG_FILE_PATH, topics={"foo", "bar"}) as subscriber:
-            await subscriber.receive()
+            with patch("aio_pika.Queue.iterator", return_value=FakeAsyncIterator([message])):
+                observed = await subscriber.receive()
 
-            self.assertEqual(1, connect_mock.call_count)
+        self.assertEqual(_Foo("foobar"), observed)
 
 
 class TestRabbitMQBrokerSubscriberBuilder(unittest.TestCase):

packages/plugins/minos-broker-rabbitmq/tests/utils.py

@@ -4,3 +4,19 @@
 
 BASE_PATH = Path(__file__).parent
 CONFIG_FILE_PATH = BASE_PATH / "test_config.yml"
+
+
+class FakeAsyncIterator:
+    """For testing purposes."""
+
+    def __init__(self, seq):
+        self.iter = iter(seq)
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            return next(self.iter)
+        except StopIteration:
+            raise StopAsyncIteration