Merge pull request #2310 from PolicyEngine/HerveRI/tests_tracer_analysis_service

Added tests for execute_analysis function

Author: anth-volk

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: minor
+  changes:
+    added:
+      - Tests for execute_analysis function 

tests/fixtures/services/tracer_analysis_service.py

@@ -1,3 +1,11 @@
+import pytest
+import json
+from policyengine_api.services.tracer_analysis_service import (
+    TracerAnalysisService,
+)
+from unittest.mock import patch
+
+
 valid_tracer_output = [
     "        snap<2027, (default)> = [6769.799]",
     "          snap<2027-01, (default)> = [561.117]",
@@ -22,3 +30,55 @@
 spliced_valid_tracer_output_leaf_variable = valid_tracer_output[8:]
 
 empty_tracer = []
+
+
+@pytest.fixture
+def sample_tracer_data():
+    return valid_tracer_output
+
+
+@pytest.fixture
+def sample_expected_segment():
+    return spliced_valid_tracer_output_nested_variable
+
+
+@pytest.fixture
+def mock_get_tracer(sample_tracer_data):
+    with patch.object(
+        TracerAnalysisService, "get_tracer", return_value=sample_tracer_data
+    ) as mock:
+        yield mock
+
+
+@pytest.fixture
+def mock_parse_tracer_output(sample_expected_segment):
+    with patch.object(
+        TracerAnalysisService,
+        "_parse_tracer_output",
+        return_value=sample_expected_segment,
+    ) as mock:
+        yield mock
+
+
+@pytest.fixture
+def mock_get_existing_analysis():
+    with patch.object(
+        TracerAnalysisService,
+        "get_existing_analysis",
+        return_value="Existing static analysis",
+    ) as mock:
+        yield mock
+
+
+@pytest.fixture
+def mock_trigger_ai_analysis():
+    def dummy_generator():
+        yield "stream chunk 1"
+        yield "stream chunk 2"
+
+    with patch.object(
+        TracerAnalysisService,
+        "trigger_ai_analysis",
+        return_value=dummy_generator(),
+    ) as mock:
+        yield mock

tests/unit/services/test_execute_analysis.py

@@ -0,0 +1,69 @@
+import pytest
+import json
+from policyengine_api.services.tracer_analysis_service import (
+    TracerAnalysisService,
+)
+from werkzeug.exceptions import NotFound
+
+from tests.fixtures.services.tracer_analysis_service import (
+    sample_tracer_data,
+    sample_expected_segment,
+    mock_get_tracer,
+    mock_get_existing_analysis,
+    mock_parse_tracer_output,
+    mock_trigger_ai_analysis,
+)
+
+service = TracerAnalysisService()
+country_id = "us"
+household_id = "71424"
+policy_id = "2"
+target_variable = "takes_up_snap_if_eligible"
+
+
+class TestExecuteAnalysis:
+    def test_execute_analysis_static(
+        self,
+        mock_get_tracer,
+        mock_parse_tracer_output,
+        mock_get_existing_analysis,
+    ):
+        """
+        GIVEN a valid tracer data and an expected parsed segment (included as fixture),
+        AND get_existing_analysis returns a static analysis (included as fixture),
+        WHEN execute_analysis is called,
+        THEN then a static analysis with the "static" flag should be returned.
+        """
+
+        analysis, analysis_type = service.execute_analysis(
+            country_id, household_id, policy_id, target_variable
+        )
+
+        assert analysis == "Existing static analysis"
+        assert analysis_type == "static"
+
+    def test_execute_analysis_streaming(
+        self,
+        mock_get_tracer,
+        mock_parse_tracer_output,
+        mock_get_existing_analysis,
+        mock_trigger_ai_analysis,
+    ):
+        """
+        GIVEN a valid tracer data and an expected parsed segment,
+        AND get_existing_analysis returns None,
+        WHEN execute_analysis is called,
+        THEN trigger_ai_analysis  is called and returns a generator with the "streaming" flag.
+        """
+
+        # When existing analysis value is None
+        mock_get_existing_analysis.return_value = None
+
+        analysis, analysis_type = service.execute_analysis(
+            country_id, household_id, policy_id, target_variable
+        )
+
+        expected_streaming_output = ["stream chunk 1", "stream chunk 2"]
+        streaming_output = list(analysis)
+        assert streaming_output == expected_streaming_output
+        assert analysis_type == "streaming"