Replaced string with bytes in broker Message.

Signed-off-by: Pavel Kirilin <[email protected]>

Author: s3rius

taskiq_aio_pika/broker.py

@@ -191,7 +191,7 @@ async def kick(self, message: BrokerMessage) -> None:
             raise ValueError("Please run startup before kicking.")
         priority = parse_val(int, message.labels.get("priority"))
         rmq_msg = Message(
-            body=message.message.encode(),
+            body=message.message,
             headers={
                 "task_id": message.task_id,
                 "task_name": message.task_name,
@@ -214,7 +214,7 @@ async def kick(self, message: BrokerMessage) -> None:
                 routing_key=self._delay_queue_name,
             )
 
-    async def listen(self) -> AsyncGenerator[BrokerMessage, None]:  # noqa: WPS210
+    async def listen(self) -> AsyncGenerator[bytes, None]:
         """
         Listen to queue.
 
@@ -231,25 +231,7 @@ async def listen(self) -> AsyncGenerator[BrokerMessage, None]:  # noqa: WPS210
         async with queue.iterator() as iterator:
             async for message in iterator:
                 async with message.process():
-                    headers = {}
-                    for header_name, header_value in message.headers.items():
-                        headers[header_name] = str(header_value)
-                    try:
-                        broker_message = BrokerMessage(
-                            task_id=headers.pop("task_id"),
-                            task_name=headers.pop("task_name"),
-                            message=message.body,
-                            labels=headers,
-                        )
-                    except (ValueError, LookupError) as exc:
-                        logger.warning(
-                            "Cannot read broker message %s",
-                            exc,
-                            exc_info=True,
-                        )
-                        continue
-
-                    yield broker_message
+                    yield message.body
 
     async def shutdown(self) -> None:
         """Close all connections on shutdown."""

tests/test_broker.py

@@ -9,7 +9,7 @@
 from taskiq_aio_pika.broker import AioPikaBroker
 
 
-async def get_first_task(broker: AioPikaBroker) -> BrokerMessage:  # type: ignore
+async def get_first_task(broker: AioPikaBroker) -> bytes:  # type: ignore
     """
     Get first message from the queue.
 
@@ -36,7 +36,7 @@ async def test_kick_success(broker: AioPikaBroker) -> None:
     sent = BrokerMessage(
         task_id=task_id,
         task_name=task_name,
-        message="my_msg",
+        message=b"my_msg",
         labels={
             "label1": "val1",
         },
@@ -46,7 +46,7 @@ async def test_kick_success(broker: AioPikaBroker) -> None:
 
     message = await asyncio.wait_for(get_first_task(broker), timeout=0.4)
 
-    assert message == sent
+    assert message == sent.message
 
 
 @pytest.mark.anyio
@@ -111,10 +111,7 @@ async def test_listen(
 
     message = await asyncio.wait_for(get_first_task(broker), timeout=0.4)
 
-    assert message.message == "test_message"
-    assert message.labels == {"label1": "label_val"}
-    assert message.task_id == "test_id"
-    assert message.task_name == "task_name"
+    assert message == b"test_message"
 
 
 @pytest.mark.anyio
@@ -124,7 +121,7 @@ async def test_wrong_format(
     test_channel: Channel,
 ) -> None:
     """
-    Tests that messages with wrong format are ignored.
+    Tests that messages with wrong format are still received.
 
     :param broker: aio-pika broker.
     :param queue_name: test queue name.
@@ -136,8 +133,9 @@ async def test_wrong_format(
         routing_key=queue_name,
     )
 
-    with pytest.raises(asyncio.TimeoutError):
-        await asyncio.wait_for(get_first_task(broker), 0.4)
+    msg_bytes = await asyncio.wait_for(get_first_task(broker), 0.4)
+
+    assert msg_bytes == b"wrong"
 
     with pytest.raises(QueueEmpty):
         await queue.get()
@@ -168,7 +166,7 @@ async def test_delayed_message(
     broker_msg = BrokerMessage(
         task_id="1",
         task_name="name",
-        message="message",
+        message=b"message",
         labels={"delay": "2"},
     )
     await broker.kick(broker_msg)