refactor: add status to the response

Author: Michael Brewer

aws_lambda_powertools/event_handler/api_gateway.py

@@ -18,72 +18,44 @@
 class ServiceError(Exception):
     """Service Error"""
 
-    def __init__(self, status_code: int, message: str):
+    def __init__(self, status_code: int, msg: str):
         """
         Parameters
         ----------
-        code: int
+        status_code: int
             Http status code
-        message: str
+        msg: str
             Error message
         """
         self.status_code = status_code
-        self.message = message
-
-    def __str__(self) -> str:
-        """To string of the message only"""
-        return self.message
+        self.msg = msg
 
 
 class BadRequestError(ServiceError):
     """Bad Request Error"""
 
-    def __init__(self, message: str):
-        """
-        Parameters
-        ----------
-        message: str
-            Error message
-        """
-        super().__init__(400, message)
+    def __init__(self, msg: str):
+        super().__init__(400, msg)
 
 
 class UnauthorizedError(ServiceError):
     """Unauthorized Error"""
 
-    def __init__(self, message: str):
-        """
-        Parameters
-        ----------
-        message: str
-            Error message
-        """
-        super().__init__(401, message)
+    def __init__(self, msg: str):
+        super().__init__(401, msg)
 
 
 class NotFoundError(ServiceError):
     """Not Found Error"""
 
-    def __init__(self, message: str = "Not found"):
-        """
-        Parameters
-        ----------
-        message: str
-            Error message
-        """
-        super().__init__(404, message)
+    def __init__(self, msg: str = "Not found"):
+        super().__init__(404, msg)
 
 
 class InternalServerError(ServiceError):
-    """Internal Serve Error"""
+    """Internal Server Error"""
 
     def __init__(self, message: str):
-        """
-        Parameters
-        ----------
-        message: str
-            Error message
-        """
         super().__init__(500, message)
 
 
@@ -542,7 +514,7 @@ def _not_found(self, method: str) -> ResponseBuilder:
                 status_code=404,
                 content_type=APPLICATION_JSON,
                 headers=headers,
-                body=json.dumps({"message": "Not found"}),
+                body=self._json_dump({"code": 404, "message": "Not found"}),
             )
         )
 
@@ -555,13 +527,12 @@ def _call_route(self, route: Route, args: Dict[str, str]) -> ResponseBuilder:
                 Response(
                     status_code=e.status_code,
                     content_type=APPLICATION_JSON,
-                    body=json.dumps({"message": str(e)}),
+                    body=self._json_dump({"code": e.status_code, "message": e.msg}),
                 ),
                 route,
             )
 
-    @staticmethod
-    def _to_response(result: Union[Dict, Response]) -> Response:
+    def _to_response(self, result: Union[Dict, Response]) -> Response:
         """Convert the route's result to a Response
 
          2 main result types are supported:
@@ -577,5 +548,10 @@ def _to_response(result: Union[Dict, Response]) -> Response:
         return Response(
             status_code=200,
             content_type=APPLICATION_JSON,
-            body=json.dumps(result, separators=(",", ":"), cls=Encoder),
+            body=self._json_dump(result),
         )
+
+    @staticmethod
+    def _json_dump(obj: Any) -> str:
+        """Does a concise json serialization"""
+        return json.dumps(obj, separators=(",", ":"), cls=Encoder)

tests/functional/event_handler/test_api_gateway.py

@@ -498,72 +498,85 @@ def custom_method():
     assert headers["Access-Control-Allow-Methods"] == "CUSTOM"
 
 
-def test_service_error_response():
-    # GIVEN a service error response
+def test_service_error_responses():
+    # SCENARIO handling different kind of service errors being raised
     app = ApiGatewayResolver(cors=CORSConfig())
 
-    @app.route(method="GET", rule="/bad-request-error", cors=False)
+    def json_dump(obj):
+        return json.dumps(obj, separators=(",", ":"))
+
+    # GIVEN an BadRequestError
+    @app.get(rule="/bad-request-error", cors=False)
     def bad_request_error():
         raise BadRequestError("Missing required parameter")
 
-    @app.route(method="GET", rule="/unauthorized-error", cors=False)
-    def unauthorized_error():
-        raise UnauthorizedError("Unauthorized")
-
-    @app.route(method="GET", rule="/service-error", cors=True)
-    def service_error():
-        raise ServiceError(403, "Unauthorized")
-
-    @app.route(method="GET", rule="/not-found-error", cors=False)
-    def not_found_error():
-        raise NotFoundError
-
-    @app.route(method="GET", rule="/internal-server-error", cors=False)
-    def internal_server_error():
-        raise InternalServerError("Internal server error")
-
     # WHEN calling the handler
     # AND path is /bad-request-error
     result = app({"path": "/bad-request-error", "httpMethod": "GET"}, None)
     # THEN return the bad request error response
     # AND status code equals 400
     assert result["statusCode"] == 400
-    assert result["body"] == json.dumps({"message": "Missing required parameter"})
     assert result["headers"]["Content-Type"] == APPLICATION_JSON
+    expected = {"code": 400, "message": "Missing required parameter"}
+    assert result["body"] == json_dump(expected)
+
+    # GIVEN an UnauthorizedError
+    @app.get(rule="/unauthorized-error", cors=False)
+    def unauthorized_error():
+        raise UnauthorizedError("Unauthorized")
 
     # WHEN calling the handler
     # AND path is /unauthorized-error
     result = app({"path": "/unauthorized-error", "httpMethod": "GET"}, None)
     # THEN return the unauthorized error response
     # AND status code equals 401
     assert result["statusCode"] == 401
-    assert result["body"] == json.dumps({"message": "Unauthorized"})
     assert result["headers"]["Content-Type"] == APPLICATION_JSON
+    expected = {"code": 401, "message": "Unauthorized"}
+    assert result["body"] == json_dump(expected)
 
-    # WHEN calling the handler
-    # AND path is /service-error
-    result = app({"path": "/service-error", "httpMethod": "GET"}, None)
-    # THEN return the service error response
-    # AND status code equals 403
-    assert result["statusCode"] == 403
-    assert result["body"] == json.dumps({"message": "Unauthorized"})
-    assert result["headers"]["Content-Type"] == APPLICATION_JSON
-    assert "Access-Control-Allow-Origin" in result["headers"]
+    # GIVEN an NotFoundError
+    @app.get(rule="/not-found-error", cors=False)
+    def not_found_error():
+        raise NotFoundError
 
     # WHEN calling the handler
     # AND path is /not-found-error
     result = app({"path": "/not-found-error", "httpMethod": "GET"}, None)
     # THEN return the not found error response
     # AND status code equals 404
     assert result["statusCode"] == 404
-    assert result["body"] == json.dumps({"message": "Not found"})
     assert result["headers"]["Content-Type"] == APPLICATION_JSON
+    expected = {"code": 404, "message": "Not found"}
+    assert result["body"] == json_dump(expected)
+
+    # GIVEN an InternalServerError
+    @app.get(rule="/internal-server-error", cors=False)
+    def internal_server_error():
+        raise InternalServerError("Internal server error")
 
     # WHEN calling the handler
     # AND path is /internal-server-error
     result = app({"path": "/internal-server-error", "httpMethod": "GET"}, None)
     # THEN return the internal server error response
     # AND status code equals 500
     assert result["statusCode"] == 500
-    assert result["body"] == json.dumps({"message": "Internal server error"})
     assert result["headers"]["Content-Type"] == APPLICATION_JSON
+    expected = {"code": 500, "message": "Internal server error"}
+    assert result["body"] == json_dump(expected)
+
+    # GIVEN an ServiceError with a custom status code
+    @app.get(rule="/service-error", cors=True)
+    def service_error():
+        raise ServiceError(502, "Something went wrong!")
+
+    # WHEN calling the handler
+    # AND path is /service-error
+    result = app({"path": "/service-error", "httpMethod": "GET"}, None)
+    # THEN return the service error response
+    # AND status code equals 502
+    assert result["statusCode"] == 502
+    assert result["headers"]["Content-Type"] == APPLICATION_JSON
+    assert "Access-Control-Allow-Origin" in result["headers"]
+    expected = {"code": 502, "message": "Something went wrong!"}
+    assert result["body"] == json_dump(expected)