Support for response multiValueHeaders locally

samcli/local/apigw/local_apigw_service.py

@@ -152,7 +152,7 @@ def _request_handler(self, **kwargs):
                                                                      route.binary_types,
                                                                      request)
         except (KeyError, TypeError, ValueError):
-            LOG.error("Function returned an invalid response (must include one of: body, headers or "
+            LOG.error("Function returned an invalid response (must include one of: body, headers, multiValueHeaders or "
                       "statusCode in the response object). Response received: %s", lambda_response)
             return ServiceErrorResponses.lambda_failure_response()
 
@@ -193,8 +193,9 @@ def _parse_lambda_output(lambda_output, binary_types, flask_request):
         if not isinstance(json_output, dict):
             raise TypeError("Lambda returned %{s} instead of dict", type(json_output))
 
+        headers = LocalApigwService._merge_response_headers(json_output.get("headers") or {},
+                                                            json_output.get("multiValueHeaders") or {})
         status_code = json_output.get("statusCode") or 200
-        headers = CaseInsensitiveDict(json_output.get("headers") or {})
         body = json_output.get("body") or "no data"
         is_base_64_encoded = json_output.get("isBase64Encoded") or False
 
@@ -234,11 +235,45 @@ def _should_base64_decode_body(binary_types, flask_request, lamba_response_heade
         True if the body from the request should be converted to binary, otherwise false
 
         """
-        best_match_mimetype = flask_request.accept_mimetypes.best_match([lamba_response_headers["Content-Type"]])
+
+        # Get the first part of the content-type header, to allow for extras such as text/html; charset=utf-8
+        content_type = lamba_response_headers['Content-Type'].split(";", 1)[0]
+        best_match_mimetype = flask_request.accept_mimetypes.best_match([content_type])
         is_best_match_in_binary_types = best_match_mimetype in binary_types or '*/*' in binary_types
 
         return best_match_mimetype and is_best_match_in_binary_types and is_base_64_encoded
 
+    @staticmethod
+    def _merge_response_headers(headers, multi_headers):
+        """
+        # Merge multiValueHeaders headers with headers
+        # Convert into CSV for Flask compatibility
+        # If you specify values for both headers and multiValueHeaders, API Gateway merges them into a single list.
+        # If the same key-value pair is specified in both, only the values from multiValueHeaders will
+        # appear in the merged list.
+
+        Parameters
+        ----------
+        headers (dict)
+            Headers map from the lambda_response_headers
+        multi_headers (dict)
+            multiValueHeaders map from the lambda_response_headers
+
+        Returns
+        -------
+        Merged list in accordance to the AWS documentation within a CaseInsensitiveDict
+
+        """
+
+        processed_headers = CaseInsensitiveDict(headers)
+
+        # Convert multi_header Lists to CSV Strings for Flask
+        # This gives multiValueHeaders precedence over the API Gateway headers property
+        for _, s in enumerate(multi_headers):
+            processed_headers[s] = ", ".join(multi_headers[s])
+
+        return processed_headers
+
     @staticmethod
     def _construct_event(flask_request, port, binary_types):
         """

tests/unit/local/apigw/test_local_apigw_service.py

@@ -260,6 +260,47 @@ def test_class_initialization(self):
         self.assertEquals(self.api_gateway.path, '/')
 
 
+class TestLambdaHeaderDictionaryMerge(TestCase):
+    def test_empty_dictionaries_produce_empty_result(self):
+        headers = {}
+        multi_value_headers = {}
+
+        result = LocalApigwService._merge_response_headers(headers, multi_value_headers)
+
+        self.assertEquals(result, {})
+
+    def test_headers_are_merged(self):
+        headers = {"h1": "value1", "h2": "value2"}
+        multi_value_headers = {"h3": ["value3"]}
+
+        result = LocalApigwService._merge_response_headers(headers, multi_value_headers)
+
+        self.assertIn("h1", result)
+        self.assertIn("h2", result)
+        self.assertIn("h3", result)
+        self.assertEquals(result["h1"], "value1")
+        self.assertEquals(result["h2"], "value2")
+        self.assertEquals(result["h3"], "value3")
+
+    def test_multivalue_headers_are_turned_into_csv(self):
+        headers = {}
+        multi_value_headers = {"h1": ["a", "b", "c"]}
+
+        result = LocalApigwService._merge_response_headers(headers, multi_value_headers)
+
+        self.assertIn("h1", result)
+        self.assertEquals(result["h1"], "a, b, c")
+
+    def test_multivalue_headers_override_headers_dict(self):
+        headers = {"h1": "ValueA"}
+        multi_value_headers = {"h1": ["ValueB"]}
+
+        result = LocalApigwService._merge_response_headers(headers, multi_value_headers)
+
+        self.assertIn("h1", result)
+        self.assertEquals(result["h1"], "ValueB")
+
+
 class TestServiceParsingLambdaOutput(TestCase):
 
     def test_default_content_type_header_added_with_no_headers(self):
@@ -289,6 +330,15 @@ def test_custom_content_type_header_is_not_modified(self):
         self.assertIn("Content-Type", headers)
         self.assertEquals(headers["Content-Type"], "text/xml")
 
+    def test_custom_content_type_multivalue_header_is_not_modified(self):
+        lambda_output = '{"statusCode": 200, "multiValueHeaders":{"Content-Type": ["text/xml"]}, "body": "{}", ' \
+                        '"isBase64Encoded": false}'
+
+        (_, headers, _) = LocalApigwService._parse_lambda_output(lambda_output, binary_types=[], flask_request=Mock())
+
+        self.assertIn("Content-Type", headers)
+        self.assertEquals(headers["Content-Type"], "text/xml")
+
     def test_extra_values_ignored(self):
         lambda_output = '{"statusCode": 200, "headers": {}, "body": "{\\"message\\":\\"Hello from Lambda\\"}", ' \
                         '"isBase64Encoded": false, "another_key": "some value"}'