feat(event_handler): Allow for cookies in response

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):

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"]