tests: fix kafka warnings

airtai · May 9, 2024 · 2ad2d1f · 2ad2d1f
1 parent 1227488
commit 2ad2d1f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/faststream/confluent/testing.py b/faststream/confluent/testing.py
@@ -1,6 +1,6 @@
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, MagicMock
 
 from typing_extensions import override
 
@@ -240,4 +240,7 @@ def build_message(
 
 
 def _fake_connection(*args: Any, **kwargs: Any) -> AsyncMock:
-    return AsyncMock()
+    mock = AsyncMock()
+    mock.getone.return_value = MagicMock()
+    mock.getmany.return_value = [MagicMock()]
+    return mock
diff --git a/faststream/kafka/testing.py b/faststream/kafka/testing.py
@@ -1,6 +1,6 @@
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, MagicMock
 
 from aiokafka import ConsumerRecord
 from typing_extensions import override
@@ -210,4 +210,7 @@ def build_message(
 
 
 def _fake_connection(*args: Any, **kwargs: Any) -> AsyncMock:
-    return AsyncMock()
+    mock = AsyncMock()
+    mock.subscribe = MagicMock
+    mock.assign = MagicMock
+    return mock