From d4f4376fc32af62ae0aab4c8e21e1dbc5ab966c1 Mon Sep 17 00:00:00 2001 From: Michael Brewer Date: Fri, 29 Apr 2022 18:28:53 -0700 Subject: [PATCH 1/2] feat(event_handler): Allow cookies in response Changes: - Add cookies field in Response - For http api gateway format 2.0 responses serialize to cookies - For alb and rest api use multiValueHeaders closes #1192 --- .../event_handler/api_gateway.py | 15 ++++++++- .../event_handler/test_api_gateway.py | 32 +++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/aws_lambda_powertools/event_handler/api_gateway.py b/aws_lambda_powertools/event_handler/api_gateway.py index e6d1af01dfc..1b529a204c7 100644 --- a/aws_lambda_powertools/event_handler/api_gateway.py +++ b/aws_lambda_powertools/event_handler/api_gateway.py @@ -140,6 +140,7 @@ def __init__( content_type: Optional[str], body: Union[str, bytes, None], headers: Optional[Dict] = None, + cookies: Optional[List[str]] = None, ): """ @@ -154,11 +155,14 @@ def __init__( Optionally set the response body. Note: bytes body will be automatically base64 encoded headers: dict Optionally set specific http headers. Setting "Content-Type" hear would override the `content_type` value. + cookies: Optional[List[str]] + Optionally set cookies """ self.status_code = status_code self.body = body self.base64_encoded = False self.headers: Dict = headers or {} + self.cookies = cookies if content_type: self.headers.setdefault("Content-Type", content_type) @@ -220,13 +224,22 @@ def build(self, event: BaseProxyEvent, cors: Optional[CORSConfig] = None) -> Dic logger.debug("Encoding bytes response with base64") self.response.base64_encoded = True self.response.body = base64.b64encode(self.response.body).decode() - return { + + response = { "statusCode": self.response.status_code, "headers": self.response.headers, "body": self.response.body, "isBase64Encoded": self.response.base64_encoded, } + if self.response.cookies: + if isinstance(event, APIGatewayProxyEventV2): + response["cookies"] = self.response.cookies + else: + response["multiValueHeaders"] = {"Set-Cookie": self.response.cookies} + + return response + class BaseRouter(ABC): current_event: BaseProxyEvent diff --git a/tests/functional/event_handler/test_api_gateway.py b/tests/functional/event_handler/test_api_gateway.py index 1ca28f869cf..9f85336c43b 100644 --- a/tests/functional/event_handler/test_api_gateway.py +++ b/tests/functional/event_handler/test_api_gateway.py @@ -1258,3 +1258,35 @@ def handler(event: APIGatewayProxyEventV2, context): # THEN result = handler(load_event("apiGatewayProxyV2Event.json"), None) assert result["statusCode"] == 200 + + +def test_api_gateway_http_cookie(): + # GIVEN + app = APIGatewayHttpResolver() + + @app.post("/my/path") + def my_path() -> Response: + return Response(200, content_types.TEXT_PLAIN, "Test", cookies=["key=value"]) + + # WHEN + result = app(load_event("apiGatewayProxyV2Event.json"), {}) + + # THEN + assert result["headers"]["Content-Type"] == content_types.TEXT_PLAIN + assert result["cookies"] == ["key=value"] + + +def test_api_gateway_rest_cookie(): + # GIVEN + app = APIGatewayRestResolver() + + @app.get("/my/path") + def get_lambda() -> Response: + return Response(200, content_types.TEXT_PLAIN, "Test", cookies=["foo=bar"]) + + # WHEN + result = app(LOAD_GW_EVENT, {}) + + # THEN + assert result["headers"]["Content-Type"] == content_types.TEXT_PLAIN + assert result["multiValueHeaders"]["Set-Cookie"] == ["foo=bar"] From 7ffeb63af8bc5be811e69a5ce22973988842f00f Mon Sep 17 00:00:00 2001 From: Michael Brewer Date: Sat, 30 Apr 2022 01:06:08 -0700 Subject: [PATCH 2/2] test: Add more coverage --- .../event_handler/api_gateway.py | 13 +++++-------- .../event_handler/test_api_gateway.py | 17 ++++++++++++----- 2 files changed, 17 insertions(+), 13 deletions(-) diff --git a/aws_lambda_powertools/event_handler/api_gateway.py b/aws_lambda_powertools/event_handler/api_gateway.py index 1b529a204c7..3ce6e785201 100644 --- a/aws_lambda_powertools/event_handler/api_gateway.py +++ b/aws_lambda_powertools/event_handler/api_gateway.py @@ -27,7 +27,7 @@ _SAFE_URI = "-._~()'!*:@,;" # https://www.ietf.org/rfc/rfc3986.txt # API GW/ALB decode non-safe URI chars; we must support them too _UNSAFE_URI = "%<> \[\]{}|^" # noqa: W605 -_NAMED_GROUP_BOUNDARY_PATTERN = fr"(?P\1[{_SAFE_URI}{_UNSAFE_URI}\\w]+)" +_NAMED_GROUP_BOUNDARY_PATTERN = rf"(?P\1[{_SAFE_URI}{_UNSAFE_URI}\\w]+)" class ProxyEventType(Enum): @@ -224,21 +224,18 @@ def build(self, event: BaseProxyEvent, cors: Optional[CORSConfig] = None) -> Dic logger.debug("Encoding bytes response with base64") self.response.base64_encoded = True self.response.body = base64.b64encode(self.response.body).decode() - - response = { + resp = { "statusCode": self.response.status_code, "headers": self.response.headers, "body": self.response.body, "isBase64Encoded": self.response.base64_encoded, } - if self.response.cookies: if isinstance(event, APIGatewayProxyEventV2): - response["cookies"] = self.response.cookies + resp["cookies"] = self.response.cookies else: - response["multiValueHeaders"] = {"Set-Cookie": self.response.cookies} - - return response + resp["multiValueHeaders"] = {"Set-Cookie": self.response.cookies} + return resp class BaseRouter(ABC): diff --git a/tests/functional/event_handler/test_api_gateway.py b/tests/functional/event_handler/test_api_gateway.py index 9f85336c43b..1b5a43c020e 100644 --- a/tests/functional/event_handler/test_api_gateway.py +++ b/tests/functional/event_handler/test_api_gateway.py @@ -1276,17 +1276,24 @@ def my_path() -> Response: assert result["cookies"] == ["key=value"] -def test_api_gateway_rest_cookie(): - # GIVEN - app = APIGatewayRestResolver() +@pytest.mark.parametrize( + "app", + [ + pytest.param(APIGatewayRestResolver(), id="rest handler"), + pytest.param(ALBResolver(), id="alb handler"), + ], +) +def test_api_gateway_rest_cookie(app: ApiGatewayResolver): + # GIVEN APIGatewayRestResolver or ALBResolver @app.get("/my/path") def get_lambda() -> Response: return Response(200, content_types.TEXT_PLAIN, "Test", cookies=["foo=bar"]) - # WHEN + # WHEN calling handler + # AND setting Response 'cookies' result = app(LOAD_GW_EVENT, {}) - # THEN + # THEN include 'Set-Cookie' in the result assert result["headers"]["Content-Type"] == content_types.TEXT_PLAIN assert result["multiValueHeaders"]["Set-Cookie"] == ["foo=bar"]