feat: add route-specific custom response validation and tests

Author: amin-farjadi

aws_lambda_powertools/event_handler/api_gateway.py

@@ -320,6 +320,7 @@ def __init__(
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         middlewares: list[Callable[..., Response]] | None = None,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
     ):
         """
         Internally used Route Configuration
@@ -362,6 +363,7 @@ def __init__(
             Whether or not to mark this route as deprecated in the OpenAPI schema
         middlewares: list[Callable[..., Response]] | None
             The list of route middlewares to be called in order.
+        # TODO
         """
         self.method = method.upper()
         self.path = "/" if path.strip() == "" else path
@@ -397,6 +399,8 @@ def __init__(
         # _body_field is used to cache the dependant model for the body field
         self._body_field: ModelField | None = None
 
+        self.custom_response_validation_http_code: int | HTTPStatus | None = custom_response_validation_http_code
+
     def __call__(
         self,
         router_middlewares: list[Callable],
@@ -565,6 +569,8 @@ def _get_openapi_path(
             },
         }
 
+        # TODO update responses
+
         # Add the response to the OpenAPI operation
         if self.responses:
             for status_code in list(self.responses):
@@ -943,6 +949,7 @@ def route(
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         middlewares: list[Callable[..., Any]] | None = None,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         raise NotImplementedError()
 
@@ -1004,6 +1011,7 @@ def get(
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         middlewares: list[Callable[..., Any]] | None = None,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Get route decorator with GET `method`
 
@@ -1044,6 +1052,7 @@ def lambda_handler(event, context):
             openapi_extensions,
             deprecated,
             middlewares,
+            custom_response_validation_http_code,
         )
 
     def post(
@@ -1063,6 +1072,7 @@ def post(
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         middlewares: list[Callable[..., Any]] | None = None,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Post route decorator with POST `method`
 
@@ -1104,6 +1114,7 @@ def lambda_handler(event, context):
             openapi_extensions,
             deprecated,
             middlewares,
+            custom_response_validation_http_code,
         )
 
     def put(
@@ -1123,6 +1134,7 @@ def put(
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         middlewares: list[Callable[..., Any]] | None = None,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Put route decorator with PUT `method`
 
@@ -1164,6 +1176,7 @@ def lambda_handler(event, context):
             openapi_extensions,
             deprecated,
             middlewares,
+            custom_response_validation_http_code,
         )
 
     def delete(
@@ -1183,6 +1196,7 @@ def delete(
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         middlewares: list[Callable[..., Any]] | None = None,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Delete route decorator with DELETE `method`
 
@@ -1223,6 +1237,7 @@ def lambda_handler(event, context):
             openapi_extensions,
             deprecated,
             middlewares,
+            custom_response_validation_http_code,
         )
 
     def patch(
@@ -1242,6 +1257,7 @@ def patch(
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         middlewares: list[Callable] | None = None,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Patch route decorator with PATCH `method`
 
@@ -1285,6 +1301,7 @@ def lambda_handler(event, context):
             openapi_extensions,
             deprecated,
             middlewares,
+            custom_response_validation_http_code,
         )
 
     def head(
@@ -1304,6 +1321,7 @@ def head(
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         middlewares: list[Callable] | None = None,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Head route decorator with HEAD `method`
 
@@ -1346,6 +1364,7 @@ def lambda_handler(event, context):
             openapi_extensions,
             deprecated,
             middlewares,
+            custom_response_validation_http_code,
         )
 
     def _push_processed_stack_frame(self, frame: str):
@@ -2126,9 +2145,14 @@ def route(
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         middlewares: list[Callable[..., Any]] | None = None,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         """Route decorator includes parameter `method`"""
 
+        custom_response_validation_http_code = self._validate_route_response_validation_error_http_code(
+            custom_response_validation_http_code,
+        )
+
         def register_resolver(func: AnyCallableT) -> AnyCallableT:
             methods = (method,) if isinstance(method, str) else method
             logger.debug(f"Adding route using rule {rule} and methods: {','.join(m.upper() for m in methods)}")
@@ -2155,6 +2179,7 @@ def register_resolver(func: AnyCallableT) -> AnyCallableT:
                     openapi_extensions,
                     deprecated,
                     middlewares,
+                    custom_response_validation_http_code,
                 )
 
                 # The more specific route wins.
@@ -2523,15 +2548,20 @@ def _call_exception_handler(self, exp: Exception, route: Route) -> ResponseBuild
             )
 
         # OpenAPIValidationMiddleware will only raise ResponseValidationError when
-        # 'self._response_validation_error_http_code' is not None
+        # 'self._response_validation_error_http_code' is not None or
+        # when route has custom_response_validation_http_code
         if isinstance(exp, ResponseValidationError):
-            http_code = self._response_validation_error_http_code
+            http_code = (
+                self._response_validation_error_http_code
+                if exp.source == "app"
+                else route.custom_response_validation_http_code
+            )
             errors = [{"loc": e["loc"], "type": e["type"]} for e in exp.errors()]
             return self._response_builder_class(
                 response=Response(
                     status_code=http_code.value,
                     content_type=content_types.APPLICATION_JSON,
-                    body={"statusCode": self._response_validation_error_http_code, "detail": errors},
+                    body={"statusCode": http_code, "detail": errors},
                 ),
                 serializer=self._serializer,
                 route=route,
@@ -2683,6 +2713,7 @@ def route(
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         middlewares: list[Callable[..., Any]] | None = None,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         def register_route(func: AnyCallableT) -> AnyCallableT:
             # All dict keys needs to be hashable. So we'll need to do some conversions:
@@ -2708,6 +2739,7 @@ def register_route(func: AnyCallableT) -> AnyCallableT:
                 frozen_security,
                 frozen_openapi_extensions,
                 deprecated,
+                custom_response_validation_http_code,
             )
 
             # Collate Middleware for routes
@@ -2795,6 +2827,7 @@ def route(
         openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         middlewares: list[Callable[..., Any]] | None = None,
+        custom_response_validation_http_code: int | HTTPStatus | None = None,
     ) -> Callable[[AnyCallableT], AnyCallableT]:
         # NOTE: see #1552 for more context.
         return super().route(
@@ -2814,6 +2847,7 @@ def route(
             openapi_extensions,
             deprecated,
             middlewares,
+            custom_response_validation_http_code,
         )
 
     # Override _compile_regex to exclude trailing slashes for route resolution

aws_lambda_powertools/event_handler/middlewares/openapi_validation.py

@@ -150,6 +150,7 @@ def _handle_response(self, *, route: Route, response: Response):
             response.body = self._serialize_response(
                 field=route.dependant.return_param,
                 response_content=response.body,
+                has_route_custom_response_validation=route.custom_response_validation_http_code is not None,
             )
 
         return response
@@ -165,6 +166,7 @@ def _serialize_response(
         exclude_unset: bool = False,
         exclude_defaults: bool = False,
         exclude_none: bool = False,
+        has_route_custom_response_validation: bool = False,
     ) -> Any:
         """
         Serialize the response content according to the field type.
@@ -174,7 +176,13 @@ def _serialize_response(
             value = _validate_field(field=field, value=response_content, loc=("response",), existing_errors=errors)
             if errors:
                 if self._has_response_validation_error:
-                    raise ResponseValidationError(errors=_normalize_errors(errors), body=response_content)
+                    raise ResponseValidationError(errors=_normalize_errors(errors), body=response_content, source="app")
+                if has_route_custom_response_validation:
+                    raise ResponseValidationError(
+                        errors=_normalize_errors(errors),
+                        body=response_content,
+                        source="route",
+                    )
                 raise RequestValidationError(errors=_normalize_errors(errors), body=response_content)
 
             if hasattr(field, "serialize"):

aws_lambda_powertools/event_handler/openapi/exceptions.py

@@ -1,4 +1,4 @@
-from typing import Any, Sequence
+from typing import Any, Literal, Sequence
 
 
 class ValidationException(Exception):
@@ -28,9 +28,10 @@ class ResponseValidationError(ValidationException):
     Raised when the response body does not match the OpenAPI schema
     """
 
-    def __init__(self, errors: Sequence[Any], *, body: Any = None) -> None:
+    def __init__(self, errors: Sequence[Any], *, body: Any = None, source: Literal["route", "app"] = "app") -> None:
         super().__init__(errors)
         self.body = body
+        self.source = source
 
 
 class SerializationError(Exception):

tests/functional/event_handler/_pydantic/test_openapi_validation_middleware.py

@@ -1378,3 +1378,67 @@ def test_custom_response_validation_error_bad_http_code(response_validation_erro
         str(exception_info.value)
         == f"'{response_validation_error_http_code}' must be an integer representing an HTTP status code."
     )
+
+
+def test_custom_route_response_validation_error_http_code_invalid_response_incomplete_model(gw_event):
+    # GIVEN an APIGatewayRestResolver with custom response validation enabled
+    app = APIGatewayRestResolver(enable_validation=True)
+
+    class Model(BaseModel):
+        name: str
+        age: int
+
+    @app.get("/incomplete_model_not_allowed")
+    def handler_incomplete_model_not_allowed() -> Model:
+        return {"age": 18}  # type: ignore
+
+    @app.get(
+        "/custom_incomplete_model_not_allowed",
+        custom_response_validation_http_code=500,
+    )
+    def handler_custom_route_response_validation_error() -> Model:
+        return {"age": 18}  # type: ignore
+
+    # WHEN returning incomplete model for a non-Optional type
+    gw_event["path"] = "/incomplete_model_not_allowed"
+    result = app(gw_event, {})
+
+    gw_event["path"] = "/custom_incomplete_model_not_allowed"
+    custom_result = app(gw_event, {})
+
+    # THEN it should return a validation error with the custom status code provided
+    assert result["statusCode"] == 422
+    assert custom_result["statusCode"] == 500
+    assert json.loads(result["body"])["detail"] == json.loads(custom_result["body"])["detail"]
+
+
+def test_custom_route_response_validation_error_sanitized_response(gw_event):
+    # GIVEN an APIGatewayRestResolver with custom response validation enabled
+    # with a sanitized response validation error response
+    app = APIGatewayRestResolver(enable_validation=True)
+
+    class Model(BaseModel):
+        name: str
+        age: int
+
+    @app.get(
+        "/custom_incomplete_model_not_allowed",
+        custom_response_validation_http_code=422,
+    )
+    def handler_custom_route_response_validation_error() -> Model:
+        return {"age": 18}  # type: ignore
+
+    @app.exception_handler(ResponseValidationError)
+    def handle_response_validation_error(ex: ResponseValidationError):
+        return Response(
+            status_code=500,
+            body="Unexpected response.",
+        )
+
+    # WHEN returning incomplete model for a non-Optional type
+    gw_event["path"] = "/custom_incomplete_model_not_allowed"
+    result = app(gw_event, {})
+
+    # THEN it should return the sanitized response
+    assert result["statusCode"] == 500
+    assert result["body"] == "Unexpected response."