Fix reform-id calibration test coverage

Author: MaxGhenis

policyengine_us_data/db/create_database_tables.py

@@ -506,15 +506,21 @@ def create_database(
     # Create validation triggers
     create_validation_triggers(engine)
 
-    # Create SQL views
+    create_or_replace_views(engine)
+
+    logger.info(f"Database and tables created successfully at {db_uri}")
+    return engine
+
+
+def create_or_replace_views(engine) -> None:
+    """Refresh SQL views so existing databases pick up schema changes."""
     with engine.connect() as conn:
+        conn.execute(text("DROP VIEW IF EXISTS stratum_domain"))
+        conn.execute(text("DROP VIEW IF EXISTS target_overview"))
         conn.execute(text(STRATUM_DOMAIN_VIEW))
         conn.execute(text(TARGET_OVERVIEW_VIEW))
         conn.commit()
 
-    logger.info(f"Database and tables created successfully at {db_uri}")
-    return engine
-
 
 if __name__ == "__main__":
     engine = create_database()

policyengine_us_data/tests/test_calibration/test_unified_calibration.py

@@ -459,7 +459,7 @@ def test_county_var_uses_county_values(self):
         person_hh_idx = np.array([0, 1, 2, 3])
 
         builder = UnifiedMatrixBuilder.__new__(UnifiedMatrixBuilder)
-        hh_vars, _ = builder._assemble_clone_values(
+        hh_vars, _, _ = builder._assemble_clone_values(
             state_values,
             clone_states,
             person_hh_idx,
@@ -499,7 +499,7 @@ def test_non_county_var_uses_state_values(self):
         person_hh_idx = np.array([0, 1, 2, 3])
 
         builder = UnifiedMatrixBuilder.__new__(UnifiedMatrixBuilder)
-        hh_vars, _ = builder._assemble_clone_values(
+        hh_vars, _, _ = builder._assemble_clone_values(
             state_values,
             clone_states,
             person_hh_idx,

policyengine_us_data/tests/test_calibration/test_unified_matrix_builder.py

@@ -212,7 +212,11 @@ def test_reform_targets_preserved(self):
     def test_inactive_targets_are_excluded(self):
         b = self._make_builder(time_period=2024)
         df = b._query_targets({"stratum_ids": [1], "variables": ["aca_ptc"]})
-        baseline_rows = df[(df["variable"] == "aca_ptc") & (df["reform_id"] == 0)]
+        baseline_rows = df[
+            (df["variable"] == "aca_ptc")
+            & (df["reform_id"] == 0)
+            & (df["stratum_id"] == 1)
+        ]
         self.assertEqual(len(baseline_rows), 1)
         self.assertEqual(int(baseline_rows.iloc[0]["period"]), 2022)
         self.assertEqual(float(baseline_rows.iloc[0]["value"]), 10000.0)
@@ -526,6 +530,7 @@ def test_return_structure_no_takeup(self, mock_msim_cls, mock_gcv):
             sim=None,
             target_vars={"snap"},
             constraint_vars={"income"},
+            reform_vars=set(),
             geography=geo,
             rerandomize_takeup=False,
         )
@@ -560,6 +565,7 @@ def test_fresh_sim_per_state(self, mock_msim_cls, mock_gcv):
             sim=None,
             target_vars={"snap"},
             constraint_vars=set(),
+            reform_vars=set(),
             geography=geo,
             rerandomize_takeup=False,
         )
@@ -582,6 +588,7 @@ def test_state_fips_set_correctly(self, mock_msim_cls, mock_gcv):
             sim=None,
             target_vars={"snap"},
             constraint_vars=set(),
+            reform_vars=set(),
             geography=geo,
             rerandomize_takeup=False,
         )
@@ -617,6 +624,7 @@ def test_takeup_vars_forced_true(self, mock_msim_cls, mock_gcv):
             sim=None,
             target_vars={"snap"},
             constraint_vars=set(),
+            reform_vars=set(),
             geography=geo,
             rerandomize_takeup=True,
         )
@@ -664,6 +672,7 @@ def test_count_vars_skipped(self, mock_msim_cls, mock_gcv):
             sim=None,
             target_vars={"snap", "snap_count"},
             constraint_vars=set(),
+            reform_vars=set(),
             geography=geo,
             rerandomize_takeup=False,
         )
@@ -914,6 +923,7 @@ def test_workers_gt1_creates_pool(self, mock_msim_cls, mock_gcv, mock_pool_cls):
                 sim=None,
                 target_vars={"snap"},
                 constraint_vars=set(),
+                reform_vars=set(),
                 geography=geo,
                 rerandomize_takeup=False,
                 workers=2,
@@ -939,6 +949,7 @@ def test_workers_1_skips_pool(self, mock_msim_cls, mock_gcv):
                 sim=None,
                 target_vars={"snap"},
                 constraint_vars=set(),
+                reform_vars=set(),
                 geography=geo,
                 rerandomize_takeup=False,
                 workers=1,

policyengine_us_data/tests/test_database.py

@@ -1,4 +1,5 @@
 import hashlib
+from sqlalchemy import text
 
 import pytest
 from sqlalchemy.exc import IntegrityError
@@ -8,6 +9,7 @@
     Stratum,
     StratumConstraint,
     Target,
+    create_or_replace_views,
     create_database,
 )
 
@@ -226,6 +228,38 @@ def test_target_with_null_source(engine):
         assert retrieved.source is None
 
 
+def test_create_database_refreshes_existing_views(tmp_path):
+    db_uri = f"sqlite:///{tmp_path / 'test.db'}"
+    engine = create_database(db_uri)
+
+    with engine.connect() as conn:
+        conn.execute(text("DROP VIEW target_overview"))
+        conn.execute(
+            text(
+                """
+                CREATE VIEW target_overview AS
+                SELECT
+                    t.target_id,
+                    t.stratum_id,
+                    t.variable,
+                    t.value,
+                    t.period,
+                    t.active
+                FROM targets t
+                """
+            )
+        )
+        conn.commit()
+
+    create_or_replace_views(engine)
+
+    with engine.connect() as conn:
+        cursor = conn.execute(text("SELECT * FROM target_overview LIMIT 0"))
+        columns = [desc[0] for desc in cursor.cursor.description]
+
+    assert "reform_id" in columns
+
+
 def test_valid_geographic_hierarchy(engine):
     """CD under its correct state should succeed."""
     with Session(engine) as session: