Merge pull request #2953 from PolicyEngine/feature/state-specific-datasets

Explicitly use default datasets for US state-level simulations

Author: anth-volk

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: minor
+  changes:
+    changed:
+      - Set dataset to None for US state-level simulations in economy service

policyengine_api/services/economy_service.py

@@ -484,11 +484,11 @@ def _setup_data(
         if dataset == "enhanced_cps":
             return "gs://policyengine-us-data/enhanced_cps_2024.h5"
 
-        # US state-level simulations must reference pooled CPS dataset
-        if country_id == "us" and region != "us":
+        # NYC simulations must reference pooled CPS dataset
+        if region == "nyc":
             return "gs://policyengine-us-data/pooled_3_year_cps_2023.h5"
 
-        # All others receive no sim API 'data' arg
+        # All others (including US state-level simulations) receive no sim API 'data' arg
         return None
 
     # Note: The following methods that interface with the ReformImpactsService

tests/unit/services/test_economy_service.py

@@ -595,10 +595,7 @@ def test__given_us_state__returns_correct_sim_options(self):
             )
             assert sim_options["time_period"] == time_period
             assert sim_options["region"] == "state/ca"
-            assert (
-                sim_options["data"]
-                == "gs://policyengine-us-data/pooled_3_year_cps_2023.h5"
-            )
+            assert sim_options["data"] is None
 
         def test__given_enhanced_cps_state__returns_correct_sim_options(self):
             # Test with enhanced_cps dataset
@@ -720,12 +717,25 @@ def test__given_enhanced_cps_dataset__returns_correct_gcp_path(self):
             # Assert the expected value
             assert result == "gs://policyengine-us-data/enhanced_cps_2024.h5"
 
-        def test__given_us_state_dataset__returns_correct_gcp_path(self):
-            # Test with US state dataset
+        def test__given_us_state_dataset__returns_none(self):
+            # Test with US state dataset - should return None
             dataset = "us_state"
             country_id = "us"
             region = "ca"
 
+            # Create an instance of the class
+            service = EconomyService()
+            # Call the method
+            result = service._setup_data(dataset, country_id, region)
+            # Assert the expected value
+            assert result is None
+
+        def test__given_nyc_region__returns_pooled_cps(self):
+            # Test with NYC region - should return pooled CPS dataset
+            dataset = None
+            country_id = "us"
+            region = "nyc"
+
             # Create an instance of the class
             service = EconomyService()
             # Call the method