|
27 | 27 | derive_clone_capped_childcare_expenses, |
28 | 28 | reconcile_ss_subcomponents, |
29 | 29 | ) |
| 30 | +from policyengine_us_data.datasets.cps.tipped_occupation import ( |
| 31 | + derive_treasury_tipped_occupation_code, |
| 32 | +) |
30 | 33 | from policyengine_us_data.datasets.org import ORG_IMPUTED_VARIABLES |
31 | 34 |
|
32 | 35 |
|
@@ -288,6 +291,15 @@ def test_se_pension_zeroed_without_se_income( |
288 | 291 | ).all(), "SE pension should be zero without SE income" |
289 | 292 |
|
290 | 293 |
|
| 294 | +class TestTreasuryTippedOccupationCode: |
| 295 | + def test_derive_treasury_tipped_occupation_code(self): |
| 296 | + derived = derive_treasury_tipped_occupation_code( |
| 297 | + np.array([4040, 4110, 4230, 2770, -1, 9999]) |
| 298 | + ) |
| 299 | + |
| 300 | + assert derived.tolist() == [101, 102, 304, 208, 0, 0] |
| 301 | + |
| 302 | + |
291 | 303 | class TestSSReconciliation: |
292 | 304 | """Post-processing SS normalization ensures sub-components sum to total.""" |
293 | 305 |
|
@@ -536,6 +548,7 @@ def test_clone_feature_imputation_rematches_outputs_and_derives_flags( |
536 | 548 | "cps_race": [2, 1], |
537 | 549 | "is_hispanic": [0, 1], |
538 | 550 | "detailed_occupation_recode": [8, 41], |
| 551 | + "treasury_tipped_occupation_code": [101, 304], |
539 | 552 | } |
540 | 553 | ) |
541 | 554 |
|
@@ -573,5 +586,7 @@ def calculate_dataframe(self, columns): |
573 | 586 | assert result["is_male"].tolist() == [1, 0] |
574 | 587 | assert result["cps_race"].tolist() == [2, 1] |
575 | 588 | assert result["is_hispanic"].tolist() == [0, 1] |
| 589 | + if "treasury_tipped_occupation_code" in result.columns: |
| 590 | + assert result["treasury_tipped_occupation_code"].tolist() == [101, 304] |
576 | 591 | assert result["is_computer_scientist"].tolist() == [True, False] |
577 | 592 | assert result["is_farmer_fisher"].tolist() == [False, True] |
0 commit comments