tests: fix in-memory mocks

airtai · Sep 20, 2024 · 1069686 · 1069686
1 parent 694757f
commit 1069686
Showing 1 changed file with 14 additions and 6 deletions.
diff --git a/tests/cli/rabbit/test_app.py b/tests/cli/rabbit/test_app.py
@@ -122,7 +122,9 @@ async def call2():
 
     test_app = FastStream(broker=broker, after_startup=[call1, call2])
 
-    with patch.object(test_app.broker, "start", async_mock.broker_start):
+    with patch.object(test_app.broker, "start", async_mock.broker_start), patch.object(
+        test_app.broker, "connect", async_mock.broker_connect
+    ):
         await test_app.start()
 
     mock.after_startup1.assert_called_once()
@@ -144,7 +146,9 @@ async def call_after():
         async_mock.before.assert_awaited_once()
         async_mock.broker_start.assert_called_once()
 
-    with patch.object(app.broker, "start", async_mock.broker_start):
+    with patch.object(app.broker, "start", async_mock.broker_start), patch.object(
+        app.broker, "connect", async_mock.broker_connect
+    ):
         await app.start()
 
     async_mock.broker_start.assert_called_once()
@@ -164,7 +168,9 @@ async def call2():
 
     test_app = FastStream(broker=broker, after_shutdown=[call1, call2])
 
-    with patch.object(test_app.broker, "start", async_mock.broker_start):
+    with patch.object(test_app.broker, "start", async_mock.broker_start), patch.object(
+        test_app.broker, "connect", async_mock.broker_connect
+    ):
         await test_app.stop()
 
     mock.after_shutdown1.assert_called_once()
@@ -341,8 +347,8 @@ async def lifespan(env: str):
     app = FastStream(app.broker, lifespan=lifespan)
 
     with patch.object(app.broker, "start", async_mock.broker_run), patch.object(
-        app.broker, "close", async_mock.broker_stopped
-    ):
+        app.broker, "connect", async_mock.broker_connect
+    ), patch.object(app.broker, "close", async_mock.broker_stopped):
         async with TestApp(app, {"env": "test"}):
             pass
 
@@ -363,7 +369,9 @@ async def lifespan(env: str):
 
     with patch.object(app.broker, "start", async_mock.broker_run), patch.object(
         app.broker, "close", async_mock.broker_stopped
-    ), TestApp(app, {"env": "test"}):
+    ), patch.object(app.broker, "connect", async_mock.broker_connect), TestApp(
+        app, {"env": "test"}
+    ):
         pass
 
     async_mock.on.assert_awaited_once_with("test")