Small follow-ups from v4 audit

changelog.d/small-follow-ups.fixed.md

@@ -0,0 +1,6 @@
+Two small follow-ups surfaced by the v4 audit:
+
+- ``Simulation.save()`` now raises a helpful ``ValueError`` when ``output_dataset`` is still ``None`` (run or ensure was never called), instead of the raw ``AttributeError: 'NoneType' object has no attribute 'save'``.
+- New ``scripts/check_data_staleness.py`` prints a one-line verdict per country comparing the bundled release manifest's pinned model version against the installed country package; exits non-zero when any country is stale. Drop into CI to get automated "bundle is behind" alerts without waiting for a human to notice.
+
+Also clarified the ``load()`` docstring — ``created_at``/``updated_at`` are filesystem ``ctime``/``mtime`` approximations, not a true round-trip of the original timestamps.

scripts/check_data_staleness.py

@@ -0,0 +1,79 @@
+"""Flag when the bundled country release manifest falls behind upstream.
+
+Each country release manifest at
+``src/policyengine/data/release_manifests/{country}.json`` pins the
+country-model version the certified microdata was built with. If
+that pin drifts behind the *currently installed* country-model
+version, downstream calculations may diverge from what the data
+was calibrated against. This script surfaces that drift so a CI
+job (or a human) can decide whether to kick off a data rebuild.
+
+Outputs a single-line verdict per country, plus a summary. Exits
+non-zero when any country is stale so CI can gate on it.
+
+No network access: reads both sides locally via
+``importlib.metadata`` + the bundled JSON.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from importlib import metadata
+from importlib.util import find_spec
+from pathlib import Path
+
+from policyengine.provenance.manifest import CountryReleaseManifest
+
+MANIFEST_DIR = (
+    Path(__file__).resolve().parent.parent
+    / "src"
+    / "policyengine"
+    / "data"
+    / "release_manifests"
+)
+
+COUNTRIES = ("us", "uk")
+
+
+def check_country(country: str) -> tuple[str, bool]:
+    """Return ``(one-line-verdict, is_stale)``."""
+    manifest_path = MANIFEST_DIR / f"{country}.json"
+    manifest_json = json.loads(manifest_path.read_text())
+    manifest = CountryReleaseManifest.model_validate(manifest_json)
+
+    pkg = manifest.model_package.name
+    pinned = manifest.model_package.version
+
+    if find_spec(pkg.replace("-", "_")) is None:
+        return (
+            f"{country}: {pkg} not installed; skipping staleness check",
+            False,
+        )
+
+    installed = metadata.version(pkg)
+    if installed == pinned:
+        return (
+            f"{country}: ok ({pkg} {installed} matches bundle pin)",
+            False,
+        )
+    return (
+        f"{country}: STALE ({pkg} installed={installed}, "
+        f"bundle pin={pinned}; consider a release-bundle refresh)",
+        True,
+    )
+
+
+def main() -> int:
+    verdicts = []
+    any_stale = False
+    for country in COUNTRIES:
+        verdict, stale = check_country(country)
+        verdicts.append(verdict)
+        any_stale = any_stale or stale
+    print("\n".join(verdicts))
+    return 1 if any_stale else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/policyengine/tax_benefit_models/common/model_version.py

@@ -290,13 +290,21 @@ def resolve_entity_variables(
 
     def save(self, simulation: Simulation) -> None:
         """Persist the simulation's output dataset to its bundled filepath."""
+        if simulation.output_dataset is None:
+            raise ValueError(
+                "Simulation.save() called with no output_dataset. Run "
+                "simulation.run() or simulation.ensure() first so there is "
+                "something to persist."
+            )
         simulation.output_dataset.save()
 
     def load(self, simulation: Simulation) -> None:
         """Rehydrate the simulation's output dataset from disk.
 
-        Loads timestamps from filesystem metadata when the file exists so
-        serialised simulations round-trip ``created_at``/``updated_at``.
+        Filesystem ``ctime``/``mtime`` on the output file are mirrored onto
+        ``simulation.created_at`` / ``updated_at`` on load — these fields
+        are not written back on ``save()``, so they're filesystem
+        approximations, not a true round-trip of the original timestamps.
         """
         filepath = str(
             Path(simulation.dataset.filepath).parent / (simulation.id + ".h5")

tests/test_small_follow_ups.py

@@ -0,0 +1,112 @@
+"""Regression tests for small follow-up fixes from the audit.
+
+Grouped together because each is small:
+
+- ``Simulation.save()`` raises a helpful ``ValueError`` when
+  ``output_dataset`` is still ``None`` (i.e. nothing has been
+  computed yet), instead of the raw
+  ``AttributeError: 'NoneType' object has no attribute 'save'``.
+
+- ``scripts/check_data_staleness.py`` detects when the bundled
+  release manifest's pinned country-model version drifts away
+  from the installed one.
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+
+
+def test__save_with_no_output_dataset_raises_helpfully() -> None:
+    """Before the fix: raw ``AttributeError`` from ``None.save()``.
+    After: ``ValueError`` that explicitly names the missing step.
+    """
+    pytest.importorskip("policyengine_us")
+
+    import policyengine as pe
+    from policyengine.core import Simulation
+    from tests.fixtures.filtering_fixtures import create_us_test_dataset
+
+    sim = Simulation(
+        dataset=create_us_test_dataset(),
+        tax_benefit_model_version=pe.us.model,
+    )
+    # ``output_dataset`` starts as ``None`` until ``run()`` / ``ensure()``.
+    assert sim.output_dataset is None
+
+    with pytest.raises(ValueError, match="no output_dataset"):
+        sim.save()
+
+
+def test__check_data_staleness_script_runs_and_reports_status() -> None:
+    """Script exits 0 when pins match the installed country packages
+    and prints one verdict line per country.
+
+    In local dev against the pinned country packages, this should
+    always pass. In CI we want a non-zero exit to gate the release
+    pipeline, which we verify separately by inspecting the script
+    logic itself (not re-running with a mismatch here, which would
+    require reinstalling a country package).
+    """
+    pytest.importorskip("policyengine_us")
+
+    result = subprocess.run(
+        [sys.executable, str(REPO_ROOT / "scripts" / "check_data_staleness.py")],
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    # Stdout has one line per configured country.
+    lines = [line for line in result.stdout.splitlines() if line.strip()]
+    assert any(line.startswith("us:") for line in lines), result.stdout
+    assert any(line.startswith("uk:") for line in lines), result.stdout
+    # Exit code mirrors "any country stale". Since we pin to specific
+    # country versions in pyproject.toml extras, matching is expected
+    # in a clean install. If this fails, either (a) the pin drifted
+    # or (b) an engineer has a newer country package installed
+    # locally — both legitimate reasons to run the refresh helper.
+    if result.returncode != 0:
+        # Still check output shape is correct even if staleness is real.
+        assert "STALE" in result.stdout, (
+            f"Non-zero exit but no STALE report: {result.stdout!r} / {result.stderr!r}"
+        )
+
+
+def test__check_data_staleness_script_detects_drift() -> None:
+    """Direct test of the country-check function so we don't have to
+    actually break the installed environment to exercise the drift
+    path.
+    """
+    pytest.importorskip("policyengine_us")
+
+    # Import the script as a module so we can call its helper.
+    import importlib.util
+
+    spec = importlib.util.spec_from_file_location(
+        "check_data_staleness",
+        REPO_ROOT / "scripts" / "check_data_staleness.py",
+    )
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+
+    # Monkey-patch the manifest path to a tmp file with a drifted pin.
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp_manifest_dir = Path(tmp)
+        real_manifest = json.loads((module.MANIFEST_DIR / "us.json").read_text())
+        real_manifest["model_package"]["version"] = "0.0.0-drift"
+        (tmp_manifest_dir / "us.json").write_text(json.dumps(real_manifest, indent=2))
+        module.MANIFEST_DIR = tmp_manifest_dir
+        verdict, stale = module.check_country("us")
+        assert stale is True
+        assert "STALE" in verdict
+        assert "0.0.0-drift" in verdict