[fix] handle int in headers fixes #209

mosquito · Apr 3, 2019 · 2cf7ff1 · 2cf7ff1
1 parent bea0065
commit 2cf7ff1
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/aio_pika/message.py b/aio_pika/message.py
@@ -161,7 +161,8 @@ def __getitem__(self, k):
         return self._cache[k]
 
     def __setitem__(self, key, value):
-        self._headers[key] = str(value).encode()
+        self._headers[key] = format_headers(value)
+        self._cache.pop(key, None)
 
     def __len__(self) -> int:
         return len(self._headers)
@@ -183,6 +184,7 @@ def header_converter(value: Any) -> bytes:
 @header_converter.register(datetime)
 @header_converter.register(NoneType)
 @header_converter.register(list)
+@header_converter.register(int)
 def _(v: bytes):
     return v
 

diff --git a/tests/test_amqp.py b/tests/test_amqp.py
@@ -1499,3 +1499,20 @@ def test_message_info(self):
         )
 
         self.assertDictEqual(info, msg.info())
+
+    def test_headers_content(self):
+        data = (
+            [42, 42, 42],
+            ['foo', b'foo', 'foo'],
+            [b'\00', b'\00', '\00'],
+        )
+
+        for src, raw, value in data:
+            msg = Message(b'', headers={'value': src})
+            self.assertEqual(
+                msg.headers_raw['value'], raw, "%r != %r" % (src, raw)
+            )
+
+            self.assertEqual(
+                msg.headers['value'], value, "%r != %r" % (src, value)
+            )