From f7ac0159f5d92cff4650eab3125dbd2c69f8e9f7 Mon Sep 17 00:00:00 2001
From: ArtyomKozyrev8 <kozirev8@gmail.com>
Date: Tue, 28 Sep 2021 17:00:12 +0300
Subject: [PATCH] change Message body argument to support str type like in pika
 module. Now Message body can be str or bytes or raise TypeError. Previous
 implementation body if isinstance(body, bytes) else bytes(body) gave error
 since there was no encoding argument in bytes(body)

---
 aio_pika/message.py   | 19 +++++++++++++++++--
 tests/test_message.py | 15 +++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/aio_pika/message.py b/aio_pika/message.py
index ff1c8eff..6c10add2 100644
--- a/aio_pika/message.py
+++ b/aio_pika/message.py
@@ -242,7 +242,7 @@ class Message:
 
     def __init__(
         self,
-        body: bytes,
+        body: Union[bytes, str],
         *,
         headers: dict = None,
         content_type: str = None,
@@ -279,7 +279,8 @@ def __init__(
         """
 
         self.__lock = False
-        self.body = body if isinstance(body, bytes) else bytes(body)
+
+        self.body = Message._prepare_msg_body(body)
         self.body_size = len(self.body) if self.body else 0
         self.headers_raw = format_headers(headers)
         self._headers = HeaderProxy(self.headers_raw)
@@ -308,6 +309,20 @@ def headers(self):
     def headers(self, value: dict):
         self.headers_raw = format_headers(value)
 
+    @staticmethod
+    def _prepare_msg_body(body: Union[str, bytes]) -> bytes:
+        """
+        If body type is bytes returns body with no changes
+        If body type is str converts it to bytes ith encoding = 'utf8'
+        If body type is not either str or bytes raises TypeError
+        """
+        if isinstance(body, bytes):
+            return body
+        elif isinstance(body, str):
+            return bytes(body, 'utf8')
+        else:
+            raise TypeError("body argument should be str or bytes, {} was provided instead".format(type(body)))
+
     @staticmethod
     def _as_bytes(value):
         if isinstance(value, bytes):
diff --git a/tests/test_message.py b/tests/test_message.py
index 9a59a1a8..88fb6308 100644
--- a/tests/test_message.py
+++ b/tests/test_message.py
@@ -5,6 +5,21 @@
 import shortuuid
 
 from aio_pika import DeliveryMode, Message
+import pytest
+
+
+def test_init_message_body():
+    # test body is string
+    msg = Message(body="test")
+    assert msg.body == b"test"
+
+    # test body is bytes
+    msg = Message(body=b"test")
+    assert msg.body == b"test"
+
+    # net not str or bytes body
+    with pytest.raises(TypeError):
+        Message(body=1)
 
 
 def test_message_copy():