From 2940e481cf7e5eb2cf93c96aeb243c6f7cf29063 Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Tue, 20 Dec 2022 22:52:51 +0000 Subject: [PATCH 1/2] fix(headers-serializer): skiping none values --- .../shared/headers_serializer.py | 25 +++++++++++-------- tests/functional/test_headers_serializer.py | 18 +++++++++++++ 2 files changed, 32 insertions(+), 11 deletions(-) diff --git a/aws_lambda_powertools/shared/headers_serializer.py b/aws_lambda_powertools/shared/headers_serializer.py index 796fd9aeae3..04965988bd9 100644 --- a/aws_lambda_powertools/shared/headers_serializer.py +++ b/aws_lambda_powertools/shared/headers_serializer.py @@ -44,7 +44,8 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Coo if isinstance(values, str): combined_headers[key] = values else: - combined_headers[key] = ", ".join(values) + if values: + combined_headers[key] = ", ".join(values) return {"headers": combined_headers, "cookies": list(map(str, cookies))} @@ -65,8 +66,9 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Coo if isinstance(values, str): payload[key].append(values) else: - for value in values: - payload[key].append(value) + if values: + for value in values: + payload[key].append(value) if cookies: payload.setdefault("Set-Cookie", []) @@ -101,13 +103,14 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Coo if isinstance(values, str): payload["headers"][key] = values else: - if len(values) > 1: - warnings.warn( - f"Can't encode more than one header value for the same key ('{key}') in the response. " - "Did you enable multiValueHeaders on the ALB Target Group?" - ) - - # We can only set one header per key, send the last one - payload["headers"][key] = values[-1] + if values: + if len(values) > 1: + warnings.warn( + f"Can't encode more than one header value for the same key ('{key}') in the response. " + "Did you enable multiValueHeaders on the ALB Target Group?" + ) + + # We can only set one header per key, send the last one + payload["headers"][key] = values[-1] return payload diff --git a/tests/functional/test_headers_serializer.py b/tests/functional/test_headers_serializer.py index 8a27ce8baa8..204df5ad614 100644 --- a/tests/functional/test_headers_serializer.py +++ b/tests/functional/test_headers_serializer.py @@ -145,3 +145,21 @@ def test_single_value_headers_with_multiple_header_values_warning(): payload = serializer.serialize(cookies=[], headers=headers) assert payload["headers"]["Foo"] == headers["Foo"][-1] + + +def test_http_api_headers_serializer_with_null_values(): + serializer = HttpApiHeadersSerializer() + payload = serializer.serialize(headers={"Foo": None}, cookies=[]) + assert payload == {"headers": {}, "cookies": []} + + +def test_multi_value_headers_serializer_with_null_values(): + serializer = MultiValueHeadersSerializer() + payload = serializer.serialize(headers={"Foo": None}, cookies=[]) + assert payload == {"multiValueHeaders": {}} + + +def test_single_value_headers_serializer_with_null_values(): + serializer = SingleValueHeadersSerializer() + payload = serializer.serialize(headers={"Foo": None}, cookies=[]) + assert payload["headers"] == {} From b886504cd675451ac0ab26d40f5f5097283c13d8 Mon Sep 17 00:00:00 2001 From: heitorlessa Date: Tue, 20 Dec 2022 21:28:28 -0300 Subject: [PATCH 2/2] chore(complexity): checks for null over falsy, list.extend over nested loop --- .../shared/headers_serializer.py | 37 +++++++++++-------- 1 file changed, 22 insertions(+), 15 deletions(-) diff --git a/aws_lambda_powertools/shared/headers_serializer.py b/aws_lambda_powertools/shared/headers_serializer.py index 04965988bd9..b4e9f7c7c5f 100644 --- a/aws_lambda_powertools/shared/headers_serializer.py +++ b/aws_lambda_powertools/shared/headers_serializer.py @@ -41,11 +41,14 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Coo # Duplicate headers are combined with commas and included in the headers field. combined_headers: Dict[str, str] = {} for key, values in headers.items(): + # omit headers with explicit null values + if values is None: + continue + if isinstance(values, str): combined_headers[key] = values else: - if values: - combined_headers[key] = ", ".join(values) + combined_headers[key] = ", ".join(values) return {"headers": combined_headers, "cookies": list(map(str, cookies))} @@ -61,14 +64,15 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Coo https://docs.aws.amazon.com/elasticloadbalancing/latest/application/lambda-functions.html#multi-value-headers-response """ payload: Dict[str, List[str]] = defaultdict(list) - for key, values in headers.items(): + # omit headers with explicit null values + if values is None: + continue + if isinstance(values, str): payload[key].append(values) else: - if values: - for value in values: - payload[key].append(value) + payload[key].extend(values) if cookies: payload.setdefault("Set-Cookie", []) @@ -100,17 +104,20 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Coo payload["headers"]["Set-Cookie"] = str(cookies[-1]) for key, values in headers.items(): + # omit headers with explicit null values + if values is None: + continue + if isinstance(values, str): payload["headers"][key] = values else: - if values: - if len(values) > 1: - warnings.warn( - f"Can't encode more than one header value for the same key ('{key}') in the response. " - "Did you enable multiValueHeaders on the ALB Target Group?" - ) - - # We can only set one header per key, send the last one - payload["headers"][key] = values[-1] + if len(values) > 1: + warnings.warn( + f"Can't encode more than one header value for the same key ('{key}') in the response. " + "Did you enable multiValueHeaders on the ALB Target Group?" + ) + + # We can only set one header per key, send the last one + payload["headers"][key] = values[-1] return payload