"""Unit tests for the parallel dataset builder (initial baseline-surrogate schema). These tests use :func:`build_dataset` with ``n_workers=1`` for reproducibility and to avoid multiprocessing fork/spawn overhead in CI. A small dedicated parallel-smoke test exercises the spawn path. """ from __future__ import annotations from pathlib import Path import numpy as np import pandas as pd import pytest from roverdevkit.surrogate.dataset import ( SCHEMA_VERSION, DatasetMetadata, build_dataset, read_parquet, read_parquet_metadata, write_parquet, ) from roverdevkit.surrogate.sampling import generate_samples # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @pytest.fixture(scope="module") def small_df() -> pd.DataFrame: """Shared tiny dataset; one evaluator run per scenario family.""" samples = generate_samples(n_per_scenario=2, seed=13) return build_dataset(samples, n_workers=1, progress=False) # --------------------------------------------------------------------------- # Schema: columns and dtypes # --------------------------------------------------------------------------- _EXPECTED_META_COLS = { "sample_index", "split", "stratum_id", "fidelity", "status", } _EXPECTED_DESIGN_COLS = { "design_wheel_radius_m", "design_wheel_width_m", "design_grouser_height_m", "design_grouser_count", "design_n_wheels", "design_chassis_mass_kg", "design_wheelbase_m", "design_solar_area_m2", "design_battery_capacity_wh", "design_avionics_power_w", "design_peak_wheel_torque_nm", } _EXPECTED_SCENARIO_COLS = { "scenario_family", "scenario_name", "scenario_latitude_deg", "scenario_traverse_distance_m", "scenario_terrain_class", "scenario_soil_simulant", "scenario_mission_duration_earth_days", "scenario_max_slope_deg", "scenario_operational_duty_cycle", "scenario_sun_geometry", "scenario_soil_n", "scenario_soil_k_c", "scenario_soil_k_phi", "scenario_soil_cohesion_kpa", "scenario_soil_friction_angle_deg", "scenario_soil_shear_modulus_k_m", "scenario_payload_mass_kg", "scenario_payload_power_w", } _EXPECTED_METRIC_COLS = { "range_km", "energy_margin_pct", "energy_margin_raw_pct", "slope_capability_deg", "total_mass_kg", "peak_motor_torque_nm", "sinkage_max_m", "stalled", } # thermal_survival was removed at SCHEMA_VERSION = v2; the surrogate no # longer consumes it (the system-level evaluator still computes it as a # diagnostic). See data/analytical/SCHEMA.md for the rationale. _THERMAL_REMOVED_COLS = {"thermal_survival"} def test_row_count_matches_sample_count(small_df: pd.DataFrame) -> None: assert len(small_df) == 8 # 2 per scenario × 4 scenarios def test_expected_columns_present(small_df: pd.DataFrame) -> None: cols = set(small_df.columns) for expected in ( _EXPECTED_META_COLS, _EXPECTED_DESIGN_COLS, _EXPECTED_SCENARIO_COLS, _EXPECTED_METRIC_COLS, ): missing = expected - cols assert not missing, f"missing columns: {missing}" def test_stat_columns_present(small_df: pd.DataFrame) -> None: stat_cols = {c for c in small_df.columns if c.startswith("stat_")} # At least 24 stat columns (20 numeric + 3 bool + 1 categorical reason). assert len(stat_cols) >= 24, stat_cols def test_categorical_columns_are_categorical(small_df: pd.DataFrame) -> None: for col in [ "split", "scenario_family", "scenario_name", "scenario_terrain_class", "scenario_soil_simulant", "scenario_sun_geometry", "fidelity", "status", "stat_terminated_reason", ]: assert isinstance(small_df[col].dtype, pd.CategoricalDtype), col def test_design_n_wheels_is_4_or_6(small_df: pd.DataFrame) -> None: assert set(small_df["design_n_wheels"].unique()) <= {4, 6} def test_rows_ordered_by_sample_index(small_df: pd.DataFrame) -> None: idx = small_df["sample_index"].to_numpy() assert np.all(idx == np.sort(idx)) def test_all_rows_succeeded_on_happy_path(small_df: pd.DataFrame) -> None: assert (small_df["status"] == "ok").all() # --------------------------------------------------------------------------- # Metric sanity # --------------------------------------------------------------------------- def test_metric_ranges_are_physically_plausible(small_df: pd.DataFrame) -> None: ok = small_df[small_df["status"] == "ok"] assert (ok["range_km"] >= 0).all() assert (ok["range_km"] < 1000).all() # sanity ceiling assert (ok["energy_margin_pct"] >= 0).all() assert (ok["energy_margin_pct"] <= 100).all() assert (ok["slope_capability_deg"] >= 0).all() assert (ok["slope_capability_deg"] <= 90).all() assert (ok["total_mass_kg"] > 0).all() def test_stalled_is_boolean(small_df: pd.DataFrame) -> None: assert small_df["stalled"].dtype == bool def test_thermal_survival_not_in_schema(small_df: pd.DataFrame) -> None: """v2 schema removes thermal_survival from the dataset (see SCHEMA.md).""" for col in _THERMAL_REMOVED_COLS: assert col not in small_df.columns, ( f"{col} should not be in the v2 schema; the surrogate does not " "predict thermal until the mass model charges RHU/MLI mass." ) def test_stat_columns_are_not_all_nan_on_ok_rows(small_df: pd.DataFrame) -> None: ok = small_df[small_df["status"] == "ok"] for col in [ "stat_power_out_mean_w", "stat_mobility_power_max_w", "stat_soc_final", ]: assert ok[col].notna().all(), col # --------------------------------------------------------------------------- # Parquet round-trip # --------------------------------------------------------------------------- def test_parquet_roundtrip_preserves_schema(tmp_path: Path, small_df: pd.DataFrame) -> None: meta = DatasetMetadata( sampler_seed=13, n_per_scenario=2, scenario_families=("equatorial_mare_traverse",), notes="roundtrip test", ) out_path = tmp_path / "tiny.parquet" write_parquet(small_df, out_path, metadata=meta) assert out_path.exists() loaded = read_parquet(out_path) assert len(loaded) == len(small_df) assert set(loaded.columns) == set(small_df.columns) # Numeric columns equal within tolerance for col in ["range_km", "energy_margin_pct", "total_mass_kg"]: np.testing.assert_allclose( loaded[col].to_numpy(), small_df[col].to_numpy(), rtol=1e-9, atol=0.0 ) def test_parquet_metadata_written_and_read_back(tmp_path: Path, small_df: pd.DataFrame) -> None: meta = DatasetMetadata( sampler_seed=13, n_per_scenario=2, scenario_families=("equatorial_mare_traverse",), notes="metadata test", ) out_path = tmp_path / "tiny.parquet" write_parquet(small_df, out_path, metadata=meta) md = read_parquet_metadata(out_path) assert md["schema_version"] == SCHEMA_VERSION assert md["sampler_seed"] == "13" assert md["notes"] == "metadata test" # --------------------------------------------------------------------------- # Failure handling # --------------------------------------------------------------------------- def test_build_dataset_rejects_empty_input() -> None: with pytest.raises(ValueError, match="No samples"): build_dataset([], n_workers=1, progress=False) def test_evaluator_failure_is_recorded_not_raised(monkeypatch: pytest.MonkeyPatch) -> None: """Inject a failure into evaluate_verbose and check the row is kept with status = exception class name and NaN numeric outputs.""" import roverdevkit.surrogate.dataset as ds_mod def boom(*_args: object, **_kwargs: object) -> None: raise RuntimeError("injected failure") monkeypatch.setattr(ds_mod, "evaluate_verbose", boom) samples = generate_samples(n_per_scenario=2, seed=0) df = build_dataset(samples[:2], n_workers=1, progress=False) assert len(df) == 2 assert (df["status"] == "RuntimeError").all() assert df["range_km"].isna().all() assert df["energy_margin_pct"].isna().all() # Schema v6: failed rows default to ``stalled=True`` (the # safe-conservative side; status != 'ok' filters them out anyway). assert (df["stalled"] == True).all() # noqa: E712 # --------------------------------------------------------------------------- # Parallel smoke test (spawn context) # --------------------------------------------------------------------------- @pytest.mark.slow def test_build_dataset_parallel_matches_serial() -> None: """Run 4 samples both serially and with 2 workers; expect identical outputs.""" samples = generate_samples(n_per_scenario=2, seed=0, scenario_names=["crater_rim_survey"]) serial = build_dataset(samples, n_workers=1, progress=False) parallel = build_dataset(samples, n_workers=2, chunksize=2, progress=False) assert list(serial["sample_index"]) == list(parallel["sample_index"]) np.testing.assert_allclose( serial["range_km"].to_numpy(), parallel["range_km"].to_numpy(), rtol=1e-9, atol=0.0 ) np.testing.assert_allclose( serial["total_mass_kg"].to_numpy(), parallel["total_mass_kg"].to_numpy(), rtol=1e-9, atol=0.0, )