Support multiValueHeaders on Response (Kludex#129)

* support multiValueHeaders

Author: koxudaxi

mangum/protocols/http.py

@@ -115,9 +115,23 @@ async def send(self, message: Message) -> None:
             and message_type == "http.response.start"
         ):
             self.response["statusCode"] = message["status"]
-            self.response["headers"] = {
-                k.decode().lower(): v.decode() for k, v in message.get("headers", [])
-            }
+            headers: typing.Dict[str, str] = {}
+            multi_value_headers: typing.Dict[str, typing.List[str]] = {}
+            for key, value in message.get("headers", []):
+                lower_key = key.decode().lower()
+                if lower_key in multi_value_headers:
+                    multi_value_headers[lower_key].append(value.decode())
+                elif lower_key in headers:
+                    multi_value_headers[lower_key] = [
+                        headers.pop(lower_key),
+                        value.decode(),
+                    ]
+                else:
+                    headers[lower_key] = value.decode()
+
+            self.response["headers"] = headers
+            if multi_value_headers:
+                self.response["multiValueHeaders"] = multi_value_headers
             self.state = HTTPCycleState.RESPONSE
 
         elif (

tests/test_http.py

@@ -682,3 +682,57 @@ async def app(scope, receive, send):
         "headers": {"content-type": "text/plain; charset=utf-8"},
         "body": "Hello, world!",
     }
+
+
+@pytest.mark.parametrize(
+    "mock_http_event,response_headers,expected_headers,expected_multi_value_headers",
+    [
+        [
+            ["GET", None, None],
+            [[b"key1", b"value1"], [b"key2", b"value2"]],
+            {"key1": "value1", "key2": "value2"},
+            {},
+        ],
+        [
+            ["GET", None, None],
+            [[b"key1", b"value1"], [b"key1", b"value2"]],
+            {},
+            {"key1": ["value1", "value2"]},
+        ],
+        [
+            ["GET", None, None],
+            [[b"key1", b"value1"], [b"key1", b"value2"], [b"key1", b"value3"]],
+            {},
+            {"key1": ["value1", "value2", "value3"]},
+        ],
+        [["GET", None, None], [], {}, {}],
+    ],
+    indirect=["mock_http_event"],
+)
+def test_http_response_headers(
+    mock_http_event, response_headers, expected_headers, expected_multi_value_headers
+) -> None:
+    async def app(scope, receive, send):
+        await send(
+            {
+                "type": "http.response.start",
+                "status": 200,
+                "headers": [[b"content-type", b"text/plain; charset=utf-8"]]
+                + response_headers,
+            }
+        )
+        await send({"type": "http.response.body", "body": b"Hello, world!"})
+
+    handler = Mangum(app, lifespan="off")
+    response = handler(mock_http_event, {})
+    expected = {
+        "statusCode": 200,
+        "isBase64Encoded": False,
+        "headers": {"content-type": "text/plain; charset=utf-8"},
+        "body": "Hello, world!",
+    }
+    if expected_headers:
+        expected["headers"].update(expected_headers)
+    if expected_multi_value_headers:
+        expected["multiValueHeaders"] = expected_multi_value_headers
+    assert response == expected