Draft for the async test client (encode#440)

Author: gbgduh

requirements.txt

@@ -15,7 +15,9 @@ databases[sqlite]
 isort
 mypy
 pytest
+pytest-asyncio
 pytest-cov
+requests_async
 
 # Documentation
 mkdocs

starlette/testclient.py

@@ -10,6 +10,7 @@
 from urllib.parse import unquote, urljoin, urlsplit
 
 import requests
+import requests_async
 
 from starlette.types import Message, Receive, Scope, Send
 from starlette.websockets import WebSocketDisconnect
@@ -472,3 +473,38 @@ async def wait_shutdown(self) -> None:
             self.task.result()
         assert message["type"] == "lifespan.shutdown.complete"
         await self.task
+
+
+class AsyncTestClient(requests_async.ASGISession):
+    async def lifespan(self) -> None:
+        scope = {"type": "lifespan"}
+        try:
+            await self.app(scope, self.receive_queue.get, self.send_queue.put)
+        finally:
+            await self.send_queue.put(None)
+
+    async def wait_startup(self) -> None:
+        await self.receive_queue.put({"type": "lifespan.startup"})
+        message = await self.send_queue.get()
+        if message is None:
+            self.task.result()
+        assert message["type"] == "lifespan.startup.complete"
+
+    async def wait_shutdown(self) -> None:
+        await self.receive_queue.put({"type": "lifespan.shutdown"})
+        message = await self.send_queue.get()
+        if message is None:
+            self.task.result()
+        assert message["type"] == "lifespan.shutdown.complete"
+        await self.task
+
+    async def __aenter__(self) -> requests_async.ASGISession:
+        loop = asyncio.get_event_loop()
+        self.send_queue = asyncio.Queue()  # type: asyncio.Queue
+        self.receive_queue = asyncio.Queue()  # type: asyncio.Queue
+        self.task = loop.create_task(self.lifespan())
+        await self.wait_startup()
+        return self
+
+    async def __aexit__(self, *args: typing.Any) -> None:
+        await self.wait_shutdown()

tests/middleware/test_lifespan.py

@@ -3,7 +3,7 @@
 from starlette.applications import Starlette
 from starlette.responses import PlainTextResponse
 from starlette.routing import Lifespan, Route, Router
-from starlette.testclient import TestClient
+from starlette.testclient import AsyncTestClient, TestClient
 
 
 def test_routed_lifespan():
@@ -106,3 +106,52 @@ async def run_shutdown():
         assert not shutdown_complete
     assert startup_complete
     assert shutdown_complete
+
+
+@pytest.mark.asyncio
+async def test_raise_on_startup_async_test_client():
+    def run_startup():
+        raise RuntimeError()
+
+    app = Router(routes=[Lifespan(on_startup=run_startup)])
+
+    with pytest.raises(RuntimeError):
+        async with AsyncTestClient(app):
+            pass  # pragma: nocover
+
+
+@pytest.mark.asyncio
+async def test_raise_on_shutdown_async_test_client():
+    def run_shutdown():
+        raise RuntimeError()
+
+    app = Router(routes=[Lifespan(on_shutdown=run_shutdown)])
+
+    with pytest.raises(RuntimeError):
+        async with AsyncTestClient(app):
+            pass  # pragma: nocover
+
+
+@pytest.mark.asyncio
+async def test_app_async_test_client_lifespan():
+    startup_complete = False
+    shutdown_complete = False
+    app = Starlette()
+
+    @app.on_event("startup")
+    async def run_startup():
+        nonlocal startup_complete
+        startup_complete = True
+
+    @app.on_event("shutdown")
+    async def run_shutdown():
+        nonlocal shutdown_complete
+        shutdown_complete = True
+
+    assert not startup_complete
+    assert not shutdown_complete
+    async with AsyncTestClient(app):
+        assert startup_complete
+        assert not shutdown_complete
+    assert startup_complete
+    assert shutdown_complete

tests/test_asynctestclient.py

@@ -0,0 +1,78 @@
+import pytest
+
+from starlette.applications import Starlette
+from starlette.responses import JSONResponse
+from starlette.testclient import AsyncTestClient
+
+mock_service = Starlette()
+
+
+@mock_service.route("/")
+def mock_service_endpoint(request):
+    return JSONResponse({"mock": "example"})
+
+
+app = Starlette()
+
+
+@pytest.mark.asyncio
+@app.route("/")
+async def homepage(request):
+    client = AsyncTestClient(mock_service)
+    response = await client.get("/")
+    return JSONResponse(response.json())
+
+
+startup_error_app = Starlette()
+
+
+@startup_error_app.on_event("startup")
+def startup():
+    raise RuntimeError()
+
+
+@pytest.mark.asyncio
+async def test_use_asynctestclient_in_endpoint():
+    """
+    We should be able to use the test client within applications.
+
+    This is useful if we need to mock out other services,
+    during tests or in development.
+    """
+    client = AsyncTestClient(app)
+    response = await client.get("/")
+    assert response.json() == {"mock": "example"}
+
+
+@pytest.mark.asyncio
+async def test_asynctestclient_as_contextmanager():
+    async with AsyncTestClient(app):
+        pass
+
+
+@pytest.mark.asyncio
+async def test_error_on_startup():
+    with pytest.raises(RuntimeError):
+        async with AsyncTestClient(startup_error_app):
+            pass  # pragma: no cover
+
+
+# TODO test_asynctestclient_asgi2
+# `requests_async` is ASGI 3 only as of now
+
+
+@pytest.mark.asyncio
+async def test_asynctestclient_asgi3():
+    async def app(scope, receive, send):
+        await send(
+            {
+                "type": "http.response.start",
+                "status": 200,
+                "headers": [[b"content-type", b"text/plain"]],
+            }
+        )
+        await send({"type": "http.response.body", "body": b"Hello, world!"})
+
+    client = AsyncTestClient(app)
+    response = await client.get("/")
+    assert response.text == "Hello, world!"

tests/test_database.py

@@ -5,7 +5,7 @@
 from starlette.applications import Starlette
 from starlette.datastructures import CommaSeparatedStrings
 from starlette.responses import JSONResponse
-from starlette.testclient import TestClient
+from starlette.testclient import AsyncTestClient, TestClient
 
 DATABASE_URL = "sqlite:///test.db"
 
@@ -165,3 +165,21 @@ def test_database_isolated_during_test_cases():
         response = client.get("/notes")
         assert response.status_code == 200
         assert response.json() == [{"text": "just one note", "completed": True}]
+
+
+@pytest.mark.asyncio
+async def test_database_with_async_test_client():
+    """
+    Using AsyncTestClient
+    """
+    async with AsyncTestClient(app) as client:
+        # Post a row to the DB
+        response = await client.post(
+            "/notes", json={"text": "just one note", "completed": True}
+        )
+        assert response.status_code == 200
+
+        # Call the DB explicitly
+        query = notes.select()
+        results = await database.fetch_all(query)
+        assert results == [(1, "just one note", True)]