fix: serialize TestClient rpc output to mock the real message (#1452)

airtai · May 16, 2024 · c555f0f · c555f0f
1 parent 987d069
commit c555f0f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/faststream/broker/core/usecase.py b/faststream/broker/core/usecase.py
@@ -333,7 +333,7 @@ async def publish(
         **kwargs: Any,
     ) -> Optional[Any]:
         """Publish message directly."""
-        assert producer, NOT_CONNECTED_YET  # nosec B101)
+        assert producer, NOT_CONNECTED_YET  # nosec B101
 
         publish: "AsyncFunc" = producer.publish
         for m in self._middlewares:

diff --git a/faststream/testing/broker.py b/faststream/testing/broker.py
@@ -16,6 +16,7 @@
 from unittest.mock import AsyncMock, MagicMock
 
 from faststream.broker.core.usecase import BrokerUsecase
+from faststream.broker.message import StreamMessage, decode_message, encode_message
 from faststream.broker.middlewares.logging import CriticalLogMiddleware
 from faststream.broker.wrapper.call import HandlerCallWrapper
 from faststream.testing.app import TestApp
@@ -212,6 +213,9 @@ async def call_handler(
         result = await handler.consume(message)
 
         if rpc:
-            return result
+            message_body, content_type = encode_message(result)
+            msg_to_publish = StreamMessage(raw_message=None, body=message_body, content_type=content_type)
+            consumed_data = decode_message(msg_to_publish)
+            return consumed_data
 
     return None