Merge pull request #2694 from PolicyEngine/test/modify-varying-your-earnings

Refactor varying your earnings test

Author: anth-volk

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    changed:
+    - Refactored test_varying_your_earnings to better serve testing needs

tests/fixtures/integration/simulations.py

@@ -0,0 +1,76 @@
+"""
+Test fixtures and constants for axes calculation tests
+"""
+
+import pytest
+from unittest.mock import Mock, MagicMock, patch
+from policyengine_api.endpoints.household import add_yearly_variables
+
+STANDARD_AXES_COUNT = 401  # Not formally defined anywhere, but this value is used throughout the API
+SMALL_AXES_COUNT = 5
+TEST_YEAR = "2025"
+TEST_STATE = "NY"
+TEST_COUNTRY_ID = "us"
+
+BASE_HOUSEHOLD = {
+    "people": {
+        "adult": {
+            "age": {TEST_YEAR: 40},
+            "employment_income": {TEST_YEAR: 30000},
+        },
+    },
+    "tax_units": {"unit": {"members": ["adult"]}},
+    "households": {
+        "household": {
+            "members": ["adult"],
+            "state_name": {TEST_YEAR: TEST_STATE},
+        }
+    },
+}
+
+SMALL_AXES_CONFIG = [
+    [
+        {
+            "name": "employment_income",
+            "period": TEST_YEAR,
+            "min": 20000,
+            "max": 40000,
+            "count": SMALL_AXES_COUNT,
+        }
+    ]
+]
+
+
+# Pytest fixtures
+def create_base_household():
+    """Returns a copy of the base household configuration"""
+    import copy
+
+    return copy.deepcopy(BASE_HOUSEHOLD)
+
+
+def create_small_axes():
+    """Returns the small axes configuration for faster tests"""
+    import copy
+
+    return copy.deepcopy(SMALL_AXES_CONFIG)
+
+
+def create_household_with_axes(base_household, axes_config):
+    """Helper function to add axes to a household"""
+    import copy
+
+    household = copy.deepcopy(base_household)
+    household["axes"] = axes_config
+    return household
+
+
+def setup_small_axes_household(base_household, small_axes_config):
+    """Fixture to setup a household with small axes for testing"""
+    household_with_axes = create_household_with_axes(
+        base_household, small_axes_config
+    )
+    household_with_axes = add_yearly_variables(
+        household_with_axes, TEST_COUNTRY_ID
+    )
+    return household_with_axes

tests/integration/test_simulations.py

@@ -0,0 +1,59 @@
+import pytest
+from unittest.mock import Mock, patch, MagicMock
+import numpy as np
+from policyengine_api.country import COUNTRIES
+from policyengine_api.endpoints.household import add_yearly_variables
+from tests.fixtures.integration.simulations import (
+    TEST_COUNTRY_ID,
+    SMALL_AXES_COUNT,
+    setup_small_axes_household,
+    create_base_household,
+    create_small_axes,
+)
+
+
+class TestSimsWithAxes:
+
+    def test__given_any_number_of_axes__sim_returns_valid_arrays(
+        self,
+    ):  # , patched_add_yearly_variables, patched_countries_get):
+        """Integration test with small axes for speed"""
+        base_household = create_base_household()
+        small_axes_config = create_small_axes()
+        household_with_axes = setup_small_axes_household(
+            base_household, small_axes_config
+        )
+        country = COUNTRIES.get(TEST_COUNTRY_ID)
+        result = country.calculate(household_with_axes, {})
+
+        # This variable does not function like others; it is a list of member names and is not calculated
+        FORBIDDEN_VARIABLES = ["members"]
+
+        # Verify we get array results
+        for entity_type in result:
+            print("Entity type: ", entity_type)
+            if entity_type == "axes":
+                continue
+            for entity_id in result[entity_type]:
+                print("Entity ID: ", entity_id)
+                for variable_name in result[entity_type][entity_id]:
+                    print("Variable name: ", variable_name)
+                    if variable_name in FORBIDDEN_VARIABLES:
+                        continue
+                    for period in result[entity_type][entity_id][
+                        variable_name
+                    ]:
+                        print("Period: ", period)
+                        value = result[entity_type][entity_id][variable_name][
+                            period
+                        ]
+                        print(f"Value: {value}")
+                        if isinstance(value, list):
+                            # Assert no Nones
+                            assert all(
+                                v is not None for v in value
+                            ), f"None found in {variable_name} for {entity_id} in {period}"
+                            # Assert correct length
+                            assert (
+                                len(value) == SMALL_AXES_COUNT
+                            ), f"Expected {SMALL_AXES_COUNT} values for {variable_name}, got {len(value)}"