Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 38 additions & 3 deletions samcli/local/apigw/local_apigw_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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):
"""
Expand Down
50 changes: 50 additions & 0 deletions tests/unit/local/apigw/test_local_apigw_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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"}'
Expand Down