fix(event_handler): keep catch-all exception handlers off 404s and adapt exception handler returns

aws_lambda_powertools/event_handler/api_gateway.py

@@ -2765,7 +2765,7 @@ def not_found_handler():
                 _headers["Access-Control-Allow-Methods"] = CORSConfig.build_allow_methods(self._cors_methods)
                 return Response(status_code=204, content_type=None, headers=_headers, body="")
 
-            custom_not_found_handler = self.exception_handler_manager.lookup_exception_handler(NotFoundError)
+            custom_not_found_handler = self._lookup_not_found_handler()
             if custom_not_found_handler:
                 return custom_not_found_handler(NotFoundError())
 
@@ -2933,7 +2933,7 @@ def not_found_handler():
                 return Response(status_code=204, content_type=None, headers=_headers, body="")
 
             # Customer registered 404 route? Call it.
-            custom_not_found_handler = self.exception_handler_manager.lookup_exception_handler(NotFoundError)
+            custom_not_found_handler = self._lookup_not_found_handler()
             if custom_not_found_handler:
                 return custom_not_found_handler(NotFoundError())
 
@@ -3008,6 +3008,21 @@ def not_found(self, func: Callable | None = None):
     def exception_handler(self, exc_class: type[Exception] | list[type[Exception]]):
         return self.exception_handler_manager.exception_handler(exc_class=exc_class)
 
+    def _lookup_not_found_handler(self) -> Callable | None:
+        """Find the handler for an unmatched route, ignoring catch-all handlers.
+
+        No exception was raised for an unmatched route, so a handler registered for `Exception` or `BaseException`
+        (e.g. a generic 500 handler) must not turn it into its own response. Handlers for `NotFoundError` and
+        `ServiceError` still apply.
+        """
+        handlers = self.exception_handler_manager.get_registered_handlers()
+        for cls in NotFoundError.__mro__:
+            if cls in (Exception, BaseException):
+                break
+            if cls in handlers:
+                return handlers[cls]
+        return None
+
     def _call_exception_handler(self, exp: Exception, route: Route) -> ResponseBuilder | None:
         # Depends() wraps what a dependency raised in DependencyResolutionError. Handle the original exception
         # as if the route had raised it, unless a handler is registered for DependencyResolutionError itself.
@@ -3022,7 +3037,12 @@ def _call_exception_handler(self, exp: Exception, route: Route) -> ResponseBuild
         handler = self.exception_handler_manager.lookup_exception_handler(type(exp))
         if handler:
             try:
-                return self._response_builder_class(response=handler(exp), serializer=self._serializer, route=route)
+                # Adapt the handler's return value like a route's (dict/tuple, ALB null body)
+                return self._response_builder_class(
+                    response=self._to_response(handler(exp)),  # type: ignore[arg-type]
+                    serializer=self._serializer,
+                    route=route,
+                )
             except ServiceError as service_error:
                 exp = service_error
 
@@ -3480,7 +3500,6 @@ def _build_response(self, response_builder: ResponseBuilder) -> dict[str, Any]:
             if handled_response is None:
                 raise
 
-            handled_response.response = cast(Response, self._to_response(handled_response.response))
             response = super()._build_response(handled_response)
             # Validate once more without invoking an exception handler recursively.
             self._validate_response_size(response)

tests/functional/event_handler/required_dependencies/test_api_gateway.py

@@ -1785,6 +1785,91 @@ def handle_not_found(_) -> Response:
     assert result["statusCode"] == 404
 
 
+def test_catch_all_exception_handler_does_not_hijack_not_found():
+    # GIVEN a resolver with a catch-all exception handler, as suggested in the docs
+    app = APIGatewayRestResolver()
+
+    @app.exception_handler(Exception)
+    def handle_everything(ex: Exception) -> Response:
+        return Response(status_code=500, content_type=content_types.TEXT_PLAIN, body="generic handler")
+
+    @app.get("/other")
+    def get_other():
+        return {}
+
+    # WHEN calling the event handler for an unmatched route
+    result = app(LOAD_GW_EVENT, {})
+
+    # THEN the default 404 is returned instead of the catch-all response
+    assert result["statusCode"] == 404
+    assert json.loads(result["body"]) == {"statusCode": 404, "message": "Not found"}
+
+
+def test_service_error_handler_still_handles_not_found():
+    # GIVEN a resolver with a handler for ServiceError, which NotFoundError inherits from
+    app = APIGatewayRestResolver()
+
+    @app.exception_handler(ServiceError)
+    def handle_service_error(ex: ServiceError) -> Response:
+        return Response(status_code=ex.status_code, content_type=content_types.TEXT_PLAIN, body="custom format")
+
+    # WHEN calling the event handler for an unmatched route
+    result = app(LOAD_GW_EVENT, {})
+
+    # THEN the ServiceError handler keeps formatting the 404
+    assert result["statusCode"] == 404
+    assert result["body"] == "custom format"
+
+
+def test_exception_handler_none_body_on_alb_is_converted_to_empty_string():
+    # GIVEN an ALB resolver whose exception handler returns a Response with no body
+    app = ALBResolver()
+
+    @app.exception_handler(ValueError)
+    def handle_value_error(ex: ValueError) -> Response:
+        return Response(status_code=400, content_type=content_types.TEXT_PLAIN, body=None)
+
+    @app.get("/lambda")
+    def get_lambda():
+        raise ValueError("Foo!")
+
+    # WHEN calling the event handler
+    result = app(load_event("albEvent.json"), {})
+
+    # THEN the body is an empty string, as it is for a route returning the same Response, since ALB rejects null
+    assert result["statusCode"] == 400
+    assert result["body"] == ""
+
+
+@pytest.mark.parametrize(
+    "handler_result, expected_status",
+    [
+        ({"error": "Foo!"}, 200),
+        (({"error": "Foo!"}, 400), 400),
+    ],
+    ids=["dict", "tuple"],
+)
+def test_exception_handler_dict_and_tuple_returns_are_serialized(handler_result, expected_status):
+    # GIVEN an exception handler that returns a dict or a (dict, status_code) tuple, like a route can
+    app = APIGatewayRestResolver()
+
+    @app.exception_handler(ValueError)
+    def handle_value_error(ex: ValueError):
+        return handler_result
+
+    @app.get("/my/path")
+    def get_lambda():
+        raise ValueError("Foo!")
+
+    # WHEN calling the event handler
+    result = app(LOAD_GW_EVENT, {})
+
+    # THEN the result is serialized the same way a route return value is
+    assert result["statusCode"] == expected_status
+    assert result["multiValueHeaders"]["Content-Type"] == [content_types.APPLICATION_JSON]
+    assert json.loads(result["body"]) == {"error": "Foo!"}
+
+
 def test_exception_handler_raises_service_error(json_dump):
     # GIVEN an exception handler raises a ServiceError (BadRequestError)
     app = ApiGatewayResolver()

tests/functional/event_handler/required_dependencies/test_resolve_async.py

@@ -185,6 +185,22 @@ def get_lambda():
         assert response["statusCode"] == 404
         assert response["body"] == '{"error": "custom 404"}'
 
+    def test_catch_all_exception_handler_does_not_hijack_not_found(self):
+        # GIVEN a catch-all exception handler
+        app = APIGatewayRestResolver()
+
+        @app.exception_handler(Exception)
+        def handle_everything(exc):
+            return Response(500, content_types.TEXT_PLAIN, "generic handler")
+
+        # WHEN resolving with no matching route
+        _setup_app(app, API_REST_EVENT)
+        result = asyncio.run(app._resolve_async())
+
+        # THEN the default 404 is returned
+        response = result.build(app.current_event, app._cors)
+        assert response["statusCode"] == 404
+
     def test_cors_preflight_returns_204(self):
         # GIVEN a resolver with CORS enabled
         app = APIGatewayRestResolver(cors=CORSConfig())