Fixed now awaited coroutine warnings

Made it returning tasks from startup and awaiting it on test
daltonmatos · Feb 17, 2023 · 4e79ec6 · 4e79ec6
1 parent 06377aa
commit 4e79ec6
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 7 deletions.
diff --git a/asyncworker/signals/handlers/rabbitmq.py b/asyncworker/signals/handlers/rabbitmq.py
@@ -1,4 +1,5 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, List
+from asyncio import Task
 
 from asyncworker.connections import AMQPConnection
 from asyncworker.consumer import Consumer
@@ -10,7 +11,9 @@
 
 
 class RabbitMQ(SignalHandler):
-    async def startup(self, app: "App"):
+    async def startup(self, app: "App") -> List[Task]:
+        tasks = []
+
         app[RouteTypes.AMQP_RABBITMQ]["consumers"] = []
         for route_info in app.routes_registry.amqp_routes:
             conn: AMQPConnection = app.get_connection_for_route(route_info)
@@ -24,4 +27,8 @@ async def startup(self, app: "App"):
             )
             app[RouteTypes.AMQP_RABBITMQ]["consumers"].append(consumer)
             conn.register(consumer.queue)
-            app.loop.create_task(consumer.start())
+            task = app.loop.create_task(consumer.start())
+
+            tasks.append(task)
+
+        return tasks
diff --git a/tests/rabbitmq/test_rabbitmq_consumer.py b/tests/rabbitmq/test_rabbitmq_consumer.py
@@ -643,8 +643,12 @@ async def test_start_always_calls_sleep(self):
         ) as sleep_mock, patch.object(
             consumer, "clock_task", side_effect=[True, True]
         ):
-            queue_mock = AsyncMock(
-                consume=AsyncMock(), connect=AsyncMock()
+            queue_mock = Mock(
+                consume=AsyncMock(),
+                connection=Mock(
+                    connect=AsyncMock(),
+                    has_channel_ready=Mock(),
+                ),
             )
             consumer.queue = queue_mock
 
@@ -682,7 +686,7 @@ async def test_start_create_clock_flusher(self):
         with patch.object(
             consumer, "keep_runnig", side_effect=[True, True, True, False]
         ):
-            queue_mock = AsyncMock(
+            queue_mock = Mock(
                 consume=AsyncMock(),
                 connection=Mock(
                     connect=AsyncMock(side_effect=[AioamqpException, True]),

diff --git a/tests/signals/handlers/test_rabbitmq.py b/tests/signals/handlers/test_rabbitmq.py
@@ -1,3 +1,4 @@
+import asyncio
 from unittest import IsolatedAsyncioTestCase
 from unittest.mock import AsyncMock, Mock, patch, call, ANY
 
@@ -52,7 +53,7 @@ async def test_startup_initializes_and_starts_one_consumer_per_route(
         # )
 
         app[RouteTypes.AMQP_RABBITMQ] = {"connections": [connection]}
-        await self.signal_handler.startup(app)
+        tasks = await self.signal_handler.startup(app)
 
         Consumer.assert_has_calls(
             [
@@ -79,6 +80,8 @@ async def test_startup_initializes_and_starts_one_consumer_per_route(
             [Consumer.return_value, Consumer.return_value],
         )
 
+        await asyncio.gather(*tasks)
+
     @patch(
         "asyncworker.signals.handlers.rabbitmq.AMQPConnection.register"
     )