Support ASGI middleware chains by only yielding http.disconnect when response has been fully received (#153)

* Yield a disconnect on the second receive call

* Stop http.disconnect being sent before all response body is sent

* Test the correct sequence of messages

* Shorter line

Co-authored-by: peterstone2017 <[email protected]>

Author: tonybaloney

azure/functions/_http_asgi.py

@@ -68,6 +68,7 @@ def __init__(self):
         self._headers: Union[Headers, Dict] = {}
         self._buffer: List[bytes] = []
         self._request_body: Optional[bytes] = b""
+        self._has_received_response: bool = False
 
     @classmethod
     async def from_app(cls, app, scope: Dict[str, Any],
@@ -95,6 +96,7 @@ def _handle_http_response_start(self, message: Dict[str, Any]):
 
     def _handle_http_response_body(self, message: Dict[str, Any]):
         self._buffer.append(message["body"])
+        self._has_received_response = not message.get("more_body", False)
         # XXX : Chunked bodies not supported, see
         # https://github.com/Azure/azure-functions-host/issues/4926
 
@@ -106,14 +108,15 @@ async def _receive(self):
                 "more_body": False,
             }
             self._request_body = None
+            return reply
         else:
-            reply = {
+            while not self._has_received_response:
+                await asyncio.sleep(0.1)
+            return {
                 "type": "http.disconnect",
             }
-        return reply
 
     async def _send(self, message):
-        logging.debug(f"Received {message} from ASGI worker.")
         if message["type"] == "http.response.start":
             self._handle_http_response_start(message)
         elif message["type"] == "http.response.body":
@@ -142,7 +145,7 @@ def __init__(self, app):
         main = func.AsgiMiddleware(app).main
         """
         if not self._usage_reported:
-            self._logger.info("Instantiating Azure Functions ASGI middleware.")
+            self._logger.debug("Starting Azure Functions ASGI middleware.")
             self._usage_reported = True
 
         self._app = app

tests/test_http_asgi.py

@@ -66,9 +66,6 @@ async def __call__(self, scope, receive, send):
         assert isinstance(self.received_request['body'], bytes)
         assert isinstance(self.received_request['more_body'], bool)
 
-        self.next_request = await receive()
-        assert self.next_request['type'] == 'http.disconnect'
-
         await send(
             {
                 "type": "http.response.start",
@@ -83,6 +80,9 @@ async def __call__(self, scope, receive, send):
             }
         )
 
+        self.next_request = await receive()
+        assert self.next_request['type'] == 'http.disconnect'
+
 
 class TestHttpAsgiMiddleware(unittest.TestCase):
     def _generate_func_request(