Fix support for generator-based WSGI apps (#887)

* Handle generator WSGI app

* Lint code

* Add type annotations

* Add more tests

* Refactor test to use application_factory

* Remove content length as it's misleading

* Add test for WSGI generator

* Add test for empty generator

* Remove previous tests

* Move docstring to a comment

* Fix whitespace

* Fix name of function

Co-Authored-By: Florimond Manca <florimond.manca@gmail.com>
* Update tests/test_wsgi.py

Co-Authored-By: Florimond Manca <florimond.manca@gmail.com>
* Update tests/test_wsgi.py

Co-Authored-By: Florimond Manca <florimond.manca@gmail.com>
* Update httpx/_dispatch/wsgi.py

Co-Authored-By: Florimond Manca <florimond.manca@gmail.com>
Co-authored-by: Florimond Manca <florimond.manca@gmail.com>

diff --git a/httpx/_dispatch/wsgi.py b/httpx/_dispatch/wsgi.py
index 60540a8b5cc9a5e793a711e0af2050698f2ce819..86227733ddf1ec9092861e13f2296ce02f313831 100644 (file)
--- a/httpx/_dispatch/wsgi.py
+++ b/httpx/_dispatch/wsgi.py
@@ -1,4 +1,5 @@
 import io
+import itertools
 import typing
 
 from .._config import TimeoutTypes
@@ -7,6 +8,14 @@ from .._models import Request, Response
 from .base import SyncDispatcher
 
 
+def _skip_leading_empty_chunks(body: typing.Iterable) -> typing.Iterable:
+    body = iter(body)
+    for chunk in body:
+        if chunk:
+            return itertools.chain([chunk], body)
+    return []
+
+
 class WSGIDispatch(SyncDispatcher):
     """
     A custom SyncDispatcher that handles sending requests directly to an WSGI app.
@@ -88,6 +97,9 @@ class WSGIDispatch(SyncDispatcher):
             seen_exc_info = exc_info
 
         result = self.app(environ, start_response)
+        # This is needed because the status returned by start_response
+        # shouldn't be used until the first non-empty chunk has been served.
+        result = _skip_leading_empty_chunks(result)
 
         assert seen_status is not None
         assert seen_response_headers is not None

diff --git a/tests/test_wsgi.py b/tests/test_wsgi.py
index 0dcf99535dc23a7c5a9623548aaf07b6ab1f2376..1786d9ef28859dc5da71d10700c11e0ffeaa5162 100644 (file)
--- a/tests/test_wsgi.py
+++ b/tests/test_wsgi.py
@@ -5,18 +5,20 @@ import pytest
 import httpx
 
 
-def hello_world(environ, start_response):
-    status = "200 OK"
-    output = b"Hello, World!"
+def application_factory(output):
+    def application(environ, start_response):
+        status = "200 OK"
 
-    response_headers = [
-        ("Content-type", "text/plain"),
-        ("Content-Length", str(len(output))),
-    ]
+        response_headers = [
+            ("Content-type", "text/plain"),
+        ]
 
-    start_response(status, response_headers)
+        start_response(status, response_headers)
 
-    return [output]
+        for item in output:
+            yield item
+
+    return application
 
 
 def echo_body(environ, start_response):
@@ -25,7 +27,6 @@ def echo_body(environ, start_response):
 
     response_headers = [
         ("Content-type", "text/plain"),
-        ("Content-Length", str(len(output))),
     ]
 
     start_response(status, response_headers)
@@ -56,7 +57,6 @@ def raise_exc(environ, start_response):
 
     response_headers = [
         ("Content-type", "text/plain"),
-        ("Content-Length", str(len(output))),
     ]
 
     try:
@@ -69,7 +69,7 @@ def raise_exc(environ, start_response):
 
 
 def test_wsgi():
-    client = httpx.Client(app=hello_world)
+    client = httpx.Client(app=application_factory([b"Hello, World!"]))
     response = client.get("http://www.example.org/")
     assert response.status_code == 200
     assert response.text == "Hello, World!"
@@ -93,3 +93,19 @@ def test_wsgi_exc():
     client = httpx.Client(app=raise_exc)
     with pytest.raises(ValueError):
         client.get("http://www.example.org/")
+
+
+def test_wsgi_generator():
+    output = [b"", b"", b"Some content", b" and more content"]
+    client = httpx.Client(app=application_factory(output))
+    response = client.get("http://www.example.org/")
+    assert response.status_code == 200
+    assert response.text == "Some content and more content"
+
+
+def test_wsgi_generator_empty():
+    output = [b"", b"", b"", b""]
+    client = httpx.Client(app=application_factory(output))
+    response = client.get("http://www.example.org/")
+    assert response.status_code == 200
+    assert response.text == ""