Add test for small ECPS (#306)

* Add test for small ECPS
Fixes #305

* Cast to dataframe

* Pin microimpute

* Handle str data type

* Rename test

Author: nikhilwoodruff

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    added:
+    - Add test for small ECPS.

policyengine_us_data/datasets/cps/cps.py

@@ -740,6 +740,7 @@ def add_tips(self, cps: h5py.File):
         ],
         2025,
     )
+    cps = pd.DataFrame(cps)
 
     cps["is_under_18"] = cps.age < 18
     cps["is_under_6"] = cps.age < 6

policyengine_us_data/datasets/cps/small_cps.py

@@ -5,6 +5,7 @@ def create_small_cps():
     from policyengine_us import Microsimulation
     from policyengine_us_data.datasets import EnhancedCPS_2024
     from policyengine_us_data.storage import STORAGE_FOLDER
+    from policyengine_core.enums import Enum
 
     simulation = Microsimulation(
         dataset=EnhancedCPS_2024,
@@ -16,8 +17,13 @@ def create_small_cps():
         for time_period in simulation.get_holder(variable).get_known_periods():
             values = simulation.get_holder(variable).get_array(time_period)
             values = np.array(values)
+            if simulation.tax_benefit_system.variables.get(
+                variable
+            ).value_type in (Enum, str):
+                values = values.astype("S")
             if values is not None:
                 data[variable][time_period] = values
+
         if len(data[variable]) == 0:
             del data[variable]
 
@@ -28,3 +34,8 @@ def create_small_cps():
             grp = f.create_group(variable)
             for period, values in periods.items():
                 grp.create_dataset(str(period), data=values)
+
+
+if __name__ == "__main__":
+    create_small_cps()
+    print("Small CPS dataset created successfully.")

policyengine_us_data/tests/test_datasets/test_small_cps.py

@@ -0,0 +1,17 @@
+import pytest
+import numpy as np
+
+
+@pytest.mark.parametrize("year", [2024])
+def test_small_cps_loads(year: int):
+    from policyengine_core.data import Dataset
+    from policyengine_us_data.storage import STORAGE_FOLDER
+    from policyengine_us import Microsimulation
+
+    sim = Microsimulation(
+        dataset=Dataset.from_file(
+            STORAGE_FOLDER / f"small_cps_{year}.h5",
+        )
+    )
+
+    assert not sim.calculate("household_net_income", 2025).isna().any()

pyproject.toml

@@ -23,7 +23,7 @@ dependencies = [
     "tqdm",
     "microdf_python>=0.4.3",
     "setuptools>=60",
-    "microimpute>=0.1.4",
+    "microimpute==0.1.4",
     "pip-system-certs",
     "google-cloud-storage",
     "google-auth",