Move fit weight test setup into fixtures

Author: anth-volk

tests/unit/fit_weights/conftest.py

@@ -0,0 +1,84 @@
+from collections.abc import Callable
+from pathlib import Path
+
+import pytest
+import yaml
+
+from policyengine_us_data.fit_weights import (
+    FitScope,
+    FittedWeightsOutputBundle,
+)
+
+
+class FakeBatch:
+    def __init__(self) -> None:
+        self.files: dict[str, bytes] = {}
+
+    def put_file(self, file_obj, destination: str) -> None:
+        self.files[destination] = file_obj.read()
+
+
+@pytest.fixture
+def artifacts_rel() -> str:
+    return "artifacts/run-1"
+
+
+@pytest.fixture
+def calibration_package_path() -> Path:
+    return Path("/pipeline/artifacts/run/calibration_package.pkl")
+
+
+@pytest.fixture
+def fake_batch() -> FakeBatch:
+    return FakeBatch()
+
+
+@pytest.fixture
+def regional_result_bytes() -> dict[str, bytes]:
+    return {
+        "weights": b"weights",
+        "geography": b"regional-geo",
+        "config": b"regional-config",
+        "log": b"regional-log",
+        "cal_log": b"regional-epoch",
+    }
+
+
+@pytest.fixture
+def national_result_bytes() -> dict[str, bytes]:
+    return {
+        "weights": b"weights",
+        "geography": b"national-geo",
+        "config": b"national-config",
+        "log": b"national-log",
+        "cal_log": b"national-epoch",
+    }
+
+
+@pytest.fixture
+def regional_output_bundle(
+    regional_result_bytes: dict[str, bytes],
+) -> FittedWeightsOutputBundle:
+    return FittedWeightsOutputBundle.from_result_bytes(
+        scope=FitScope.REGIONAL,
+        result_bytes=regional_result_bytes,
+        run_id="run-1",
+    )
+
+
+@pytest.fixture
+def national_output_bundle(
+    national_result_bytes: dict[str, bytes],
+) -> FittedWeightsOutputBundle:
+    return FittedWeightsOutputBundle.from_result_bytes(
+        scope=FitScope.NATIONAL,
+        result_bytes=national_result_bytes,
+        run_id="run-1",
+    )
+
+
+@pytest.fixture
+def stage_3_substage() -> Callable[[str], dict]:
+    data = yaml.safe_load(Path("docs/pipeline_map.yaml").read_text())
+    substages = {substage["id"]: substage for substage in data["stages"]}
+    return substages.__getitem__

tests/unit/fit_weights/test_bundles.py

@@ -10,93 +10,72 @@
 )
 
 
-class FakeBatch:
-    def __init__(self) -> None:
-        self.files: dict[str, bytes] = {}
-
-    def put_file(self, file_obj, destination: str) -> None:
-        self.files[destination] = file_obj.read()
-
-
-def test_input_bundle_exposes_calibration_package_identity_path() -> None:
+def test_input_bundle_exposes_calibration_package_identity_path(
+    calibration_package_path: Path,
+) -> None:
     bundle = FittedWeightsInputBundle(
         scope="regional",
-        calibration_package_path=Path(
-            "/pipeline/artifacts/run/calibration_package.pkl"
-        ),
+        calibration_package_path=calibration_package_path,
     )
 
     assert bundle.scope == FitScope.REGIONAL
     assert bundle.artifact_identity_paths() == {
-        "calibration_package": Path("/pipeline/artifacts/run/calibration_package.pkl")
+        "calibration_package": calibration_package_path
     }
 
 
-def test_regional_output_bundle_writes_expected_paths() -> None:
-    bundle = FittedWeightsOutputBundle.from_result_bytes(
-        scope=FitScope.REGIONAL,
-        result_bytes={
-            "weights": b"weights",
-            "geography": b"geo",
-            "config": b"config",
-            "log": b"log",
-            "cal_log": b"epoch",
-        },
-        run_id="run-1",
-    )
-    batch = FakeBatch()
-
-    written = bundle.write_artifacts(batch, "artifacts/run-1")
+def test_regional_output_bundle_writes_expected_paths(
+    artifacts_rel: str,
+    fake_batch,
+    regional_output_bundle: FittedWeightsOutputBundle,
+) -> None:
+    written = regional_output_bundle.write_artifacts(fake_batch, artifacts_rel)
 
     assert written == [
         "artifacts/run-1/calibration_weights.npy",
         "artifacts/run-1/geography_assignment.npz",
         "artifacts/run-1/unified_run_config.json",
     ]
-    assert batch.files["artifacts/run-1/calibration_weights.npy"] == b"weights"
-    assert bundle.artifact_paths("/pipeline/artifacts/run-1") == [
+    assert fake_batch.files["artifacts/run-1/calibration_weights.npy"] == b"weights"
+    assert regional_output_bundle.artifact_paths("/pipeline/artifacts/run-1") == [
         Path("/pipeline/artifacts/run-1/calibration_weights.npy"),
         Path("/pipeline/artifacts/run-1/geography_assignment.npz"),
         Path("/pipeline/artifacts/run-1/unified_run_config.json"),
     ]
 
 
-def test_national_output_bundle_writes_expected_paths() -> None:
-    bundle = FittedWeightsOutputBundle.from_result_bytes(
-        scope="national",
-        result_bytes={
-            "weights": b"weights",
-            "geography": b"geo",
-            "config": b"config",
-        },
-    )
-    batch = FakeBatch()
-
-    written = bundle.write_artifacts(batch, "artifacts/run-1")
+def test_national_output_bundle_writes_expected_paths(
+    artifacts_rel: str,
+    fake_batch,
+    national_output_bundle: FittedWeightsOutputBundle,
+) -> None:
+    written = national_output_bundle.write_artifacts(fake_batch, artifacts_rel)
 
     assert written == [
         "artifacts/run-1/national_calibration_weights.npy",
         "artifacts/run-1/national_geography_assignment.npz",
         "artifacts/run-1/national_unified_run_config.json",
     ]
-    assert batch.files["artifacts/run-1/national_calibration_weights.npy"] == b"weights"
+    assert (
+        fake_batch.files["artifacts/run-1/national_calibration_weights.npy"]
+        == b"weights"
+    )
 
 
-def test_missing_optional_epoch_log_is_allowed() -> None:
+def test_missing_optional_epoch_log_is_allowed(
+    regional_result_bytes: dict[str, bytes],
+) -> None:
+    result_bytes = dict(regional_result_bytes)
+    result_bytes.pop("cal_log")
     bundle = FittedWeightsOutputBundle.from_result_bytes(
         scope=FitScope.REGIONAL,
-        result_bytes={
-            "weights": b"weights",
-            "geography": b"geo",
-            "config": b"config",
-            "log": b"log",
-        },
+        result_bytes=result_bytes,
     )
 
     assert bundle.diagnostic_result_bytes() == {
-        "log": b"log",
+        "log": b"regional-log",
         "cal_log": None,
-        "config": b"config",
+        "config": b"regional-config",
     }
 
 
@@ -118,45 +97,32 @@ def test_missing_weights_is_a_hard_failure() -> None:
 def test_missing_required_primary_artifacts_fail_before_writes(
     missing_key: str,
     expected_role: str,
+    artifacts_rel: str,
+    fake_batch,
+    regional_result_bytes: dict[str, bytes],
 ) -> None:
-    result_bytes = {
-        "weights": b"weights",
-        "geography": b"geo",
-        "config": b"config",
-    }
+    result_bytes = dict(regional_result_bytes)
     result_bytes.pop(missing_key)
     bundle = FittedWeightsOutputBundle.from_result_bytes(
         scope=FitScope.REGIONAL,
         result_bytes=result_bytes,
     )
 
     with pytest.raises(MissingFitWeightsOutputError, match=expected_role):
-        bundle.write_artifacts(FakeBatch(), "artifacts/run-1")
+        bundle.write_artifacts(fake_batch, artifacts_rel)
 
 
-def test_diagnostics_are_scoped_to_the_output_bundle() -> None:
-    regional = FittedWeightsOutputBundle.from_result_bytes(
-        scope=FitScope.REGIONAL,
-        result_bytes={
-            "weights": b"weights",
-            "geography": b"regional-geo",
-            "config": b"regional-config",
-            "log": b"regional-log",
-            "cal_log": b"regional-epoch",
-        },
+def test_diagnostics_are_scoped_to_the_output_bundle(
+    regional_output_bundle: FittedWeightsOutputBundle,
+    national_output_bundle: FittedWeightsOutputBundle,
+) -> None:
+    assert (
+        regional_output_bundle.artifacts.diagnostics.filename
+        == "unified_diagnostics.csv"
     )
-    national = FittedWeightsOutputBundle.from_result_bytes(
-        scope=FitScope.NATIONAL,
-        result_bytes={
-            "weights": b"weights",
-            "geography": b"national-geo",
-            "config": b"national-config",
-            "log": b"national-log",
-            "cal_log": b"national-epoch",
-        },
+    assert (
+        national_output_bundle.artifacts.diagnostics.filename
+        == "national_unified_diagnostics.csv"
     )
-
-    assert regional.artifacts.diagnostics.filename == "unified_diagnostics.csv"
-    assert national.artifacts.diagnostics.filename == "national_unified_diagnostics.csv"
-    assert regional.diagnostic_result_bytes()["log"] == b"regional-log"
-    assert national.diagnostic_result_bytes()["log"] == b"national-log"
+    assert regional_output_bundle.diagnostic_result_bytes()["log"] == b"regional-log"
+    assert national_output_bundle.diagnostic_result_bytes()["log"] == b"national-log"

tests/unit/fit_weights/test_pipeline_docs.py

@@ -1,18 +1,7 @@
-from pathlib import Path
-
-import yaml
-
 from policyengine_us_data.fit_weights import FitScope, fit_artifacts_for_scope
 from scripts.extract_pipeline_docs import scan_decorated_objects
 
 
-def _substage(substage_id: str) -> dict:
-    data = yaml.safe_load(Path("docs/pipeline_map.yaml").read_text())
-    return next(
-        substage for substage in data["stages"] if substage["id"] == substage_id
-    )
-
-
 def test_fit_weights_identity_nodes_are_in_generated_pipeline_docs() -> None:
     decorated = scan_decorated_objects()
 
@@ -26,11 +15,11 @@ def test_fit_weights_identity_nodes_are_in_generated_pipeline_docs() -> None:
     )
 
 
-def test_stage_3_pipeline_map_labels_match_scoped_artifacts() -> None:
+def test_stage_3_pipeline_map_labels_match_scoped_artifacts(stage_3_substage) -> None:
     regional_artifacts = fit_artifacts_for_scope(FitScope.REGIONAL)
     national_artifacts = fit_artifacts_for_scope(FitScope.NATIONAL)
-    regional = _substage("3a_weight_fitting_regional")
-    national = _substage("3b_weight_fitting_national")
+    regional = stage_3_substage("3a_weight_fitting_regional")
+    national = stage_3_substage("3b_weight_fitting_national")
 
     regional_nodes = {node["id"]: node for node in regional["extra_nodes"]}
     national_nodes = {node["id"]: node for node in national["extra_nodes"]}