Allow raw CPS source artifact through publication contract (#1012)

Author: MaxGhenis

changelog.d/1012.fixed.md

@@ -0,0 +1 @@
+Allow raw CPS source artifacts to pass publication validation while retaining the leaf-input export guard for final enhanced CPS datasets.

policyengine_us_data/storage/upload_completed_datasets.py

@@ -46,6 +46,15 @@
     "cps_2024.h5",
 }
 
+# Raw CPS is a source artifact used to construct final simulation datasets.
+# It intentionally carries some Census source measures whose names now map to
+# formula variables in policyengine-us. The leaf-input export contract is
+# enforced on final simulation datasets instead.
+ENFORCE_LEAF_INPUT_EXPORT_BY_FILENAME = {
+    "enhanced_cps_2024.h5": True,
+    "cps_2024.h5": False,
+}
+
 # Minimum file sizes in bytes for validated datasets.
 MIN_FILE_SIZES = {
     "enhanced_cps_2024.h5": 95 * 1024 * 1024,  # 95 MB
@@ -424,7 +433,13 @@ def validate_dataset(file_path: Path) -> None:
         )
 
     try:
-        contract_summary = validate_dataset_contract(file_path)
+        contract_summary = validate_dataset_contract(
+            file_path,
+            enforce_no_computed_policyengine_us_variables=ENFORCE_LEAF_INPUT_EXPORT_BY_FILENAME.get(
+                filename,
+                True,
+            ),
+        )
     except DatasetContractError as e:
         errors.append(f"Dataset contract validation failed: {e}")
         raise DatasetValidationError(

policyengine_us_data/utils/dataset_validation.py

@@ -268,6 +268,7 @@ def validate_dataset_contract(
     microsimulation_cls=None,
     dataset_loader=None,
     smoke_test_variable: str = "household_weight",
+    enforce_no_computed_policyengine_us_variables: bool = True,
 ) -> DatasetContractSummary:
     file_path = Path(file_path)
     policyengine_us_info = assert_locked_policyengine_us_version()
@@ -281,7 +282,7 @@ def validate_dataset_contract(
 
     dataset_lengths = _dataset_lengths(file_path)
     time_period = _infer_time_period_from_file(file_path)
-    if time_period is not None:
+    if time_period is not None and enforce_no_computed_policyengine_us_variables:
         assert_no_computed_policyengine_us_variables_exported(
             variable_names=dataset_lengths.keys(),
             time_period=time_period,

tests/unit/test_dataset_validation.py

@@ -239,6 +239,39 @@ def test_validate_dataset_contract_rejects_computed_policyengine_variables(
         )
 
 
+def test_validate_dataset_contract_can_skip_computed_variable_check(
+    tmp_path,
+    monkeypatch,
+):
+    file_path = tmp_path / "cps_2024.h5"
+    _write_test_h5(
+        file_path,
+        {
+            "person_id": np.array([101, 102], dtype=np.int32),
+            "household_id": np.array([501], dtype=np.int32),
+            "computed_income": np.array([10_000.0, 20_000.0], dtype=np.float32),
+            "household_weight": np.array([1.5], dtype=np.float32),
+        },
+    )
+    monkeypatch.setattr(
+        "policyengine_us_data.utils.dataset_validation.assert_locked_policyengine_us_version",
+        lambda: PolicyEngineUSBuildInfo(version="1.587.0"),
+    )
+    tbs = _fake_tax_benefit_system()
+    tbs.variables["computed_income"] = _fake_variable(
+        "person",
+        adds=["computed_income_before_response"],
+    )
+
+    validate_dataset_contract(
+        file_path,
+        tax_benefit_system=tbs,
+        microsimulation_cls=_FakeMicrosimulation,
+        dataset_loader=lambda path: path,
+        enforce_no_computed_policyengine_us_variables=False,
+    )
+
+
 def test_validate_dataset_contract_allows_future_period_formulas(tmp_path, monkeypatch):
     file_path = tmp_path / "enhanced_cps_2024.h5"
     _write_test_h5(

tests/unit/test_upload_completed_datasets.py

@@ -109,6 +109,15 @@ def _fake_tax_benefit_system():
     )
 
 
+def _fake_variable(entity_key, *, formulas=None, adds=None, subtracts=None):
+    return SimpleNamespace(
+        entity=SimpleNamespace(key=entity_key),
+        formulas=formulas or {},
+        adds=adds,
+        subtracts=subtracts,
+    )
+
+
 def _write_h5(path, datasets: dict[str, np.ndarray]) -> None:
     with h5py.File(path, "w") as h5_file:
         for name, values in datasets.items():
@@ -151,11 +160,12 @@ def patch_contract_validation(monkeypatch):
     monkeypatch.setattr(
         upload_module,
         "validate_dataset_contract",
-        lambda file_path: validate_dataset_contract(
+        lambda file_path, **kwargs: validate_dataset_contract(
             file_path,
             tax_benefit_system=_fake_tax_benefit_system(),
             microsimulation_cls=_FakeMicrosimulation,
             dataset_loader=lambda path: path,
+            **kwargs,
         ),
     )
 
@@ -221,6 +231,80 @@ def test_validate_dataset_infers_time_period_for_flat_h5(tmp_path, monkeypatch):
     assert _TimePeriodCheckingAggregateMicrosimulation.last_dataset.time_period == 2024
 
 
+def test_validate_cps_allows_source_computed_policyengine_variables(
+    tmp_path,
+    monkeypatch,
+):
+    file_path = tmp_path / "cps_2024.h5"
+    _write_h5(
+        file_path,
+        {
+            "person_id": np.array([101], dtype=np.int32),
+            "household_id": np.array([201], dtype=np.int32),
+            "employment_income": np.array([50_000.0], dtype=np.float32),
+            "household_weight": np.array([1.0], dtype=np.float32),
+        },
+    )
+    tbs = _fake_tax_benefit_system()
+    tbs.variables["employment_income"] = _fake_variable(
+        "person",
+        adds=["employment_income_before_lsr"],
+    )
+    monkeypatch.setattr(
+        upload_module,
+        "validate_dataset_contract",
+        lambda file_path, **kwargs: validate_dataset_contract(
+            file_path,
+            tax_benefit_system=tbs,
+            microsimulation_cls=_FakeMicrosimulation,
+            dataset_loader=lambda path: path,
+            **kwargs,
+        ),
+    )
+    monkeypatch.setattr(
+        "policyengine_us.Microsimulation",
+        _TimePeriodCheckingAggregateMicrosimulation,
+    )
+
+    validate_dataset(file_path)
+
+
+def test_validate_enhanced_cps_rejects_computed_policyengine_variables(
+    tmp_path,
+    monkeypatch,
+):
+    file_path = tmp_path / "enhanced_cps_2024.h5"
+    _write_h5(file_path, _minimal_enhanced_cps_contract_datasets())
+    tbs = _fake_tax_benefit_system()
+    tbs.variables["employment_income"] = _fake_variable(
+        "person",
+        adds=["employment_income_before_lsr"],
+    )
+    monkeypatch.setattr(
+        upload_module,
+        "REQUIRED_VARIABLES_BY_FILENAME",
+        {},
+    )
+    monkeypatch.setattr(
+        upload_module,
+        "validate_dataset_contract",
+        lambda file_path, **kwargs: validate_dataset_contract(
+            file_path,
+            tax_benefit_system=tbs,
+            microsimulation_cls=_FakeMicrosimulation,
+            dataset_loader=lambda path: path,
+            **kwargs,
+        ),
+    )
+    monkeypatch.setattr(
+        "policyengine_us.Microsimulation",
+        _TimePeriodCheckingAggregateMicrosimulation,
+    )
+
+    with pytest.raises(DatasetValidationError, match="employment_income"):
+        validate_dataset(file_path)
+
+
 def test_validate_dataset_rejects_temporary_reported_source_variables(
     tmp_path,
     monkeypatch,