Updated broker's interface to work with raw bytes.

Signed-off-by: Pavel Kirilin <win10@list.ru>

Author: s3rius

taskiq_nats/broker.py

@@ -58,23 +58,19 @@ async def kick(self, message: BrokerMessage) -> None:
         """
         await self.client.publish(
             self.subject,
-            payload=message.json().encode(),
+            payload=message.message,
             headers=message.labels,
         )
 
-    async def listen(self) -> AsyncGenerator[BrokerMessage, None]:
+    async def listen(self) -> AsyncGenerator[bytes, None]:
         """
         Start listen to new messages.
 
         :yield: incoming messages.
         """
         subscribe = await self.client.subscribe(self.subject, queue=self.queue or "")
         async for message in subscribe.messages:
-            try:
-                yield BrokerMessage.parse_raw(message.data)
-            except ValueError:
-                data = message.data.decode("utf-8")
-                logger.warning(f"Cannot parse message: {data}")
+            yield message.data
 
     async def shutdown(self) -> None:
         """Close connections to NATS."""

tests/test_broker.py

@@ -8,7 +8,7 @@
 from taskiq_nats import NatsBroker
 
 
-async def read_message(broker: NatsBroker) -> BrokerMessage:  # type: ignore
+async def read_message(broker: NatsBroker) -> bytes:  # type: ignore
     """
     Read signle message from the broker's listen method.
 
@@ -31,12 +31,12 @@ async def test_success_broadcast(nats_urls: List[str], nats_subject: str) -> Non
     sent_message = BrokerMessage(
         task_id=uuid.uuid4().hex,
         task_name="meme",
-        message="some",
+        message=b"some",
         labels={},
     )
     asyncio.create_task(broker.kick(sent_message))
     for received_message in await asyncio.gather(*tasks):
-        assert received_message == sent_message
+        assert received_message == sent_message.message
 
 
 @pytest.mark.anyio
@@ -49,8 +49,8 @@ async def test_success_queued(nats_urls: List[str], nats_subject: str) -> None:
     sent_message = BrokerMessage(
         task_id=uuid.uuid4().hex,
         task_name="meme",
-        message="some",
+        message=b"some",
         labels={},
     )
     asyncio.create_task(broker.kick(sent_message))
-    assert await reading_task == sent_message
+    assert await reading_task == sent_message.message