""" Tests for Week 2 feature engineering. Run: python -m pytest tests/test_feature_engineering.py -v """ import numpy as np import pandas as pd import pytest from ml.feature_engineering import ( add_driver_form, add_era_normalized_sector_times, create_time_series_splits, load_encoders, prepare_model_data, save_encoders, validate_input_schema, ) @pytest.fixture def sample_race_df(): return pd.DataFrame( { "season": [2021, 2021, 2022, 2022, 2023, 2023, 2024, 2024], "round": [1, 2, 1, 2, 1, 2, 1, 2], "driver_id": ["VER", "HAM", "VER", "HAM", "VER", "HAM", "VER", "HAM"], "driver_name": [ "Max Verstappen", "Lewis Hamilton", "Max Verstappen", "Lewis Hamilton", "Max Verstappen", "Lewis Hamilton", "Max Verstappen", "Lewis Hamilton", ], "team": [ "Red Bull Racing", "Mercedes", "Red Bull Racing", "Mercedes", "Red Bull Racing", "Mercedes", "Red Bull Racing", None, ], "grid_position": [1, 2, 1, 3, 1, 4, 2, 1], "finish_position": [1, 2, 1, 3, 2, 1, 1, 2], "circuit_id": [ "Monaco Grand Prix", "Silverstone Grand Prix", "Bahrain Grand Prix", "Australian Grand Prix", "Singapore Grand Prix", "Italian Grand Prix", "Canadian Grand Prix", "Las Vegas Grand Prix", ], "regulation_era": [ "hybrid_era", "hybrid_era", "ground_effect_era", "ground_effect_era", "ground_effect_era", "ground_effect_era", "ground_effect_era", "ground_effect_era", ], "track_type": [ "street", "permanent", "permanent", "hybrid", "street", "permanent", "hybrid", "street", ], "sector_1_time": [28.0, np.nan, 29.0, 30.0, 28.5, 29.2, 28.8, 29.1], "sector_2_time": [35.0, 35.5, 36.0, np.nan, 35.8, 36.1, 35.6, 35.9], "sector_3_time": [22.0, 22.2, np.nan, 22.8, 22.5, 22.7, 22.4, 22.6], "avg_lap_time_practice": [93.0, np.nan, 94.0, 95.0, 93.5, 94.2, 93.8, 94.1], "tire_compound": [0, 1, 0, 1, np.nan, 1, 2, 0], "tire_age_laps": [14, 12, 18, 16, 20, np.nan, 15, 17], "fresh_tire": [1, 1, 1, 0, 1, 1, np.nan, 0], "pit_stop_count": [2, 1, 2, 2, 3, 2, 1, np.nan], "team_pit_speed": [np.nan] * 8, "weather_temp_track": [30.0, 28.0, 31.0, 27.0, 29.0, 32.0, 26.0, 25.0], "weather_rainfall": [0, 0, 0, 1, 0, 0, 1, 0], "championship_standing": [1, 2, 1, 3, 1, 2, 1, 2], "driver_form_last3": [np.nan, np.nan, 1.0, 2.0, 1.0, 2.5, 1.3, 2.3], "safety_car_probability": [0.2, 0.0, 0.1, 0.3, 0.4, 0.1, 0.2, 0.5], "telemetry_available": [True] * 8, "win_probability": [1, 0, 1, 0, 0, 1, 1, 0], } ) def test_validate_input_schema_rejects_missing_required_column(sample_race_df): broken = sample_race_df.drop(columns=["driver_id"]) with pytest.raises(ValueError, match="driver_id"): validate_input_schema(broken) def test_prepare_model_data_imputes_missing_values(sample_race_df): bundle, encoders = prepare_model_data(sample_race_df) assert not bundle.X.isna().any().any() assert len(bundle.X) == len(sample_race_df) assert len(bundle.y) == len(sample_race_df) assert set(encoders) == {"team", "track_type", "regulation_era"} def test_prepare_model_data_excludes_leakage_columns(sample_race_df): bundle, _ = prepare_model_data(sample_race_df) assert "finish_position" not in bundle.feature_columns assert "driver_id" not in bundle.feature_columns assert "win_probability" not in bundle.feature_columns assert "round" not in bundle.feature_columns assert "round" in bundle.metadata.columns def test_prepare_model_data_label_encodes_categoricals(sample_race_df): bundle, encoders = prepare_model_data(sample_race_df) assert "team" in bundle.feature_columns assert "track_type" in bundle.feature_columns assert "regulation_era" in bundle.feature_columns assert "unknown" in encoders["team"].classes_ assert pd.api.types.is_integer_dtype(bundle.X["team"]) def test_era_normalized_sector_times_center_within_era(sample_race_df): normalized = add_era_normalized_sector_times(sample_race_df) for era, group in normalized.groupby("regulation_era"): mean_value = group["sector_1_time_era_norm"].mean() assert mean_value == pytest.approx(0.0, abs=1e-12), era def test_time_series_splits_are_chronological(sample_race_df): bundle, _ = prepare_model_data(sample_race_df) splits = list(create_time_series_splits(bundle.X, n_splits=3)) assert len(splits) == 3 for train_idx, test_idx in splits: assert train_idx.max() < test_idx.min() def test_prepare_model_data_can_accept_fastf1_only_shape(sample_race_df): fastf1_only = sample_race_df.drop( columns=[ "regulation_era", "track_type", "championship_standing", "driver_form_last3", "safety_car_probability", "win_probability", ] ) bundle, _ = prepare_model_data(fastf1_only) assert "win_probability" not in bundle.feature_columns assert set(bundle.y.unique()) <= {0, 1} assert "safety_car_probability" in bundle.feature_columns def test_prepare_model_data_reuses_fitted_encoders_for_inference(sample_race_df): train_df = sample_race_df.iloc[:6].copy() inference_df = sample_race_df.iloc[6:].copy() inference_df.loc[:, "team"] = "New Team" _, encoders = prepare_model_data(train_df) inference_bundle, reused_encoders = prepare_model_data(inference_df, encoders) unknown_code = encoders["team"].transform(["unknown"])[0] assert reused_encoders is encoders assert (inference_bundle.X["team"] == unknown_code).all() def test_add_driver_form_uses_only_prior_races(): df = pd.DataFrame( { "driver_id": ["VER", "VER", "VER", "VER"], "season": [2023, 2023, 2023, 2023], "round": [1, 2, 3, 4], "finish_position": [1, 3, 5, 7], } ) result = add_driver_form(df) assert pd.isna(result.loc[0, "driver_form_last3"]) assert result.loc[1, "driver_form_last3"] == pytest.approx(1.0) assert result.loc[2, "driver_form_last3"] == pytest.approx(2.0) assert result.loc[3, "driver_form_last3"] == pytest.approx(3.0) def test_prepare_model_data_computes_driver_form_only_when_missing(sample_race_df): missing_form = sample_race_df.copy() missing_form["driver_form_last3"] = np.nan bundle, _ = prepare_model_data(missing_form) assert not bundle.X["driver_form_last3"].isna().any() def test_championship_standing_uses_season_max_fill(sample_race_df): df = sample_race_df.copy() df.loc[(df["season"] == 2023) & (df["driver_id"] == "HAM"), "championship_standing"] = np.nan bundle, _ = prepare_model_data(df) row = bundle.metadata[ (bundle.metadata["season"] == 2023) & (bundle.metadata["driver_id"] == "HAM") ].index[0] assert bundle.X.loc[row, "championship_standing"] == 1 def test_encoder_persistence_round_trip(sample_race_df, tmp_path): _, encoders = prepare_model_data(sample_race_df) path = tmp_path / "label_encoders.pkl" save_encoders(encoders, str(path)) loaded = load_encoders(str(path)) assert set(loaded) == set(encoders) assert list(loaded["team"].classes_) == list(encoders["team"].classes_)