From 084f35648b64b70f7a8255faa196b09d432622fe Mon Sep 17 00:00:00 2001 From: Tom Christie Date: Wed, 17 Feb 2021 11:10:21 +0000 Subject: [PATCH] Allow handler to optionally be async when MockTransport is used with AsyncClient (#1449) --- httpx/_transports/mock.py | 11 +++++++++++ tests/client/test_async_client.py | 13 +++++++++++++ 2 files changed, 24 insertions(+) diff --git a/httpx/_transports/mock.py b/httpx/_transports/mock.py index 6f9ebc1e0f..a55a88b7a2 100644 --- a/httpx/_transports/mock.py +++ b/httpx/_transports/mock.py @@ -1,3 +1,4 @@ +import asyncio from typing import Callable, List, Optional, Tuple import httpcore @@ -47,7 +48,17 @@ async def arequest( stream=stream, ) await request.aread() + response = self.handler(request) + + # Allow handler to *optionally* be an `async` function. + # If it is, then the `response` variable need to be awaited to actually + # return the result. + + # https://simonwillison.net/2020/Sep/2/await-me-maybe/ + if asyncio.iscoroutine(response): + response = await response + return ( response.status_code, response.headers.raw, diff --git a/tests/client/test_async_client.py b/tests/client/test_async_client.py index 0cc9d8a4ff..62464865e2 100644 --- a/tests/client/test_async_client.py +++ b/tests/client/test_async_client.py @@ -301,3 +301,16 @@ async def test_mounted_transport(): response = await client.get("custom://www.example.com") assert response.status_code == 200 assert response.json() == {"app": "mounted"} + + +@pytest.mark.usefixtures("async_environment") +async def test_async_mock_transport(): + async def hello_world(request): + return httpx.Response(200, text="Hello, world!") + + transport = httpx.MockTransport(hello_world) + + async with httpx.AsyncClient(transport=transport) as client: + response = await client.get("https://www.example.com") + assert response.status_code == 200 + assert response.text == "Hello, world!"