add test for ASF request streaming

alexrashed · alexrashed · commit f4389cff9734 · 2023-01-25T08:43:58.000+01:00
diff --git a/localstack/aws/handlers/logging.py b/localstack/aws/handlers/logging.py
@@ -3,7 +3,8 @@
 import logging
 from functools import cached_property
 from typing import Type
-from wsgiref.headers import Headers
+
+from werkzeug.datastructures import Headers
 
 from localstack.aws.api import RequestContext, ServiceException
 from localstack.aws.chain import ExceptionHandler, HandlerChain
@@ -81,7 +82,7 @@ def _prepare_logger(self, logger: logging.Logger, formatter: Type):
             logger.addHandler(handler)
         return logger
 
-    def log_http_response(
+    def _log_http_response(
         self,
         logger: logging.Logger,
         request: Request,
@@ -94,13 +95,14 @@ def log_http_response(
         The given response data is returned by this function, which allows the usage as a log interceptor for streamed
         response data.
 
+        :param logger: HTTP logger to log the request onto
         :param request: HTTP request data (containing useful metadata like the HTTP method and path)
         :param response_status: HTTP status of the response to log
         :param response_headers: HTTP headers of the response to log
         :param response_data: HTTP body of the response to log
         :return: response data
         """
-        self.http_logger.info(
+        logger.info(
             "%s %s => %d",
             request.method,
             request.path,
@@ -173,7 +175,7 @@ def _log(self, context: RequestContext, response: Response):
             if hasattr(response.response, "__iter__"):
                 # If the response is streamed, wrap the response data's iterator which logs all values when they are consumed
                 log_partial = functools.partial(
-                    self.log_http_response,
+                    self._log_http_response,
                     http_logger,
                     context.request,
                     response.status_code,
@@ -183,7 +185,7 @@ def _log(self, context: RequestContext, response: Response):
                 response.set_response(wrapped_response_iterator)
             else:
                 # If the response is synchronous, we log the data directly
-                self.log_http_response(
+                self._log_http_response(
                     http_logger,
                     context.request,
                     response.status_code,
diff --git a/tests/integration/test_edge.py b/tests/integration/test_edge.py
@@ -15,6 +15,7 @@
 from localstack import config
 from localstack.aws.accounts import get_aws_account_id
 from localstack.constants import APPLICATION_JSON, HEADER_LOCALSTACK_EDGE_URL
+from localstack.http import Request as HttpRequest
 from localstack.http import Response as HttpResponse
 from localstack.http import route
 from localstack.http.request import get_full_raw_path
@@ -452,3 +453,39 @@ def chunk_generator():
                 assert element == next(chunk_iterator).decode("utf-8")
                 # make sure the queue is empty
                 assert queue.empty()
+
+    def test_streaming_request(self):
+        """Test if request data is correctly streamed when the HTTP request uses a generator."""
+        queue = Queue()
+
+        # generate test events
+        elements = [bytes(f"element-{n:02d}", "latin-1") for n in range(100)]
+
+        @route("/streaming-endpoint-test")
+        def streaming_endpoint(request: HttpRequest):
+            # take the first element from the list and add it to the queue
+            next_element = elements.pop(0)
+            queue.put(next_element)
+            # process each element, and add the next element to the generating queue after the element is processed
+            while data := request.input_stream.read(10):
+                assert next_element == data
+                if len(elements) > 0:
+                    next_element = elements.pop(0)
+                    queue.put(next_element)
+
+        def data_generator():
+            try:
+                # wait for a new element in the queue, if one is received, send it immediately
+                while chunk := queue.get(timeout=0.1):
+                    yield chunk
+            except Empty:
+                # the queue is empty for more than 100 ms, we end here
+                pass
+
+        # register the streaming test endpoint
+        with self.register_edge_route(streaming_endpoint):
+            # send a streamed request data to the registered endpoint
+            requests.post(f"{config.get_edge_url()}/streaming-endpoint-test", data=data_generator())
+
+        # ensure that all elements have been processed
+        assert len(elements) == 0
