Repository: lundberg/respx
Pull request #27 (2019-11-21 15:24:13)

Add support for async decorator and context manager
Makes the `respx.mock` decorator and context manager work correctly in async context.

```py
@respx.mock
async def test_something():
    ...

async with respx.mock():
    ...
```


Patch:
diff --git a/respx/mock.py b/respx/mock.py
index 5f71368..0305717 100644
--- a/respx/mock.py
+++ b/respx/mock.py
@@ -1,3 +1,4 @@
+import inspect
 import io
 import ssl
 import typing
@@ -51,6 +52,9 @@ class HTTPXMock:
         self.start()
         return self
 
+    async def __aenter__(self) -> "HTTPXMock":
+        return self.__enter__()
+
     def __exit__(self, *args: typing.Any) -> None:
         try:
             if self._assert_all_called:
@@ -58,17 +62,31 @@ class HTTPXMock:
         finally:
             self.stop()
 
+    async def __aexit__(self, *args: typing.Any) -> None:
+        self.__exit__(*args)
+
     def mock(self, func=None):
         """
         Starts mocking and stops once wrapped function, or context, is executed.
         """
+        # Context Manager
+        if func is None:
+            return self
+
+        # Async Decorator
+        @wraps(func)
+        async def async_decorator(*args, **kwargs):
+            async with self:
+                return await func(*args, **kwargs)
 
+        # Sync Decorator
         @wraps(func)
-        def wrapper(*args, **kwargs):
+        def sync_decorator(*args, **kwargs):
             with self:
                 return func(*args, **kwargs)
 
-        return wrapper if func else self
+        # Dispatch depening on decorated function
+        return async_decorator if inspect.iscoroutinefunction(func) else sync_decorator
 
     def start(self) -> None:
         """
@@ -323,7 +341,7 @@ class HTTPXMock:
         return await self._mock_socket_stream(response)
 
     async def _mock_socket_stream(self, response: ResponseTemplate) -> BaseSocketStream:
-        content = response.content
+        content = await response.content
         headers = response.headers
 
         # Build raw bytes data
diff --git a/respx/models.py b/respx/models.py
index faaf64c..d8c2a50 100644
--- a/respx/models.py
+++ b/respx/models.py
@@ -1,3 +1,4 @@
+import inspect
 import json as jsonlib
 import re
 import typing
@@ -37,18 +38,21 @@ class ResponseTemplate:
         return self._headers
 
     @property
-    def content(self) -> bytes:
+    async def content(self) -> bytes:
         content = self._content
 
+        if callable(content):
+            if inspect.iscoroutinefunction(content):
+                content = await content(**self.context)
+            else:
+                content = content(**self.context)
+
         if isinstance(content, Exception):
             raise content
 
         if isinstance(content, bytes):
             return content
 
-        if callable(content):
-            content = content(**self.context)
-
         if isinstance(content, (list, dict)):
             content = jsonlib.dumps(content)
             if "Content-Type" not in self._headers:


Test patch:
diff --git a/tests/test_mock.py b/tests/test_mock.py
index 6b31cf8..d510cf1 100644
--- a/tests/test_mock.py
+++ b/tests/test_mock.py
@@ -1,3 +1,4 @@
+import asyncio
 import re
 
 import asynctest
@@ -35,6 +36,16 @@ class HTTPXMockTestCase(asynctest.TestCase):
         response = httpx.get("https://foo/bar/")
         self.assertEqual(response.status_code, 202)
 
+    @respx.mock
+    async def test_mock_decorator_async(self):
+        respx.get("https://foo/bar/", status_code=202)
+        async with httpx.AsyncClient() as client:
+            response = await client.get("https://foo/bar/")
+        self.assertEqual(response.status_code, 202)
+
+    def foo(self):
+        pass
+
     def test_mock_contextmanager(self):
         self.assertEqual(respx.stats.call_count, 0)
 
@@ -47,6 +58,19 @@ class HTTPXMockTestCase(asynctest.TestCase):
 
         self.assertEqual(respx.stats.call_count, 0)
 
+    async def test_mock_contextmanager_async(self):
+        self.assertEqual(respx.stats.call_count, 0)
+
+        async with respx.mock():
+            respx.get("https://foo/bar/", status_code=202)
+            async with httpx.AsyncClient() as client:
+                response = await client.get("https://foo/bar/")
+
+            self.assertEqual(response.status_code, 202)
+            self.assertEqual(respx.stats.call_count, 1)
+
+        self.assertEqual(respx.stats.call_count, 0)
+
     def test_http_methods(self):
         with respx.HTTPXMock() as httpx_mock:
             url = "https://foo/bar/"
@@ -412,61 +436,61 @@ class HTTPXMockTestCase(asynctest.TestCase):
             self.assertTrue(request.called)
             self.assertIsNone(request.pass_through)
 
+    @respx.mock
     async def test_async_stats(self, backend=None):
-        with respx.mock() as m:
-            url = "https://foo/bar/1/"
-            respx.get(re.compile("http://some/url"))
-            respx.delete("http://some/url")
-
-            foobar1 = respx.get(url, status_code=202, alias="get_foobar")
-            foobar2 = respx.delete(url, status_code=200, alias="del_foobar")
-
-            self.assertFalse(foobar1.called)
-            self.assertEqual(foobar1.call_count, len(foobar1.calls))
-            self.assertEqual(foobar1.call_count, 0)
-            self.assertEqual(m.stats.call_count, len(respx.calls))
-            self.assertEqual(m.stats.call_count, 0)
-
-            async with httpx.AsyncClient(backend=backend) as client:
-                get_response = await client.get(url)
-                del_response = await client.delete(url)
+        url = "https://foo/bar/1/"
+        respx.get(re.compile("http://some/url"))
+        respx.delete("http://some/url")
 
-            self.assertTrue(foobar1.called)
-            self.assertTrue(foobar2.called)
-            self.assertEqual(foobar1.call_count, 1)
-            self.assertEqual(foobar2.call_count, 1)
-
-            _request, _response = foobar1.calls[-1]
-            self.assertIsInstance(_request, httpx.AsyncRequest)
-            self.assertIsInstance(_response, httpx.AsyncResponse)
-            self.assertEqual(_request.method, "GET")
-            self.assertEqual(_request.url, url)
-            self.assertEqual(_response.status_code, 202)
-            self.assertEqual(_response.status_code, get_response.status_code)
-            self.assertEqual(_response.content, get_response.content)
-            self.assertEqual(id(_response), id(get_response))
-
-            _request, _response = foobar2.calls[-1]
-            self.assertIsInstance(_request, httpx.AsyncRequest)
-            self.assertIsInstance(_response, httpx.AsyncResponse)
-            self.assertEqual(_request.method, "DELETE")
-            self.assertEqual(_request.url, url)
-            self.assertEqual(_response.status_code, 200)
-            self.assertEqual(_response.status_code, del_response.status_code)
-            self.assertEqual(_response.content, del_response.content)
-            self.assertEqual(id(_response), id(del_response))
+        foobar1 = respx.get(url, status_code=202, alias="get_foobar")
+        foobar2 = respx.delete(url, status_code=200, alias="del_foobar")
 
-            self.assertEqual(respx.stats.call_count, 2)
-            self.assertEqual(respx.calls[0], foobar1.calls[-1])
-            self.assertEqual(respx.calls[1], foobar2.calls[-1])
-
-            alias = respx.aliases["get_foobar"]
-            self.assertEqual(alias, foobar1)
-            self.assertEqual(alias.alias, foobar1.alias)
+        self.assertFalse(foobar1.called)
+        self.assertEqual(foobar1.call_count, len(foobar1.calls))
+        self.assertEqual(foobar1.call_count, 0)
+        self.assertEqual(respx.stats.call_count, len(respx.calls))
+        self.assertEqual(respx.stats.call_count, 0)
 
-            alias = respx.aliases["del_foobar"]
-            self.assertEqual(alias, foobar2)
-            self.assertEqual(alias.alias, foobar2.alias)
+        async with httpx.AsyncClient(backend=backend) as client:
+            get_response = await client.get(url)
+            del_response = await client.delete(url)
+
+        self.assertTrue(foobar1.called)
+        self.assertTrue(foobar2.called)
+        self.assertEqual(foobar1.call_count, 1)
+        self.assertEqual(foobar2.call_count, 1)
+
+        _request, _response = foobar1.calls[-1]
+        self.assertIsInstance(_request, httpx.AsyncRequest)
+        self.assertIsInstance(_response, httpx.AsyncResponse)
+        self.assertEqual(_request.method, "GET")
+        self.assertEqual(_request.url, url)
+        self.assertEqual(_response.status_code, 202)
+        self.assertEqual(_response.status_code, get_response.status_code)
+        self.assertEqual(_response.content, get_response.content)
+        self.assertEqual(id(_response), id(get_response))
+
+        _request, _response = foobar2.calls[-1]
+        self.assertIsInstance(_request, httpx.AsyncRequest)
+        self.assertIsInstance(_response, httpx.AsyncResponse)
+        self.assertEqual(_request.method, "DELETE")
+        self.assertEqual(_request.url, url)
+        self.assertEqual(_response.status_code, 200)
+        self.assertEqual(_response.status_code, del_response.status_code)
+        self.assertEqual(_response.content, del_response.content)
+        self.assertEqual(id(_response), id(del_response))
+
+        self.assertEqual(respx.stats.call_count, 2)
+        self.assertEqual(respx.calls[0], foobar1.calls[-1])
+        self.assertEqual(respx.calls[1], foobar2.calls[-1])
+
+        alias = respx.aliases["get_foobar"]
+        self.assertEqual(alias, foobar1)
+        self.assertEqual(alias.alias, foobar1.alias)
+
+        alias = respx.aliases["del_foobar"]
+        self.assertEqual(alias, foobar2)
+        self.assertEqual(alias.alias, foobar2.alias)
 
     def test_sync_stats(self, backend=None):
         with respx.mock():
@@ -502,3 +526,23 @@ class HTTPXMockTestCase(asynctest.TestCase):
 
     def test_trio_backend(self):
         trio.run(self.test_async_stats, TrioBackend())
+
+    @respx.mock
+    async def test_parallel_requests(self):
+        async def content(request, page):
+            await asyncio.sleep(0.2 if page == "one" else 0.1)
+            return page
+
+        url_pattern = re.compile(r"https://foo/(?P<page>[a-z]+)/$")
+        respx.get(url_pattern, content=content)
+
+        async with httpx.AsyncClient() as client:
+            responses = await asyncio.gather(
+                client.get("https://foo/one/"), client.get("https://foo/two/")
+            )
+            response_one, response_two = responses
+
+            self.assertEqual(response_one.text, "one")
+            self.assertEqual(response_two.text, "two")
+
+        self.assertEqual(respx.stats.call_count, 2)