test: Fix tests

PolicyEngine · Dec 18, 2024 · 0d931a3 · 0d931a3
1 parent adbfaa9
commit 0d931a3
Show file tree

Hide file tree

Showing 3 changed files with 108 additions and 127 deletions.
diff --git a/policyengine_api/api.py b/policyengine_api/api.py
@@ -13,7 +13,7 @@
 # from werkzeug.middleware.profiler import ProfilerMiddleware
 
 # Endpoints
-from policyengine_api.routes.error_routes import ErrorRoutes
+from policyengine_api.routes.error_routes import error_bp
 from policyengine_api.routes.economy_routes import economy_bp
 from policyengine_api.routes.household_routes import household_bp
 from policyengine_api.routes.simulation_analysis_routes import (
@@ -55,7 +55,7 @@
 
 CORS(app)
 
-ErrorRoutes.init_app(app)
+app.register_blueprint(error_bp)
 
 app.route("/", methods=["GET"])(get_home)
 

diff --git a/policyengine_api/routes/error_routes.py b/policyengine_api/routes/error_routes.py
@@ -1,139 +1,120 @@
 import json
-from flask import Response, Flask
+from flask import Response, Blueprint
 from werkzeug.exceptions import (
     HTTPException,
     NotFound,
-    BadRequest,
-    Unauthorized,
-    Forbidden,
-    InternalServerError,
 )
 
+error_bp = Blueprint("error", __name__)
 
-class ErrorRoutes:
-    """
-    Error routing class
-    """
 
-    @staticmethod
-    def init_app(app: Flask) -> None:
-        """
-        Register all error handlers with the Flask app
+@error_bp.app_errorhandler(404)
+def handle_404(error) -> Response:
+    """Specific handler for 404 Not Found errors"""
+    return Response(
+        json.dumps(
+            {
+                "status": "error",
+                "message": str(error),
+                "result": None,
+            }
+        ),
+        404,
+        mimetype="application/json",
+    )
 
-        Args:
-            app (Flask): The Flask app to register error handlers
-        """
-        app.register_error_handler(404, ErrorRoutes.handle_404)
-        app.register_error_handler(400, ErrorRoutes.handle_400)
-        app.register_error_handler(401, ErrorRoutes.handle_401)
-        app.register_error_handler(403, ErrorRoutes.handle_403)
-        app.register_error_handler(500, ErrorRoutes.handle_500)
-        app.register_error_handler(
-            HTTPException, ErrorRoutes.handle_http_exception
-        )
-        app.register_error_handler(Exception, ErrorRoutes.handle_generic_error)
 
-    @staticmethod
-    def handle_http_exception(error: HTTPException) -> Response:
-        """Generic handler for HTTPException; should be raised if no specific handler is found"""
-        return Response(
-            json.dumps(
-                {
-                    "status": "error",
-                    "message": error.description,
-                    "result": None,
-                }
-            ),
-            error.code,
-            mimetype="application/json",
-        )
+@error_bp.app_errorhandler(400)
+def handle_400(error) -> Response:
+    """Specific handler for 400 Bad Request errors"""
+    return Response(
+        json.dumps(
+            {
+                "status": "error",
+                "message": str(error),
+                "result": None,
+            }
+        ),
+        400,
+        mimetype="application/json",
+    )
 
-    @staticmethod
-    def handle_404(error: NotFound) -> Response:
-        """Specific handler for 404 Not Found errors"""
-        return Response(
-            json.dumps(
-                {
-                    "status": "error",
-                    "message": str(error),
-                    "result": None,
-                }
-            ),
-            404,
-            mimetype="application/json",
-        )
 
-    @staticmethod
-    def handle_400(error: BadRequest) -> Response:
-        """Specific handler for 400 Bad Request errors"""
-        return Response(
-            json.dumps(
-                {
-                    "status": "error",
-                    "message": str(error),
-                    "result": None,
-                }
-            ),
-            400,
-            mimetype="application/json",
-        )
+@error_bp.app_errorhandler(401)
+def handle_401(error) -> Response:
+    """Specific handler for 401 Unauthorized errors"""
+    return Response(
+        json.dumps(
+            {
+                "status": "error",
+                "message": str(error),
+                "result": None,
+            }
+        ),
+        401,
+        mimetype="application/json",
+    )
 
-    @staticmethod
-    def handle_401(error: Unauthorized) -> Response:
-        """Specific handler for 401 Unauthorized errors"""
-        return Response(
-            json.dumps(
-                {
-                    "status": "error",
-                    "message": str(error),
-                    "result": None,
-                }
-            ),
-            401,
-            mimetype="application/json",
-        )
 
-    @staticmethod
-    def handle_403(error: Forbidden) -> Response:
-        """Specific handler for 403 Forbidden errors"""
-        return Response(
-            json.dumps(
-                {
-                    "status": "error",
-                    "message": str(error),
-                    "result": None,
-                }
-            ),
-            403,
-            mimetype="application/json",
-        )
+@error_bp.app_errorhandler(403)
+def handle_403(error) -> Response:
+    """Specific handler for 403 Forbidden errors"""
+    return Response(
+        json.dumps(
+            {
+                "status": "error",
+                "message": str(error),
+                "result": None,
+            }
+        ),
+        403,
+        mimetype="application/json",
+    )
 
-    @staticmethod
-    def handle_500(error: InternalServerError) -> Response:
-        """Specific handler for 500 Internal Server Error"""
-        return Response(
-            json.dumps(
-                {
-                    "status": "error",
-                    "message": str(error),
-                    "result": None,
-                }
-            ),
-            500,
-            mimetype="application/json",
-        )
 
-    @staticmethod
-    def handle_generic_error(error: Exception) -> Response:
-        """Handler for any unhandled exceptions"""
-        return Response(
-            json.dumps(
-                {
-                    "status": "error",
-                    "message": str(error),
-                    "result": None,
-                }
-            ),
-            500,
-            mimetype="application/json",
-        )
+@error_bp.app_errorhandler(500)
+def handle_500(error) -> Response:
+    """Specific handler for 500 Internal Server Error"""
+    return Response(
+        json.dumps(
+            {
+                "status": "error",
+                "message": str(error),
+                "result": None,
+            }
+        ),
+        500,
+        mimetype="application/json",
+    )
+
+
+@error_bp.app_errorhandler(HTTPException)
+def handle_http_exception(error: HTTPException) -> Response:
+    """Generic handler for HTTPException; should be raised if no specific handler is found"""
+    return Response(
+        json.dumps(
+            {
+                "status": "error",
+                "message": error.description,
+                "result": None,
+            }
+        ),
+        error.code,
+        mimetype="application/json",
+    )
+
+
+@error_bp.app_errorhandler(Exception)
+def handle_generic_error(error: Exception) -> Response:
+    """Handler for any unhandled exceptions"""
+    return Response(
+        json.dumps(
+            {
+                "status": "error",
+                "message": str(error),
+                "result": None,
+            }
+        ),
+        500,
+        mimetype="application/json",
+    )
diff --git a/tests/python/test_errors.py b/tests/python/test_errors.py
@@ -1,20 +1,20 @@
 import pytest
 from flask import Flask
+from policyengine_api.routes.error_routes import error_bp
 from werkzeug.exceptions import (
     NotFound,
     BadRequest,
     Unauthorized,
     Forbidden,
     InternalServerError,
 )
-from policyengine_api.routes.error_routes import ErrorRoutes
 
 
 @pytest.fixture
 def app():
     """Create and configure a new app instance for each test."""
     app = Flask(__name__)
-    ErrorRoutes.init_app(app)
+    app.register_blueprint(error_bp)
     return app