From 2deb7c88639d52ab0d2faa96c9030457d4e36c00 Mon Sep 17 00:00:00 2001 From: Ruben Fonseca Date: Fri, 2 Sep 2022 17:21:50 +0200 Subject: [PATCH] feat(event_handler): add cookies as 1st class citizen in v2 (#1487) Co-authored-by: Heitor Lessa Co-authored-by: Heitor Lessa --- .../event_handler/api_gateway.py | 5 +- aws_lambda_powertools/shared/cookies.py | 118 +++++++++++ .../shared/headers_serializer.py | 16 +- docs/core/event_handler/api_gateway.md | 2 +- docs/upgrade.md | 2 +- .../src/fine_grained_responses.py | 3 +- .../src/fine_grained_responses_output.json | 2 +- .../e2e/event_handler/handlers/alb_handler.py | 20 +- .../handlers/api_gateway_http_handler.py | 20 +- .../handlers/api_gateway_rest_handler.py | 20 +- .../handlers/lambda_function_url_handler.py | 20 +- tests/e2e/event_handler/infrastructure.py | 4 +- .../event_handler/test_header_serializer.py | 188 ++++++++++++------ .../event_handler/test_api_gateway.py | 9 +- .../event_handler/test_lambda_function_url.py | 5 +- 15 files changed, 333 insertions(+), 101 deletions(-) create mode 100644 aws_lambda_powertools/shared/cookies.py diff --git a/aws_lambda_powertools/event_handler/api_gateway.py b/aws_lambda_powertools/event_handler/api_gateway.py index 2d315fcc434..126eee8b0aa 100644 --- a/aws_lambda_powertools/event_handler/api_gateway.py +++ b/aws_lambda_powertools/event_handler/api_gateway.py @@ -15,6 +15,7 @@ from aws_lambda_powertools.event_handler import content_types from aws_lambda_powertools.event_handler.exceptions import NotFoundError, ServiceError from aws_lambda_powertools.shared import constants +from aws_lambda_powertools.shared.cookies import Cookie from aws_lambda_powertools.shared.functions import resolve_truthy_env_var_choice from aws_lambda_powertools.shared.json_encoder import Encoder from aws_lambda_powertools.utilities.data_classes import ( @@ -147,7 +148,7 @@ def __init__( content_type: Optional[str], body: Union[str, bytes, None], headers: Optional[Dict[str, Union[str, List[str]]]] = None, - cookies: Optional[List[str]] = None, + cookies: Optional[List[Cookie]] = None, ): """ @@ -162,7 +163,7 @@ def __init__( Optionally set the response body. Note: bytes body will be automatically base64 encoded headers: dict[str, Union[str, List[str]]] Optionally set specific http headers. Setting "Content-Type" here would override the `content_type` value. - cookies: list[str] + cookies: list[Cookie] Optionally set cookies. """ self.status_code = status_code diff --git a/aws_lambda_powertools/shared/cookies.py b/aws_lambda_powertools/shared/cookies.py new file mode 100644 index 00000000000..944bcb5dc9f --- /dev/null +++ b/aws_lambda_powertools/shared/cookies.py @@ -0,0 +1,118 @@ +from datetime import datetime +from enum import Enum +from io import StringIO +from typing import List, Optional + + +class SameSite(Enum): + """ + SameSite allows a server to define a cookie attribute making it impossible for + the browser to send this cookie along with cross-site requests. The main + goal is to mitigate the risk of cross-origin information leakage, and provide + some protection against cross-site request forgery attacks. + + See https://tools.ietf.org/html/draft-ietf-httpbis-cookie-same-site-00 for details. + """ + + DEFAULT_MODE = "" + LAX_MODE = "Lax" + STRICT_MODE = "Strict" + NONE_MODE = "None" + + +def _format_date(timestamp: datetime) -> str: + # Specification example: Wed, 21 Oct 2015 07:28:00 GMT + return timestamp.strftime("%a, %d %b %Y %H:%M:%S GMT") + + +class Cookie: + """ + A Cookie represents an HTTP cookie as sent in the Set-Cookie header of an + HTTP response or the Cookie header of an HTTP request. + + See https://tools.ietf.org/html/rfc6265 for details. + """ + + def __init__( + self, + name: str, + value: str, + path: str = "", + domain: str = "", + secure: bool = True, + http_only: bool = False, + max_age: Optional[int] = None, + expires: Optional[datetime] = None, + same_site: Optional[SameSite] = None, + custom_attributes: Optional[List[str]] = None, + ): + """ + + Parameters + ---------- + name: str + The name of this cookie, for example session_id + value: str + The cookie value, for instance an uuid + path: str + The path for which this cookie is valid. Optional + domain: str + The domain for which this cookie is valid. Optional + secure: bool + Marks the cookie as secure, only sendable to the server with an encrypted request over the HTTPS protocol + http_only: bool + Enabling this attribute makes the cookie inaccessible to the JavaScript `Document.cookie` API + max_age: Optional[int] + Defines the period of time after which the cookie is invalid. Use negative values to force cookie deletion. + expires: Optional[datetime] + Defines a date where the permanent cookie expires. + same_site: Optional[SameSite] + Determines if the cookie should be sent to third party websites + custom_attributes: Optional[List[str]] + List of additional custom attributes to set on the cookie + """ + self.name = name + self.value = value + self.path = path + self.domain = domain + self.secure = secure + self.expires = expires + self.max_age = max_age + self.http_only = http_only + self.same_site = same_site + self.custom_attributes = custom_attributes + + def __str__(self) -> str: + payload = StringIO() + payload.write(f"{self.name}={self.value}") + + if self.path: + payload.write(f"; Path={self.path}") + + if self.domain: + payload.write(f"; Domain={self.domain}") + + if self.expires: + payload.write(f"; Expires={_format_date(self.expires)}") + + if self.max_age: + if self.max_age > 0: + payload.write(f"; MaxAge={self.max_age}") + else: + # negative or zero max-age should be set to 0 + payload.write("; MaxAge=0") + + if self.http_only: + payload.write("; HttpOnly") + + if self.secure: + payload.write("; Secure") + + if self.same_site: + payload.write(f"; SameSite={self.same_site.value}") + + if self.custom_attributes: + for attr in self.custom_attributes: + payload.write(f"; {attr}") + + return payload.getvalue() diff --git a/aws_lambda_powertools/shared/headers_serializer.py b/aws_lambda_powertools/shared/headers_serializer.py index 4db7effe81b..796fd9aeae3 100644 --- a/aws_lambda_powertools/shared/headers_serializer.py +++ b/aws_lambda_powertools/shared/headers_serializer.py @@ -2,6 +2,8 @@ from collections import defaultdict from typing import Any, Dict, List, Union +from aws_lambda_powertools.shared.cookies import Cookie + class BaseHeadersSerializer: """ @@ -9,7 +11,7 @@ class BaseHeadersSerializer: ALB and Lambda Function URL response payload. """ - def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str]) -> Dict[str, Any]: + def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Cookie]) -> Dict[str, Any]: """ Serializes headers and cookies according to the request type. Returns a dict that can be merged with the response payload. @@ -25,7 +27,7 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str class HttpApiHeadersSerializer(BaseHeadersSerializer): - def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str]) -> Dict[str, Any]: + def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Cookie]) -> Dict[str, Any]: """ When using HTTP APIs or LambdaFunctionURLs, everything is taken care automatically for us. We can directly assign a list of cookies and a dict of headers to the response payload, and the @@ -44,11 +46,11 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str else: combined_headers[key] = ", ".join(values) - return {"headers": combined_headers, "cookies": cookies} + return {"headers": combined_headers, "cookies": list(map(str, cookies))} class MultiValueHeadersSerializer(BaseHeadersSerializer): - def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str]) -> Dict[str, Any]: + def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Cookie]) -> Dict[str, Any]: """ When using REST APIs, headers can be encoded using the `multiValueHeaders` key on the response. This is also the case when using an ALB integration with the `multiValueHeaders` option enabled. @@ -69,13 +71,13 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str if cookies: payload.setdefault("Set-Cookie", []) for cookie in cookies: - payload["Set-Cookie"].append(cookie) + payload["Set-Cookie"].append(str(cookie)) return {"multiValueHeaders": payload} class SingleValueHeadersSerializer(BaseHeadersSerializer): - def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str]) -> Dict[str, Any]: + def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Cookie]) -> Dict[str, Any]: """ The ALB integration has `multiValueHeaders` disabled by default. If we try to set multiple headers with the same key, or more than one cookie, print a warning. @@ -93,7 +95,7 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[str ) # We can only send one cookie, send the last one - payload["headers"]["Set-Cookie"] = cookies[-1] + payload["headers"]["Set-Cookie"] = str(cookies[-1]) for key, values in headers.items(): if isinstance(values, str): diff --git a/docs/core/event_handler/api_gateway.md b/docs/core/event_handler/api_gateway.md index c4cae718289..934465d6b96 100644 --- a/docs/core/event_handler/api_gateway.md +++ b/docs/core/event_handler/api_gateway.md @@ -323,7 +323,7 @@ You can use the `Response` class to have full control over the response. For exa === "fine_grained_responses.py" - ```python hl_lines="7 24-29" + ```python hl_lines="7 25-30" --8<-- "examples/event_handler_rest/src/fine_grained_responses.py" ``` diff --git a/docs/upgrade.md b/docs/upgrade.md index 20cf4aa25a6..3d1257f1c12 100644 --- a/docs/upgrade.md +++ b/docs/upgrade.md @@ -53,7 +53,7 @@ def get_todos(): return Response( # ... headers={"Content-Type": ["text/plain"]}, - cookies=["CookieName=CookieValue"] + cookies=[Cookie(name="session_id", value="12345", secure=True, http_only=True)], ) ``` diff --git a/examples/event_handler_rest/src/fine_grained_responses.py b/examples/event_handler_rest/src/fine_grained_responses.py index 4892de9c798..639b6a5b120 100644 --- a/examples/event_handler_rest/src/fine_grained_responses.py +++ b/examples/event_handler_rest/src/fine_grained_responses.py @@ -6,6 +6,7 @@ from aws_lambda_powertools import Logger, Tracer from aws_lambda_powertools.event_handler import APIGatewayRestResolver, Response, content_types from aws_lambda_powertools.logging import correlation_paths +from aws_lambda_powertools.shared.cookies import Cookie from aws_lambda_powertools.utilities.typing import LambdaContext tracer = Tracer() @@ -26,7 +27,7 @@ def get_todos(): content_type=content_types.APPLICATION_JSON, body=todos.json()[:10], headers=custom_headers, - cookies=["=; Secure; Expires="], + cookies=[Cookie(name="session_id", value="12345")], ) diff --git a/examples/event_handler_rest/src/fine_grained_responses_output.json b/examples/event_handler_rest/src/fine_grained_responses_output.json index 1ce606839b1..0b33bd91542 100644 --- a/examples/event_handler_rest/src/fine_grained_responses_output.json +++ b/examples/event_handler_rest/src/fine_grained_responses_output.json @@ -3,7 +3,7 @@ "multiValueHeaders": { "Content-Type": ["application/json"], "X-Transaction-Id": ["3490eea9-791b-47a0-91a4-326317db61a9"], - "Set-Cookie": ["=; Secure; Expires="] + "Set-Cookie": ["session_id=12345; Secure"] }, "body": "{\"todos\":[{\"userId\":1,\"id\":1,\"title\":\"delectus aut autem\",\"completed\":false},{\"userId\":1,\"id\":2,\"title\":\"quis ut nam facilis et officia qui\",\"completed\":false},{\"userId\":1,\"id\":3,\"title\":\"fugiat veniam minus\",\"completed\":false},{\"userId\":1,\"id\":4,\"title\":\"et porro tempora\",\"completed\":true},{\"userId\":1,\"id\":5,\"title\":\"laboriosam mollitia et enim quasi adipisci quia provident illum\",\"completed\":false},{\"userId\":1,\"id\":6,\"title\":\"qui ullam ratione quibusdam voluptatem quia omnis\",\"completed\":false},{\"userId\":1,\"id\":7,\"title\":\"illo expedita consequatur quia in\",\"completed\":false},{\"userId\":1,\"id\":8,\"title\":\"quo adipisci enim quam ut ab\",\"completed\":true},{\"userId\":1,\"id\":9,\"title\":\"molestiae perspiciatis ipsa\",\"completed\":false},{\"userId\":1,\"id\":10,\"title\":\"illo est ratione doloremque quia maiores aut\",\"completed\":true}]}", "isBase64Encoded": false diff --git a/tests/e2e/event_handler/handlers/alb_handler.py b/tests/e2e/event_handler/handlers/alb_handler.py index 4c3f4f9dac3..0e386c82c51 100644 --- a/tests/e2e/event_handler/handlers/alb_handler.py +++ b/tests/e2e/event_handler/handlers/alb_handler.py @@ -3,14 +3,22 @@ app = ALBResolver() -@app.get("/todos") +@app.post("/todos") def hello(): + payload = app.current_event.json_body + + body = payload.get("body", "Hello World") + status_code = payload.get("status_code", 200) + headers = payload.get("headers", {}) + cookies = payload.get("cookies", []) + content_type = headers.get("Content-Type", content_types.TEXT_PLAIN) + return Response( - status_code=200, - content_type=content_types.TEXT_PLAIN, - body="Hello world", - cookies=["CookieMonster", "MonsterCookie"], - headers={"Foo": ["bar", "zbr"]}, + status_code=status_code, + content_type=content_type, + body=body, + cookies=cookies, + headers=headers, ) diff --git a/tests/e2e/event_handler/handlers/api_gateway_http_handler.py b/tests/e2e/event_handler/handlers/api_gateway_http_handler.py index 1a20b730285..990761cd3b9 100644 --- a/tests/e2e/event_handler/handlers/api_gateway_http_handler.py +++ b/tests/e2e/event_handler/handlers/api_gateway_http_handler.py @@ -3,14 +3,22 @@ app = APIGatewayHttpResolver() -@app.get("/todos") +@app.post("/todos") def hello(): + payload = app.current_event.json_body + + body = payload.get("body", "Hello World") + status_code = payload.get("status_code", 200) + headers = payload.get("headers", {}) + cookies = payload.get("cookies", []) + content_type = headers.get("Content-Type", content_types.TEXT_PLAIN) + return Response( - status_code=200, - content_type=content_types.TEXT_PLAIN, - body="Hello world", - cookies=["CookieMonster", "MonsterCookie"], - headers={"Foo": ["bar", "zbr"]}, + status_code=status_code, + content_type=content_type, + body=body, + cookies=cookies, + headers=headers, ) diff --git a/tests/e2e/event_handler/handlers/api_gateway_rest_handler.py b/tests/e2e/event_handler/handlers/api_gateway_rest_handler.py index 2f5ad0b94fa..0aa836cfe74 100644 --- a/tests/e2e/event_handler/handlers/api_gateway_rest_handler.py +++ b/tests/e2e/event_handler/handlers/api_gateway_rest_handler.py @@ -3,14 +3,22 @@ app = APIGatewayRestResolver() -@app.get("/todos") +@app.post("/todos") def hello(): + payload = app.current_event.json_body + + body = payload.get("body", "Hello World") + status_code = payload.get("status_code", 200) + headers = payload.get("headers", {}) + cookies = payload.get("cookies", []) + content_type = headers.get("Content-Type", content_types.TEXT_PLAIN) + return Response( - status_code=200, - content_type=content_types.TEXT_PLAIN, - body="Hello world", - cookies=["CookieMonster", "MonsterCookie"], - headers={"Foo": ["bar", "zbr"]}, + status_code=status_code, + content_type=content_type, + body=body, + cookies=cookies, + headers=headers, ) diff --git a/tests/e2e/event_handler/handlers/lambda_function_url_handler.py b/tests/e2e/event_handler/handlers/lambda_function_url_handler.py index 3fd4b46ea28..c9c825c38d2 100644 --- a/tests/e2e/event_handler/handlers/lambda_function_url_handler.py +++ b/tests/e2e/event_handler/handlers/lambda_function_url_handler.py @@ -3,14 +3,22 @@ app = LambdaFunctionUrlResolver() -@app.get("/todos") +@app.post("/todos") def hello(): + payload = app.current_event.json_body + + body = payload.get("body", "Hello World") + status_code = payload.get("status_code", 200) + headers = payload.get("headers", {}) + cookies = payload.get("cookies", []) + content_type = headers.get("Content-Type", content_types.TEXT_PLAIN) + return Response( - status_code=200, - content_type=content_types.TEXT_PLAIN, - body="Hello world", - cookies=["CookieMonster", "MonsterCookie"], - headers={"Foo": ["bar", "zbr"]}, + status_code=status_code, + content_type=content_type, + body=body, + cookies=cookies, + headers=headers, ) diff --git a/tests/e2e/event_handler/infrastructure.py b/tests/e2e/event_handler/infrastructure.py index 62421b8aac9..735261138f3 100644 --- a/tests/e2e/event_handler/infrastructure.py +++ b/tests/e2e/event_handler/infrastructure.py @@ -61,7 +61,7 @@ def _create_api_gateway_http(self, function: Function): apigw = apigwv2.HttpApi(self.stack, "APIGatewayHTTP", create_default_stage=True) apigw.add_routes( path="/todos", - methods=[apigwv2.HttpMethod.GET], + methods=[apigwv2.HttpMethod.POST], integration=apigwv2integrations.HttpLambdaIntegration("TodosIntegration", function), ) @@ -71,7 +71,7 @@ def _create_api_gateway_rest(self, function: Function): apigw = apigwv1.RestApi(self.stack, "APIGatewayRest", deploy_options=apigwv1.StageOptions(stage_name="dev")) todos = apigw.root.add_resource("todos") - todos.add_method("GET", apigwv1.LambdaIntegration(function, proxy=True)) + todos.add_method("POST", apigwv1.LambdaIntegration(function, proxy=True)) CfnOutput(self.stack, "APIGatewayRestUrl", value=apigw.url) diff --git a/tests/e2e/event_handler/test_header_serializer.py b/tests/e2e/event_handler/test_header_serializer.py index 2b1d51bfb3d..eedb69ccaad 100644 --- a/tests/e2e/event_handler/test_header_serializer.py +++ b/tests/e2e/event_handler/test_header_serializer.py @@ -1,6 +1,9 @@ +from uuid import uuid4 + import pytest from requests import Request +from aws_lambda_powertools.shared.cookies import Cookie from tests.e2e.utils import data_fetcher @@ -36,106 +39,179 @@ def lambda_function_url_endpoint(infrastructure: dict) -> str: def test_alb_headers_serializer(alb_basic_listener_endpoint): # GIVEN url = f"{alb_basic_listener_endpoint}/todos" + body = "Hello World" + status_code = 200 + headers = {"Content-Type": "text/plain", "Vary": ["Accept-Encoding", "User-Agent"]} + cookies = [ + Cookie(name="session_id", value=str(uuid4()), secure=True, http_only=True), + Cookie(name="ab_experiment", value="3"), + ] + last_cookie = cookies[-1] # WHEN - response = data_fetcher.get_http_response(Request(method="GET", url=url)) + response = data_fetcher.get_http_response( + Request( + method="POST", + url=url, + json={"body": body, "status_code": status_code, "headers": headers, "cookies": list(map(str, cookies))}, + ) + ) # THEN - assert response.status_code == 200 - assert response.content == b"Hello world" - assert response.headers["content-type"] == "text/plain" + assert response.status_code == status_code + # response.content is a binary string, needs to be decoded to compare with the real string + assert response.content.decode("ascii") == body - # Only the last header for key "Foo" should be set - assert response.headers["Foo"] == "zbr" + # Only the last header should be set + for key, value in headers.items(): + assert key in response.headers + value = value if isinstance(value, str) else sorted(value)[-1] + assert response.headers[key] == value # Only the last cookie should be set - assert "MonsterCookie" in response.cookies.keys() - assert "CookieMonster" not in response.cookies.keys() + assert len(response.cookies.items()) == 1 + assert last_cookie.name in response.cookies + assert response.cookies.get(last_cookie.name) == last_cookie.value def test_alb_multi_value_headers_serializer(alb_multi_value_header_listener_endpoint): # GIVEN url = f"{alb_multi_value_header_listener_endpoint}/todos" + body = "Hello World" + status_code = 200 + headers = {"Content-Type": "text/plain", "Vary": ["Accept-Encoding", "User-Agent"]} + cookies = [ + Cookie(name="session_id", value=str(uuid4()), secure=True, http_only=True), + Cookie(name="ab_experiment", value="3"), + ] # WHEN - response = data_fetcher.get_http_response(Request(method="GET", url=url)) + response = data_fetcher.get_http_response( + Request( + method="POST", + url=url, + json={"body": body, "status_code": status_code, "headers": headers, "cookies": list(map(str, cookies))}, + ) + ) # THEN - assert response.status_code == 200 - assert response.content == b"Hello world" - assert response.headers["content-type"] == "text/plain" + assert response.status_code == status_code + # response.content is a binary string, needs to be decoded to compare with the real string + assert response.content.decode("ascii") == body - # Only the last header for key "Foo" should be set - assert "Foo" in response.headers - foo_headers = [x.strip() for x in response.headers["Foo"].split(",")] - assert sorted(foo_headers) == ["bar", "zbr"] + for key, value in headers.items(): + assert key in response.headers + value = value if isinstance(value, str) else ", ".join(sorted(value)) - # Only the last cookie should be set - assert "MonsterCookie" in response.cookies.keys() - assert "CookieMonster" in response.cookies.keys() + # ALB sorts the header values randomly, so we have to re-order them for comparison here + returned_value = ", ".join(sorted(response.headers[key].split(", "))) + assert returned_value == value + + for cookie in cookies: + assert cookie.name in response.cookies + assert response.cookies.get(cookie.name) == cookie.value def test_api_gateway_rest_headers_serializer(apigw_rest_endpoint): # GIVEN - url = f"{apigw_rest_endpoint}/todos" + url = f"{apigw_rest_endpoint}todos" + body = "Hello World" + status_code = 200 + headers = {"Content-Type": "text/plain", "Vary": ["Accept-Encoding", "User-Agent"]} + cookies = [ + Cookie(name="session_id", value=str(uuid4()), secure=True, http_only=True), + Cookie(name="ab_experiment", value="3"), + ] # WHEN - response = data_fetcher.get_http_response(Request(method="GET", url=url)) + response = data_fetcher.get_http_response( + Request( + method="POST", + url=url, + json={"body": body, "status_code": status_code, "headers": headers, "cookies": list(map(str, cookies))}, + ) + ) # THEN - assert response.status_code == 200 - assert response.content == b"Hello world" - assert response.headers["content-type"] == "text/plain" + assert response.status_code == status_code + # response.content is a binary string, needs to be decoded to compare with the real string + assert response.content.decode("ascii") == body - # Only the last header for key "Foo" should be set - assert "Foo" in response.headers - foo_headers = [x.strip() for x in response.headers["Foo"].split(",")] - assert sorted(foo_headers) == ["bar", "zbr"] + for key, value in headers.items(): + assert key in response.headers + value = value if isinstance(value, str) else ", ".join(sorted(value)) + assert response.headers[key] == value - # Only the last cookie should be set - assert "MonsterCookie" in response.cookies.keys() - assert "CookieMonster" in response.cookies.keys() + for cookie in cookies: + assert cookie.name in response.cookies + assert response.cookies.get(cookie.name) == cookie.value def test_api_gateway_http_headers_serializer(apigw_http_endpoint): # GIVEN - url = f"{apigw_http_endpoint}/todos" + url = f"{apigw_http_endpoint}todos" + body = "Hello World" + status_code = 200 + headers = {"Content-Type": "text/plain", "Vary": ["Accept-Encoding", "User-Agent"]} + cookies = [ + Cookie(name="session_id", value=str(uuid4()), secure=True, http_only=True), + Cookie(name="ab_experiment", value="3"), + ] # WHEN - response = data_fetcher.get_http_response(Request(method="GET", url=url)) + response = data_fetcher.get_http_response( + Request( + method="POST", + url=url, + json={"body": body, "status_code": status_code, "headers": headers, "cookies": list(map(str, cookies))}, + ) + ) # THEN - assert response.status_code == 200 - assert response.content == b"Hello world" - assert response.headers["content-type"] == "text/plain" + assert response.status_code == status_code + # response.content is a binary string, needs to be decoded to compare with the real string + assert response.content.decode("ascii") == body - # Only the last header for key "Foo" should be set - assert "Foo" in response.headers - foo_headers = [x.strip() for x in response.headers["Foo"].split(",")] - assert sorted(foo_headers) == ["bar", "zbr"] + for key, value in headers.items(): + assert key in response.headers + value = value if isinstance(value, str) else ", ".join(sorted(value)) + assert response.headers[key] == value - # Only the last cookie should be set - assert "MonsterCookie" in response.cookies.keys() - assert "CookieMonster" in response.cookies.keys() + for cookie in cookies: + assert cookie.name in response.cookies + assert response.cookies.get(cookie.name) == cookie.value def test_lambda_function_url_headers_serializer(lambda_function_url_endpoint): # GIVEN url = f"{lambda_function_url_endpoint}todos" # the function url endpoint already has the trailing / + body = "Hello World" + status_code = 200 + headers = {"Content-Type": "text/plain", "Vary": ["Accept-Encoding", "User-Agent"]} + cookies = [ + Cookie(name="session_id", value=str(uuid4()), secure=True, http_only=True), + Cookie(name="ab_experiment", value="3"), + ] # WHEN - response = data_fetcher.get_http_response(Request(method="GET", url=url)) + response = data_fetcher.get_http_response( + Request( + method="POST", + url=url, + json={"body": body, "status_code": status_code, "headers": headers, "cookies": list(map(str, cookies))}, + ) + ) # THEN - assert response.status_code == 200 - assert response.content == b"Hello world" - assert response.headers["content-type"] == "text/plain" - - # Only the last header for key "Foo" should be set - assert "Foo" in response.headers - foo_headers = [x.strip() for x in response.headers["Foo"].split(",")] - assert sorted(foo_headers) == ["bar", "zbr"] - - # Only the last cookie should be set - assert "MonsterCookie" in response.cookies.keys() - assert "CookieMonster" in response.cookies.keys() + assert response.status_code == status_code + # response.content is a binary string, needs to be decoded to compare with the real string + assert response.content.decode("ascii") == body + + for key, value in headers.items(): + assert key in response.headers + value = value if isinstance(value, str) else ", ".join(sorted(value)) + assert response.headers[key] == value + + for cookie in cookies: + assert cookie.name in response.cookies + assert response.cookies.get(cookie.name) == cookie.value diff --git a/tests/functional/event_handler/test_api_gateway.py b/tests/functional/event_handler/test_api_gateway.py index 125a0f8c147..989475a934e 100644 --- a/tests/functional/event_handler/test_api_gateway.py +++ b/tests/functional/event_handler/test_api_gateway.py @@ -30,6 +30,7 @@ UnauthorizedError, ) from aws_lambda_powertools.shared import constants +from aws_lambda_powertools.shared.cookies import Cookie from aws_lambda_powertools.shared.json_encoder import Encoder from aws_lambda_powertools.utilities.data_classes import ( ALBEvent, @@ -98,7 +99,7 @@ def get_lambda() -> Response: def test_api_gateway_v1_cookies(): # GIVEN a Http API V1 proxy type event app = APIGatewayRestResolver() - cookie = "CookieMonster" + cookie = Cookie(name="CookieMonster", value="MonsterCookie") @app.get("/my/path") def get_lambda() -> Response: @@ -111,7 +112,7 @@ def get_lambda() -> Response: # THEN process event correctly # AND set the current_event type as APIGatewayProxyEvent assert result["statusCode"] == 200 - assert result["multiValueHeaders"]["Set-Cookie"] == [cookie] + assert result["multiValueHeaders"]["Set-Cookie"] == ["CookieMonster=MonsterCookie; Secure"] def test_api_gateway(): @@ -158,7 +159,7 @@ def my_path() -> Response: def test_api_gateway_v2_cookies(): # GIVEN a Http API V2 proxy type event app = APIGatewayHttpResolver() - cookie = "CookieMonster" + cookie = Cookie(name="CookieMonster", value="MonsterCookie") @app.post("/my/path") def my_path() -> Response: @@ -172,7 +173,7 @@ def my_path() -> Response: # AND set the current_event type as APIGatewayProxyEventV2 assert result["statusCode"] == 200 assert result["headers"]["Content-Type"] == content_types.TEXT_PLAIN - assert result["cookies"] == [cookie] + assert result["cookies"] == ["CookieMonster=MonsterCookie; Secure"] def test_include_rule_matching(): diff --git a/tests/functional/event_handler/test_lambda_function_url.py b/tests/functional/event_handler/test_lambda_function_url.py index ae0a231d46b..c87d0ecb854 100644 --- a/tests/functional/event_handler/test_lambda_function_url.py +++ b/tests/functional/event_handler/test_lambda_function_url.py @@ -1,4 +1,5 @@ from aws_lambda_powertools.event_handler import LambdaFunctionUrlResolver, Response, content_types +from aws_lambda_powertools.shared.cookies import Cookie from aws_lambda_powertools.utilities.data_classes import LambdaFunctionUrlEvent from tests.functional.utils import load_event @@ -28,7 +29,7 @@ def foo(): def test_lambda_function_url_event_with_cookies(): # GIVEN a Lambda Function Url type event app = LambdaFunctionUrlResolver() - cookie = "CookieMonster" + cookie = Cookie(name="CookieMonster", value="MonsterCookie") @app.get("/") def foo(): @@ -42,7 +43,7 @@ def foo(): # THEN process event correctly # AND set the current_event type as LambdaFunctionUrlEvent assert result["statusCode"] == 200 - assert result["cookies"] == [cookie] + assert result["cookies"] == ["CookieMonster=MonsterCookie; Secure"] def test_lambda_function_url_no_matches():