Fixed #34342, Refs #33735 -- Fixed test client handling of async stre…

…aming responses. Bug in 0bd2c0c. Co-authored-by: Carlton Gibson <[email protected]>
Yako3 · Feb 17, 2023 · 52b0548 · 52b0548
1 parent bfb8fda
commit 52b0548
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 9 deletions.
diff --git a/django/test/client.py b/django/test/client.py
@@ -116,6 +116,16 @@ def closing_iterator_wrapper(iterable, close):
         request_finished.connect(close_old_connections)
 
 
+async def aclosing_iterator_wrapper(iterable, close):
+    try:
+        async for chunk in iterable:
+            yield chunk
+    finally:
+        request_finished.disconnect(close_old_connections)
+        close()  # will fire request_finished
+        request_finished.connect(close_old_connections)
+
+
 def conditional_content_removal(request, response):
     """
     Simulate the behavior of most web servers by removing the content of
@@ -174,9 +184,14 @@ def __call__(self, environ):
 
         # Emulate a WSGI server by calling the close method on completion.
         if response.streaming:
-            response.streaming_content = closing_iterator_wrapper(
-                response.streaming_content, response.close
-            )
+            if response.is_async:
+                response.streaming_content = aclosing_iterator_wrapper(
+                    response.streaming_content, response.close
+                )
+            else:
+                response.streaming_content = closing_iterator_wrapper(
+                    response.streaming_content, response.close
+                )
         else:
             request_finished.disconnect(close_old_connections)
             response.close()  # will fire request_finished
@@ -223,12 +238,14 @@ async def __call__(self, scope):
         response.asgi_request = request
         # Emulate a server by calling the close method on completion.
         if response.streaming:
-            response.streaming_content = await sync_to_async(
-                closing_iterator_wrapper, thread_sensitive=False
-            )(
-                response.streaming_content,
-                response.close,
-            )
+            if response.is_async:
+                response.streaming_content = aclosing_iterator_wrapper(
+                    response.streaming_content, response.close
+                )
+            else:
+                response.streaming_content = closing_iterator_wrapper(
+                    response.streaming_content, response.close
+                )
         else:
             request_finished.disconnect(close_old_connections)
             # Will fire request_finished.

diff --git a/tests/handlers/tests.py b/tests/handlers/tests.py
@@ -253,6 +253,16 @@ def test_streaming(self):
         self.assertEqual(response.status_code, 200)
         self.assertEqual(b"".join(list(response)), b"streaming content")
 
+    def test_async_streaming(self):
+        response = self.client.get("/async_streaming/")
+        self.assertEqual(response.status_code, 200)
+        msg = (
+            "StreamingHttpResponse must consume asynchronous iterators in order to "
+            "serve them synchronously. Use a synchronous iterator instead."
+        )
+        with self.assertWarnsMessage(Warning, msg):
+            self.assertEqual(b"".join(list(response)), b"streaming content")
+
 
 class ScriptNameTests(SimpleTestCase):
     def test_get_script_name(self):
@@ -329,3 +339,10 @@ async def test_sync_streaming(self):
             self.assertEqual(
                 b"".join([chunk async for chunk in response]), b"streaming content"
             )
+
+    async def test_async_streaming(self):
+        response = await self.async_client.get("/async_streaming/")
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(
+            b"".join([chunk async for chunk in response]), b"streaming content"
+        )
diff --git a/tests/handlers/urls.py b/tests/handlers/urls.py
@@ -8,6 +8,7 @@
     path("no_response_fbv/", views.no_response),
     path("no_response_cbv/", views.NoResponse()),
     path("streaming/", views.streaming),
+    path("async_streaming/", views.async_streaming),
     path("in_transaction/", views.in_transaction),
     path("not_in_transaction/", views.not_in_transaction),
     path("not_in_transaction_using_none/", views.not_in_transaction_using_none),

diff --git a/tests/handlers/views.py b/tests/handlers/views.py
@@ -65,6 +65,15 @@ async def async_regular(request):
     return HttpResponse(b"regular content")
 
 
+async def async_streaming(request):
+    async def async_streaming_generator():
+        yield b"streaming"
+        yield b" "
+        yield b"content"
+
+    return StreamingHttpResponse(async_streaming_generator())
+
+
 class CoroutineClearingView:
     def __call__(self, request):
         """Return an unawaited coroutine (common error for async views)."""