Applied Pull Requests Suggestions

Author: HerveRI

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/fixtures/services/tracer_analysis_service_fixture.py

@@ -5,21 +5,28 @@
 )
 from werkzeug.exceptions import NotFound
 
-from tests.fixtures.services.tracer_analysis_service_fixture import (
+from tests.fixtures.services.tracer_analysis_service import (
     sample_tracer_data,
     sample_expected_segment,
-    get_tracer,
-    get_existing_analysis,
-    parse_tracer_output,
-    trigger_ai_analysis,
+    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, get_tracer, parse_tracer_output, get_existing_analysis
+        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),
@@ -29,33 +36,34 @@ def test_execute_analysis_static(
         """
 
         analysis, analysis_type = service.execute_analysis(
-            "us", "71424", "2", "market_income"
+            country_id, household_id, policy_id, target_variable
         )
 
         assert analysis == "Existing static analysis"
         assert analysis_type == "static"
 
     def test_execute_analysis_streaming(
         self,
-        get_tracer,
-        parse_tracer_output,
-        get_existing_analysis,
-        trigger_ai_analysis,
+        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 ",
+        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
-        get_existing_analysis.return_value = None
+        mock_get_existing_analysis.return_value = None
 
         analysis, analysis_type = service.execute_analysis(
-            "us", "71424", "2", "market_income"
+            country_id, household_id, policy_id, target_variable
         )
 
+        expected_streaming_output = ["stream chunk 1", "stream chunk 2"]
         streaming_output = list(analysis)
-        assert streaming_output == ["stream chunk 1", "stream chunk 2"]
+        assert streaming_output == expected_streaming_output
         assert analysis_type == "streaming"