diff --git a/aws_lambda_powertools/event_handler/api_gateway.py b/aws_lambda_powertools/event_handler/api_gateway.py index e6d1af01dfc..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): @@ -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,12 +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() - return { + 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): + resp["cookies"] = self.response.cookies + else: + 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 1ca28f869cf..1b5a43c020e 100644 --- a/tests/functional/event_handler/test_api_gateway.py +++ b/tests/functional/event_handler/test_api_gateway.py @@ -1258,3 +1258,42 @@ 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"] + + +@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 calling handler + # AND setting Response 'cookies' + result = app(LOAD_GW_EVENT, {}) + + # THEN include 'Set-Cookie' in the result + assert result["headers"]["Content-Type"] == content_types.TEXT_PLAIN + assert result["multiValueHeaders"]["Set-Cookie"] == ["foo=bar"]