add boolean option for ASGITransport streaming

Author: jhominal

httpx/_transports/asgi.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 import typing
 
 from .._models import Request, Response
@@ -101,11 +102,9 @@ def __init__(
         self,
         ignore_body: bool,
         asgi_generator: typing.AsyncGenerator[_Message, None],
-        disconnect_request_event: Event,
     ) -> None:
         self._ignore_body = ignore_body
         self._asgi_generator = asgi_generator
-        self._disconnect_request_event = disconnect_request_event
 
     async def __aiter__(self) -> typing.AsyncIterator[bytes]:
         more_body = True
@@ -118,13 +117,10 @@ async def __aiter__(self) -> typing.AsyncIterator[bytes]:
                     more_body = message.get("more_body", False)
                     if chunk and not self._ignore_body:
                         yield chunk
-                    if not more_body:
-                        self._disconnect_request_event.set()
         finally:
             await self.aclose()
 
     async def aclose(self) -> None:
-        self._disconnect_request_event.set()
         await self._asgi_generator.aclose()
 
 
@@ -149,6 +145,9 @@ class ASGITransport(AsyncBaseTransport):
        such as testing the content of a client 500 response.
     * `root_path` - The root path on which the ASGI application should be mounted.
     * `client` - A two-tuple indicating the client IP and port of incoming requests.
+    * `streaming` - Set to `True` to enable streaming of response content. Default to
+      `False`, as activating this feature means that the ASGI `app` will run in a
+      sub-task, which has observable side effects for context variables.
     ```
     """
 
@@ -158,18 +157,20 @@ def __init__(
         raise_app_exceptions: bool = True,
         root_path: str = "",
         client: tuple[str, int] = ("127.0.0.1", 123),
+        *,
+        streaming: bool = False,
     ) -> None:
         self.app = app
         self.raise_app_exceptions = raise_app_exceptions
         self.root_path = root_path
         self.client = client
+        self.streaming = streaming
 
     async def handle_async_request(
         self,
         request: Request,
     ) -> Response:
-        disconnect_request_event = create_event()
-        asgi_generator = self._stream_asgi_messages(request, disconnect_request_event)
+        asgi_generator = self._stream_asgi_messages(request)
 
         async for message in asgi_generator:
             if message["type"] == "http.response.start":
@@ -179,15 +180,13 @@ async def handle_async_request(
                     stream=ASGIResponseStream(
                         ignore_body=request.method == "HEAD",
                         asgi_generator=asgi_generator,
-                        disconnect_request_event=disconnect_request_event,
                     ),
                 )
         else:
-            disconnect_request_event.set()
             return Response(status_code=500, headers=[])
 
     async def _stream_asgi_messages(
-        self, request: Request, disconnect_request_event: Event
+        self, request: Request
     ) -> typing.AsyncGenerator[typing.MutableMapping[str, typing.Any]]:
         assert isinstance(request.stream, AsyncByteStream)
 
@@ -211,9 +210,13 @@ async def _stream_asgi_messages(
         request_body_chunks = request.stream.__aiter__()
         request_complete = False
 
+        # Response.
+        response_complete = create_event()
+
         # ASGI response messages stream
+        stream_size = 0 if self.streaming else float("inf")
         response_message_send_stream, response_message_recv_stream = (
-            create_memory_object_stream(0)
+            create_memory_object_stream(stream_size)
         )
 
         # ASGI app exception
@@ -225,7 +228,7 @@ async def receive() -> _Message:
             nonlocal request_complete
 
             if request_complete:
-                await disconnect_request_event.wait()
+                await response_complete.wait()
                 return {"type": "http.disconnect"}
 
             try:
@@ -235,17 +238,29 @@ async def receive() -> _Message:
                 return {"type": "http.request", "body": b"", "more_body": False}
             return {"type": "http.request", "body": body, "more_body": True}
 
+        async def send(message: _Message) -> None:
+            await response_message_send_stream.send(message)
+            if message["type"] == "http.response.body" and not message.get(
+                "more_body", False
+            ):
+                response_complete.set()
+
         async def run_app() -> None:
             nonlocal app_exception
             try:
-                await self.app(scope, receive, response_message_send_stream.send)
+                await self.app(scope, receive, send)
             except Exception as ex:
                 app_exception = ex
             finally:
                 await response_message_send_stream.aclose()
 
-        async with create_task_group() as task_group:
-            task_group.start_soon(run_app)
+        async with contextlib.AsyncExitStack() as exit_stack:
+            exit_stack.callback(response_complete.set)
+            if self.streaming:
+                task_group = await exit_stack.enter_async_context(create_task_group())
+                task_group.start_soon(run_app)
+            else:
+                await run_app()
 
             async with response_message_recv_stream:
                 try:

tests/test_asgi.py

@@ -1,11 +1,20 @@
 import json
 
 import anyio
+import contextvars
 import pytest
 
 import httpx
 
 
+test_asgi_contextvar = contextvars.ContextVar("test_asgi_contextvar")
+
+
+@pytest.fixture(params=[False, True], ids=["no_streaming", "with_streaming"])
+def streaming(request):
+    return request.param
+
+
 def run_in_task_group(app):
     """A decorator that runs an ASGI callable in a task group"""
 
@@ -91,8 +100,8 @@ async def raise_exc_after_response(scope, receive, send):
 
 
 @pytest.mark.anyio
-async def test_asgi_transport():
-    async with httpx.ASGITransport(app=hello_world) as transport:
+async def test_asgi_transport(streaming: bool):
+    async with httpx.ASGITransport(app=hello_world, streaming=streaming) as transport:
         request = httpx.Request("GET", "http://www.example.com/")
         response = await transport.handle_async_request(request)
         await response.aread()
@@ -101,8 +110,8 @@ async def test_asgi_transport():
 
 
 @pytest.mark.anyio
-async def test_asgi_transport_no_body():
-    async with httpx.ASGITransport(app=echo_body) as transport:
+async def test_asgi_transport_no_body(streaming: bool):
+    async with httpx.ASGITransport(app=echo_body, streaming=streaming) as transport:
         request = httpx.Request("GET", "http://www.example.com/")
         response = await transport.handle_async_request(request)
         await response.aread()
@@ -111,8 +120,8 @@ async def test_asgi_transport_no_body():
 
 
 @pytest.mark.anyio
-async def test_asgi():
-    transport = httpx.ASGITransport(app=hello_world)
+async def test_asgi(streaming: bool):
+    transport = httpx.ASGITransport(app=hello_world, streaming=streaming)
     async with httpx.AsyncClient(transport=transport) as client:
         response = await client.get("http://www.example.org/")
 
@@ -121,8 +130,8 @@ async def test_asgi():
 
 
 @pytest.mark.anyio
-async def test_asgi_urlencoded_path():
-    transport = httpx.ASGITransport(app=echo_path)
+async def test_asgi_urlencoded_path(streaming: bool):
+    transport = httpx.ASGITransport(app=echo_path, streaming=streaming)
     async with httpx.AsyncClient(transport=transport) as client:
         url = httpx.URL("http://www.example.org/").copy_with(path="/user@example.org")
         response = await client.get(url)
@@ -132,8 +141,8 @@ async def test_asgi_urlencoded_path():
 
 
 @pytest.mark.anyio
-async def test_asgi_raw_path():
-    transport = httpx.ASGITransport(app=echo_raw_path)
+async def test_asgi_raw_path(streaming: bool):
+    transport = httpx.ASGITransport(app=echo_raw_path, streaming=streaming)
     async with httpx.AsyncClient(transport=transport) as client:
         url = httpx.URL("http://www.example.org/").copy_with(path="/user@example.org")
         response = await client.get(url)
@@ -143,11 +152,11 @@ async def test_asgi_raw_path():
 
 
 @pytest.mark.anyio
-async def test_asgi_raw_path_should_not_include_querystring_portion():
+async def test_asgi_raw_path_should_not_include_querystring_portion(streaming: bool):
     """
     See https://github.com/encode/httpx/issues/2810
     """
-    transport = httpx.ASGITransport(app=echo_raw_path)
+    transport = httpx.ASGITransport(app=echo_raw_path, streaming=streaming)
     async with httpx.AsyncClient(transport=transport) as client:
         url = httpx.URL("http://www.example.org/path?query")
         response = await client.get(url)
@@ -157,8 +166,8 @@ async def test_asgi_raw_path_should_not_include_querystring_portion():
 
 
 @pytest.mark.anyio
-async def test_asgi_upload():
-    transport = httpx.ASGITransport(app=echo_body)
+async def test_asgi_upload(streaming: bool):
+    transport = httpx.ASGITransport(app=echo_body, streaming=streaming)
     async with httpx.AsyncClient(transport=transport) as client:
         response = await client.post("http://www.example.org/", content=b"example")
 
@@ -167,8 +176,8 @@ async def test_asgi_upload():
 
 
 @pytest.mark.anyio
-async def test_asgi_headers():
-    transport = httpx.ASGITransport(app=echo_headers)
+async def test_asgi_headers(streaming: bool):
+    transport = httpx.ASGITransport(app=echo_headers, streaming=streaming)
     async with httpx.AsyncClient(transport=transport) as client:
         response = await client.get("http://www.example.org/")
 
@@ -185,31 +194,33 @@ async def test_asgi_headers():
 
 
 @pytest.mark.anyio
-async def test_asgi_exc():
-    transport = httpx.ASGITransport(app=raise_exc)
+async def test_asgi_exc(streaming: bool):
+    transport = httpx.ASGITransport(app=raise_exc, streaming=streaming)
     async with httpx.AsyncClient(transport=transport) as client:
         with pytest.raises(RuntimeError):
             await client.get("http://www.example.org/")
 
 
 @pytest.mark.anyio
-async def test_asgi_exc_after_response_start():
-    transport = httpx.ASGITransport(app=raise_exc_after_response_start)
+async def test_asgi_exc_after_response_start(streaming: bool):
+    transport = httpx.ASGITransport(
+        app=raise_exc_after_response_start, streaming=streaming
+    )
     async with httpx.AsyncClient(transport=transport) as client:
         with pytest.raises(RuntimeError):
             await client.get("http://www.example.org/")
 
 
 @pytest.mark.anyio
-async def test_asgi_exc_after_response():
-    transport = httpx.ASGITransport(app=raise_exc_after_response)
+async def test_asgi_exc_after_response(streaming: bool):
+    transport = httpx.ASGITransport(app=raise_exc_after_response, streaming=streaming)
     async with httpx.AsyncClient(transport=transport) as client:
         with pytest.raises(RuntimeError):
             await client.get("http://www.example.org/")
 
 
 @pytest.mark.anyio
-async def test_asgi_disconnect_after_response_complete():
+async def test_asgi_disconnect_after_response_complete(streaming: bool):
     disconnect = False
 
     async def read_body(scope, receive, send):
@@ -235,7 +246,7 @@ async def read_body(scope, receive, send):
         message = await receive()
         disconnect = message.get("type") == "http.disconnect"
 
-    transport = httpx.ASGITransport(app=read_body)
+    transport = httpx.ASGITransport(app=read_body, streaming=streaming)
     async with httpx.AsyncClient(transport=transport) as client:
         response = await client.post("http://www.example.org/", content=b"example")
 
@@ -244,18 +255,33 @@ async def read_body(scope, receive, send):
 
 
 @pytest.mark.anyio
-async def test_asgi_exc_no_raise():
-    transport = httpx.ASGITransport(app=raise_exc, raise_app_exceptions=False)
+async def test_asgi_exc_no_raise(streaming: bool):
+    transport = httpx.ASGITransport(
+        app=raise_exc, raise_app_exceptions=False, streaming=streaming
+    )
     async with httpx.AsyncClient(transport=transport) as client:
         response = await client.get("http://www.example.org/")
 
         assert response.status_code == 500
 
 
 @pytest.mark.anyio
-async def test_asgi_exc_no_raise_after_response_start():
+async def test_asgi_exc_no_raise_after_response_start(streaming: bool):
+    transport = httpx.ASGITransport(
+        app=raise_exc_after_response_start,
+        raise_app_exceptions=False,
+        streaming=streaming,
+    )
+    async with httpx.AsyncClient(transport=transport) as client:
+        response = await client.get("http://www.example.org/")
+
+        assert response.status_code == 200
+
+
+@pytest.mark.anyio
+async def test_asgi_exc_no_raise_after_response(streaming: bool):
     transport = httpx.ASGITransport(
-        app=raise_exc_after_response_start, raise_app_exceptions=False
+        app=raise_exc_after_response, raise_app_exceptions=False, streaming=streaming
     )
     async with httpx.AsyncClient(transport=transport) as client:
         response = await client.get("http://www.example.org/")
@@ -264,15 +290,32 @@ async def test_asgi_exc_no_raise_after_response_start():
 
 
 @pytest.mark.anyio
-async def test_asgi_exc_no_raise_after_response():
+async def test_asgi_app_runs_in_same_context_as_caller():
+    async def set_contextvar_in_app(scope, receive, send):
+        test_asgi_contextvar.set("value_from_app")
+
+        status = 200
+        output = b"Hello, World!"
+        headers = [
+            (b"content-type", "text/plain"),
+            (b"content-length", str(len(output))),
+        ]
+
+        await send(
+            {"type": "http.response.start", "status": status, "headers": headers}
+        )
+        await send({"type": "http.response.body", "body": output})
+
     transport = httpx.ASGITransport(
-        app=raise_exc_after_response, raise_app_exceptions=False
+        app=set_contextvar_in_app, raise_app_exceptions=False, streaming=False
     )
     async with httpx.AsyncClient(transport=transport) as client:
         response = await client.get("http://www.example.org/")
 
         assert response.status_code == 200
 
+        assert test_asgi_contextvar.get(None) == "value_from_app"
+
 
 @pytest.mark.parametrize(
     "send_in_sub_task",
@@ -296,7 +339,7 @@ async def send_response_body_after_event(scope, receive, send):
             send_response_body_after_event
         )
 
-    transport = httpx.ASGITransport(app=send_response_body_after_event)
+    transport = httpx.ASGITransport(app=send_response_body_after_event, streaming=True)
     async with httpx.AsyncClient(transport=transport) as client:
         with anyio.fail_after(0.1):
             async with client.stream("GET", "http://www.example.org/") as response:
@@ -335,7 +378,7 @@ async def send_response_body_after_event(scope, receive, send):
             send_response_body_after_event
         )
 
-    transport = httpx.ASGITransport(app=send_response_body_after_event)
+    transport = httpx.ASGITransport(app=send_response_body_after_event, streaming=True)
     async with httpx.AsyncClient(transport=transport) as client:
         with anyio.fail_after(0.1):
             async with client.stream("GET", "http://www.example.org/") as response: