Merge pull request #200 from marcinsulikowski/allow-binary-bodies

Support binary bodies in responses
pnuckowski · Jan 10, 2022 · 9303268 · 9303268
2 parents 51f44f1 + f922b7f
commit 9303268
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 4 deletions.
diff --git a/aioresponses/core.py b/aioresponses/core.py
@@ -35,7 +35,7 @@ class CallbackResult:
 
     def __init__(self, method: str = hdrs.METH_GET,
                  status: int = 200,
-                 body: str = '',
+                 body: Union[str, bytes] = '',
                  content_type: str = 'application/json',
                  payload: Dict = None,
                  headers: Dict = None,
@@ -57,7 +57,7 @@ class RequestMatch(object):
     def __init__(self, url: Union[URL, str, Pattern],
                  method: str = hdrs.METH_GET,
                  status: int = 200,
-                 body: str = '',
+                 body: Union[str, bytes] = '',
                  payload: Dict = None,
                  exception: 'Exception' = None,
                  headers: Dict = None,
@@ -118,7 +118,7 @@ def _build_response(self, url: 'Union[URL, str]',
                         method: str = hdrs.METH_GET,
                         request_headers: Dict = None,
                         status: int = 200,
-                        body: str = '',
+                        body: Union[str, bytes] = '',
                         content_type: str = 'application/json',
                         payload: Dict = None,
                         headers: Dict = None,
@@ -290,7 +290,7 @@ def options(self, url: 'Union[URL, str, Pattern]', **kwargs):
 
     def add(self, url: 'Union[URL, str, Pattern]', method: str = hdrs.METH_GET,
             status: int = 200,
-            body: str = '',
+            body: Union[str, bytes] = '',
             exception: 'Exception' = None,
             content_type: str = 'application/json',
             payload: Dict = None,

diff --git a/tests/test_aioresponses.py b/tests/test_aioresponses.py
@@ -166,6 +166,26 @@ async def test_streaming_up_to(self, m):
         content = await resp.content.read(2)
         self.assertEqual(content, b'st')
 
+    @aioresponses()
+    async def test_binary_body(self, m):
+        body = b'Invalid utf-8: \x95\x00\x85'
+        m.get(self.url, body=body)
+        resp = await self.session.get(self.url)
+        content = await resp.read()
+        self.assertEqual(content, body)
+
+    @aioresponses()
+    async def test_binary_body_via_callback(self, m):
+        body = b'\x00\x01\x02\x80\x81\x82\x83\x84\x85'
+
+        def callback(url, **kwargs):
+            return CallbackResult(body=body)
+
+        m.get(self.url, callback=callback)
+        resp = await self.session.get(self.url)
+        content = await resp.read()
+        self.assertEqual(content, body)
+
     async def test_mocking_as_context_manager(self):
         with aioresponses() as aiomock:
             aiomock.add(self.url, payload={'foo': 'bar'})