]> git.ipfire.org Git - thirdparty/starlette.git/commitdiff
Test StreamingResponse behavior on receiving http.disconnect (#1727)
authorJean Hominal <jhominal@gmail.com>
Sat, 2 Jul 2022 18:11:46 +0000 (20:11 +0200)
committerGitHub <noreply@github.com>
Sat, 2 Jul 2022 18:11:46 +0000 (20:11 +0200)
tests/test_responses.py

index a272559ebbccaa89215dd6b8790cafa46165c6a7..2030a73824d0fcc1ddf189542272265ccaaef787 100644 (file)
@@ -391,3 +391,34 @@ def test_streaming_response_known_size(test_client_factory):
     client: TestClient = test_client_factory(app)
     response = client.get("/")
     assert response.headers["content-length"] == "10"
+
+
+@pytest.mark.anyio
+async def test_streaming_response_stops_if_receiving_http_disconnect():
+    streamed = 0
+
+    disconnected = anyio.Event()
+
+    async def receive_disconnect():
+        await disconnected.wait()
+        return {"type": "http.disconnect"}
+
+    async def send(message):
+        nonlocal streamed
+        if message["type"] == "http.response.body":
+            streamed += len(message.get("body", b""))
+            # Simulate disconnection after download has started
+            if streamed >= 16:
+                disconnected.set()
+
+    async def stream_indefinitely():
+        while True:
+            # Need a sleep for the event loop to switch to another task
+            await anyio.sleep(0)
+            yield b"chunk "
+
+    response = StreamingResponse(content=stream_indefinitely())
+
+    with anyio.move_on_after(1) as cancel_scope:
+        await response({}, receive_disconnect, send)
+    assert not cancel_scope.cancel_called, "Content streaming should stop itself."