diff --git a/faircode/strategies.py b/faircode/strategies.py index 2ec9342..c47b189 100644 --- a/faircode/strategies.py +++ b/faircode/strategies.py @@ -99,7 +99,9 @@ def strategy_features(strategy: str, core: list, proxies: list, protected: list) return list(dict.fromkeys(core + proxies + protected)) if strategy == "unawareness": return list(dict.fromkeys(core + proxies)) - return list(core) # unawareness_proxy_removal, in_processing, post_processing + if strategy in ("unawareness_proxy_removal", "in_processing", "post_processing"): + return list(dict.fromkeys(core)) + raise ValueError(f"unknown strategy: {strategy!r}") def fit_in_processing(base_model, X_train, y_train, sensitive_train): diff --git a/tests/test_strategies.py b/tests/test_strategies.py index 58178e1..d89b721 100644 --- a/tests/test_strategies.py +++ b/tests/test_strategies.py @@ -65,6 +65,16 @@ def test_strategy_features_deduplicates_overlapping_names(): assert set(cols) == {"a", "b", "c"} +def test_strategy_features_raises_on_unknown_strategy(): + with pytest.raises(ValueError, match="unknown strategy: 'in_procesing'"): + strategy_features("in_procesing", CORE, PROXIES, PROTECTED) + + +def test_strategy_features_raises_on_empty_string(): + with pytest.raises(ValueError, match="unknown strategy: ''"): + strategy_features("", CORE, PROXIES, PROTECTED) + + # ── encode_features ────────────────────────────────────────────────────────── def test_encode_features_passes_numeric_columns_through(): df = pd.DataFrame({"age": [20.0, 30.0, 40.0]})