Merge pull request #2352 from PolicyEngine/SonaliBedge/2237-tracer-get-method-handling-invalid-input-incorrectly

Added validation in get_tracer method to handle inputs

Author: anth-volk

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    fixed:
+    - Added functions and tests to handle invalid or incorrect input parameters.

policyengine_api/routes/tracer_analysis_routes.py

@@ -8,6 +8,8 @@
     TracerAnalysisService,
 )
 import json
+from policyengine_api.country import COUNTRY_PACKAGE_VERSIONS
+import re
 
 tracer_analysis_bp = Blueprint("tracer_analysis", __name__)
 tracer_analysis_service = TracerAnalysisService()
@@ -26,6 +28,10 @@ def execute_tracer_analysis(country_id):
     household_id = payload.get("household_id")
     policy_id = payload.get("policy_id")
     variable = payload.get("variable")
+    api_version = COUNTRY_PACKAGE_VERSIONS[country_id]
+
+    if not isinstance(variable, str):
+        raise BadRequest("variable must be a string")
 
     analysis, analysis_type = tracer_analysis_service.execute_analysis(
         country_id,

policyengine_api/services/tracer_analysis_service.py

@@ -112,10 +112,10 @@ def _parse_tracer_output(self, tracer_output, target_variable):
         # Create a regex pattern to match the exact variable name
         # This will match the variable name followed by optional whitespace,
         # then optional angle brackets with any content, then optional whitespace
-
         pattern = (
             rf"^(\s*)({re.escape(target_variable)})(?!\w)\s*(?:<[^>]*>)?\s*"
         )
+
         for line in tracer_output:
             # Count leading spaces to determine indentation level
             indent = len(line) - len(line.strip())

policyengine_api/utils/payload_validators/validate_tracer_analysis_payload.py

@@ -8,4 +8,22 @@ def validate_tracer_analysis_payload(payload: dict):
         if key not in payload:
             return False, f"Missing required key: {key}"
 
+    # Validate types and formats
+    household_id = payload["household_id"]
+    policy_id = payload["policy_id"]
+    variable = payload["variable"]
+
+    if not isinstance(household_id, (str, int)) or (
+        isinstance(household_id, str) and not household_id.isdigit()
+    ):
+        return False, "household_id must be a numeric integer or string"
+
+    if not isinstance(policy_id, (str, int)) or (
+        isinstance(policy_id, str) and not policy_id.isdigit()
+    ):
+        return False, "policy_id must be a numeric integer or string"
+
+    if not isinstance(variable, str):
+        return False, "variable must be a string"
+
     return True, None

tests/to_refactor/python/test_tracer_analysis_routes.py

@@ -1,6 +1,15 @@
 import pytest
 from flask import json
 from unittest.mock import patch
+from werkzeug.exceptions import BadRequest
+
+# constants
+VALID_HOUSEHOLD_ID = 123
+VALID_POLICY_ID = 456
+INVALID_HOUSEHOLD_ID = "abc123"
+INVALID_POLICY_ID = "invalid-id"
+TEST_VARIABLE = "disposable_income"
+INVALID_VARIABLE = 123
 
 
 @patch("policyengine_api.services.tracer_analysis_service.local_database")
@@ -41,8 +50,8 @@ def test_execute_tracer_analysis_no_tracer(mock_db, rest_client):
     response = rest_client.post(
         "/us/tracer-analysis",
         json={
-            "household_id": "test_household",
-            "policy_id": "test_policy",
+            "household_id": VALID_HOUSEHOLD_ID,
+            "policy_id": VALID_POLICY_ID,
             "variable": "disposable_income",
         },
     )
@@ -85,6 +94,32 @@ def test_execute_tracer_analysis_ai_error(
     assert json.loads(response.data)["status"] == "error"
 
 
+@patch("policyengine_api.services.tracer_analysis_service.local_database")
+def test_invalid_variable_types(mock_db, rest_client):
+    """Test that different non-string variable types are rejected"""
+    invalid_variables = [
+        123,
+        None,
+        {"key": "value"},
+        ["list"],
+        True,
+    ]
+
+    for invalid_var in invalid_variables:
+        response = rest_client.post(
+            "/us/tracer-analysis",
+            json={
+                "household_id": VALID_HOUSEHOLD_ID,
+                "policy_id": VALID_POLICY_ID,
+                "variable": invalid_var,
+            },
+        )
+        assert response.status_code == 400
+        assert (
+            "variable must be a string" in json.loads(response.data)["message"]
+        )
+
+
 # Test invalid country
 def test_invalid_country(rest_client):
     response = rest_client.post(
@@ -97,3 +132,93 @@ def test_invalid_country(rest_client):
     )
     assert response.status_code == 400
     assert b"Country invalid_country not found" in response.data
+
+
+def test_invalid_household_id_format(rest_client):
+    """Test that non-numeric household_id is rejected"""
+    response = rest_client.post(
+        "/us/tracer-analysis",
+        json={
+            "household_id": INVALID_HOUSEHOLD_ID,
+            "policy_id": VALID_POLICY_ID,
+            "variable": "disposable_income",
+        },
+    )
+    assert response.status_code == 400
+    assert (
+        "household_id must be a numeric integer or string"
+        in json.loads(response.data)["message"]
+    )
+
+
+def test_invalid_policy_id_format(rest_client):
+    """Test that non-numeric policy_id is rejected"""
+    response = rest_client.post(
+        "/us/tracer-analysis",
+        json={
+            "household_id": VALID_HOUSEHOLD_ID,
+            "policy_id": INVALID_POLICY_ID,
+            "variable": "disposable_income",
+        },
+    )
+    assert response.status_code == 400
+    assert (
+        "policy_id must be a numeric integer or string"
+        in json.loads(response.data)["message"]
+    )
+
+
+def test_empty_household_id(rest_client):
+    """Test that empty household_id is rejected"""
+    response = rest_client.post(
+        "/us/tracer-analysis",
+        json={
+            "household_id": "",
+            "policy_id": VALID_POLICY_ID,
+            "variable": "disposable_income",
+        },
+    )
+    assert response.status_code == 400
+
+
+def test_missing_required_fields(rest_client):
+    """Test that missing required fields are rejected"""
+    response = rest_client.post(
+        "/us/tracer-analysis",
+        json={
+            # household_id missing
+            "policy_id": VALID_POLICY_ID,
+            "variable": "disposable_income",
+        },
+    )
+    assert response.status_code == 400
+
+
+def test_invalid_types(rest_client):
+    """Test that invalid types are rejected"""
+    response = rest_client.post(
+        "/us/tracer-analysis",
+        json={
+            "household_id": None,  # Invalid type
+            "policy_id": INVALID_POLICY_ID,
+            "variable": "disposable_income",
+        },
+    )
+    assert response.status_code == 400
+
+
+def test_validate_tracer_analysis_payload_failure(rest_client):
+    """Test handling of invalid payload from validate_tracer_analysis_payload"""
+    response = rest_client.post(
+        "/us/tracer-analysis",
+        json={
+            # Missing required field 'variable'
+            "household_id": VALID_HOUSEHOLD_ID,
+            "policy_id": VALID_POLICY_ID,
+        },
+    )
+    assert response.status_code == 400
+    assert (
+        "Missing required key: variable"
+        in json.loads(response.data)["message"]
+    )