Test SSI disability training usecols

Author: MaxGhenis

tests/unit/datasets/test_sipp_ssi_disability.py

@@ -11,6 +11,7 @@
     preserve_under_65_ssi_disability_criteria,
     prepare_ssi_disability_receiver,
 )
+from policyengine_us_data.datasets.sipp.sipp import SSI_DISABILITY_COLUMNS
 
 
 def _base_sipp_frame() -> pd.DataFrame:
@@ -66,6 +67,10 @@ def test_build_ssi_disability_training_frame_uses_all_disability_amounts():
     assert result["has_disability_income"].iloc[0]
 
 
+def test_ssi_disability_training_usecols_include_label_and_income_columns():
+    assert {"TPTOTINC", "RSSI_YRYN"} <= set(SSI_DISABILITY_COLUMNS)
+
+
 def test_prepare_ssi_disability_receiver_fills_missing_predictors():
     result = prepare_ssi_disability_receiver(
         pd.DataFrame(