feat: Basic Werkzeug route validators

PolicyEngine · Dec 12, 2024 · f789559 · f789559
1 parent c0fdde9
commit f789559
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 16 deletions.
diff --git a/policyengine_api/routes/economy_routes.py b/policyengine_api/routes/economy_routes.py
@@ -12,7 +12,7 @@
 
 @validate_country
 @economy_bp.route(
-    "/<country_id>/economy/<policy_id>/over/<baseline_policy_id>",
+    "/<country_id>/economy/<int:policy_id>/over/<int:baseline_policy_id>",
     methods=["GET"],
 )
 def get_economic_impact(country_id, policy_id, baseline_policy_id):

diff --git a/policyengine_api/routes/household_routes.py b/policyengine_api/routes/household_routes.py
@@ -14,27 +14,18 @@
 household_service = HouseholdService()
 
 
-@household_bp.route("/<country_id>/household/<household_id>", methods=["GET"])
+@household_bp.route("/<country_id>/household/<int:household_id>", methods=["GET"])
 @validate_country
-def get_household(country_id: str, household_id: str) -> Response:
+def get_household(country_id: str, household_id: int) -> Response:
     """
     Get a household's input data with a given ID.
 
     Args:
         country_id (str): The country ID.
-        household_id (str): The household ID.
+        household_id (int): The household ID.
     """
     print(f"Got request for household {household_id} in country {country_id}")
 
-    # Ensure that household ID is a number
-    try:
-        household_id = int(household_id)
-    except ValueError:
-        return Response(
-            status=400,
-            response=f"Invalid household ID; household ID must be a number",
-        )
-
     try:
         household: dict | None = household_service.get_household(
             country_id, household_id
@@ -128,15 +119,15 @@ def post_household(country_id: str) -> Response:
         )
 
 
-@household_bp.route("/<country_id>/household/<household_id>", methods=["PUT"])
+@household_bp.route("/<country_id>/household/<int:household_id>", methods=["PUT"])
 @validate_country
-def update_household(country_id: str, household_id: str) -> Response:
+def update_household(country_id: str, household_id: int) -> Response:
     """
     Update a household's input data.
 
     Args:
         country_id (str): The country ID.
-        household_id (str): The household ID.
+        household_id (int): The household ID.
     """
 
     # Validate payload