SummerProject2026/tests/test_baselines.py

560 lines
22 KiB
Python

"""Tests for the private, leakage-safe baseline training contract."""
from __future__ import annotations
import csv
import gzip
import importlib.util
import sys
import tempfile
import unittest
import warnings
from pathlib import Path
from typing import Dict, Iterable, List, Mapping
from unittest import mock
PROJECT_ROOT = Path(__file__).resolve().parents[1]
MODULE_PATH = PROJECT_ROOT / "scripts/train_baselines.py"
SPEC = importlib.util.spec_from_file_location("train_baselines", MODULE_PATH)
if SPEC is None or SPEC.loader is None: # pragma: no cover - import machinery guard
raise RuntimeError("Could not load scripts/train_baselines.py")
baselines = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = baselines
SPEC.loader.exec_module(baselines)
SKLEARN_AVAILABLE = importlib.util.find_spec("numpy") is not None and importlib.util.find_spec(
"sklearn"
) is not None
DUCKDB_AVAILABLE = importlib.util.find_spec("duckdb") is not None
def valid_record(
episode_start: str = "2020-06-01T09:00:00",
partition: str = "train",
target: int = 0,
bucket: int = 25,
token_number: int = 1,
) -> Dict[str, object]:
return {
"vehicle_token": "{:064x}".format(token_number),
"vehicle_bucket": bucket,
"is_vin_audit": "true" if bucket < 10 else "false",
"episode_number": 2,
"episode_start": episode_start,
"first_outcome": "fail" if target else "pass",
"target_nonpass": target,
"eligible_returning_target": "true",
"temporal_partition": partition,
"vehicle_age": 10,
"prior_episode_count": 1,
"prior_total_attempt_count": 1,
"prior_attempt_count": 1,
"days_since_prior_episode": 365,
"days_since_prior_adverse": "",
"prior_nonpass_rate": 0,
"public_county": "salt_lake",
"source_era": "slc",
"target_outcome_label_source": "overall_result",
"target_season": "summer",
"prior_first_outcome": "pass",
"prior_final_outcome": "pass",
"last_observed_make": "toyota",
"last_observed_model": "camry",
}
def write_csv(path: Path, records: Iterable[Mapping[str, object]]) -> None:
rows = list(records)
if not rows:
raise ValueError("Tests require at least one record")
opener = gzip.open if path.name.endswith(".gz") else open
with opener(path, "wt", encoding="utf-8", newline="") as handle:
writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys()))
writer.writeheader()
writer.writerows(rows)
class SchemaTests(unittest.TestCase):
def test_required_schema_is_accepted(self) -> None:
baselines.validate_schema(list(valid_record().keys()))
def test_missing_required_column_fails_closed(self) -> None:
columns = list(valid_record().keys())
columns.remove("prior_first_outcome")
with self.assertRaisesRegex(baselines.SchemaError, "prior_first_outcome"):
baselines.validate_schema(columns)
def test_target_label_source_column_is_required(self) -> None:
columns = list(valid_record().keys())
columns.remove("target_outcome_label_source")
with self.assertRaisesRegex(
baselines.SchemaError, "target_outcome_label_source"
):
baselines.validate_schema(columns)
def test_forbidden_column_fails_even_when_not_a_feature(self) -> None:
columns = list(valid_record().keys()) + ["vin"]
with self.assertRaisesRegex(baselines.SchemaError, "forbidden"):
baselines.validate_schema(columns)
def test_unknown_sensitive_column_fails_closed(self) -> None:
columns = list(valid_record().keys()) + ["owner_zip_code"]
with self.assertRaisesRegex(baselines.SchemaError, "owner_zip_code"):
baselines.validate_schema(columns)
def test_prohibited_sensitive_columns_fail_closed(self) -> None:
for column in (
"technician_id",
"inspection_technician_code",
"inspector_identifier",
"database_password",
"api_secret_value",
"raw_inspection_json",
"row_prediction_probability",
):
with self.subTest(column=column):
with self.assertRaisesRegex(baselines.SchemaError, column):
baselines.validate_schema(list(valid_record().keys()) + [column])
def test_duplicate_columns_fail_closed(self) -> None:
columns = list(valid_record().keys()) + ["vehicle_token"]
with self.assertRaisesRegex(baselines.SchemaError, "duplicate"):
baselines.validate_schema(columns)
class MartValidationTests(unittest.TestCase):
def test_loads_csv_gzip_and_assigns_fixed_partitions(self) -> None:
records = [
valid_record("2022-12-31T23:59:59", "train", 0, 25, 1),
valid_record("2023-01-01T00:00:00", "tune", 1, 25, 2),
valid_record("2024-01-01T00:00:00", "calibrate", 0, 5, 3),
valid_record("2025-01-01T00:00:00", "test", 1, 25, 4),
]
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "mart.csv.gz"
write_csv(path, records)
mart = baselines.load_mart(path)
self.assertEqual(mart.eligible_rows, 4)
self.assertEqual(len(mart.rows_by_partition["train"]), 1)
self.assertEqual(len(mart.rows_by_partition["tune"]), 1)
self.assertEqual(len(mart.rows_by_partition["calibrate"]), 1)
self.assertEqual(len(mart.rows_by_partition["locked_test"]), 1)
self.assertTrue(mart.rows_by_partition["calibrate"][0].audit_vehicle)
@unittest.skipUnless(DUCKDB_AVAILABLE, "duckdb is not installed")
def test_loads_parquet_through_duckdb_without_pyarrow(self) -> None:
import duckdb
with tempfile.TemporaryDirectory() as directory:
csv_path = Path(directory) / "mart.csv"
parquet_path = Path(directory) / "mart.parquet"
write_csv(csv_path, [valid_record()])
connection = duckdb.connect(database=":memory:")
try:
connection.execute(
"CREATE TABLE mart AS SELECT * FROM read_csv_auto(?)", [str(csv_path)]
)
escaped_path = str(parquet_path).replace("'", "''")
connection.execute(
"COPY mart TO '{}' (FORMAT PARQUET)".format(escaped_path)
)
finally:
connection.close()
mart = baselines.load_mart(parquet_path)
self.assertEqual(mart.eligible_rows, 1)
self.assertEqual(len(mart.rows_by_partition["train"]), 1)
def test_declared_partition_must_match_timestamp(self) -> None:
record = valid_record("2025-05-01T00:00:00", "train")
with self.assertRaisesRegex(baselines.DataValidationError, "belongs to locked_test"):
baselines.episode_from_record(record, 2)
def test_label_column_must_match_canonical_outcome(self) -> None:
record = valid_record(target=1)
record["target_nonpass"] = 0
with self.assertRaisesRegex(baselines.DataValidationError, "disagrees"):
baselines.episode_from_record(record, 2)
def test_audit_boolean_must_match_bucket_contract(self) -> None:
record = valid_record(bucket=5)
record["is_vin_audit"] = "false"
with self.assertRaisesRegex(baselines.DataValidationError, "bucket<10"):
baselines.episode_from_record(record, 2)
def test_source_era_is_required_for_audit_but_not_modeling(self) -> None:
record = valid_record()
record["source_era"] = ""
with self.assertRaisesRegex(baselines.DataValidationError, "drift auditing"):
baselines.episode_from_record(record, 2)
def test_utah_proxy_is_accepted_only_for_utah_source_era(self) -> None:
record = valid_record()
record["source_era"] = "utah"
record["target_outcome_label_source"] = "utah_obd_proxy"
episode = baselines.episode_from_record(record, 2)
self.assertEqual(episode.target_outcome_label_source, "utah_obd_proxy")
record["source_era"] = "slc"
with self.assertRaisesRegex(baselines.DataValidationError, "outside source_era"):
baselines.episode_from_record(record, 2)
record["source_era"] = "utah"
record["first_outcome"] = "reject"
record["target_nonpass"] = 1
rejected = baselines.episode_from_record(record, 2)
self.assertEqual(rejected.target, 1)
self.assertEqual(rejected.target_outcome_label_source, "utah_obd_proxy")
record["first_outcome"] = "abort"
aborted = baselines.episode_from_record(record, 2)
self.assertEqual(aborted.target, 1)
self.assertEqual(aborted.target_outcome_label_source, "utah_obd_proxy")
def test_unapproved_target_label_source_fails_closed(self) -> None:
record = valid_record()
record["target_outcome_label_source"] = "inferred_four_class"
with self.assertRaisesRegex(baselines.DataValidationError, "unapproved"):
baselines.episode_from_record(record, 2)
def test_first_episode_cannot_be_eligible_returning_target(self) -> None:
record = valid_record()
record["episode_number"] = 1
with self.assertRaisesRegex(baselines.DataValidationError, "not a returning"):
baselines.episode_from_record(record, 2)
def test_ineligible_row_is_not_parsed_as_a_supervised_target(self) -> None:
record = valid_record()
record["eligible_returning_target"] = "false"
record["first_outcome"] = ""
self.assertIsNone(baselines.episode_from_record(record, 2))
def test_private_output_path_guard(self) -> None:
accepted = PROJECT_ROOT / "artifacts/private/baselines"
rejected = PROJECT_ROOT / "artifacts/public/baselines"
self.assertEqual(baselines.require_private_path(accepted, "output"), accepted)
with self.assertRaisesRegex(baselines.DataValidationError, "must remain"):
baselines.require_private_path(rejected, "output")
class MetricTests(unittest.TestCase):
def test_perfect_ranking_metrics(self) -> None:
targets = [0, 1, 0, 1]
probabilities = [0.1, 0.9, 0.2, 0.8]
metrics = baselines.binary_metrics(targets, probabilities)
self.assertAlmostEqual(metrics["average_precision"], 1.0)
self.assertAlmostEqual(metrics["roc_auc"], 1.0)
self.assertAlmostEqual(metrics["brier"], 0.025)
self.assertAlmostEqual(metrics["top_5_precision"], 1.0)
self.assertAlmostEqual(metrics["top_5_capture"], 0.1)
def test_capacity_metrics_fractionally_handle_boundary_ties(self) -> None:
targets = [0, 1] * 50
probabilities = [0.2] * 100
top_5 = baselines.top_capacity_metrics(targets, probabilities, 0.05)
top_10 = baselines.top_capacity_metrics(targets, probabilities, 0.10)
self.assertAlmostEqual(top_5["precision"], 0.5)
self.assertAlmostEqual(top_5["capture"], 0.05)
self.assertAlmostEqual(top_10["precision"], 0.5)
self.assertAlmostEqual(top_10["capture"], 0.10)
def test_constant_predictions_have_half_auc(self) -> None:
targets = [0, 1, 0, 1]
probabilities = [0.25] * 4
self.assertAlmostEqual(baselines.roc_auc(targets, probabilities), 0.5)
self.assertAlmostEqual(baselines.average_precision(targets, probabilities), 0.5)
def test_calibration_bins_cover_each_row_once(self) -> None:
targets = [0, 1, 0, 1, 1]
probabilities = [0.1, 0.2, 0.3, 0.8, 0.9]
bins = baselines.calibration_bins(targets, probabilities, 3)
self.assertEqual(sum(item["count"] for item in bins), len(targets))
self.assertEqual([item["count"] for item in bins], [2, 2, 1])
class BaselineBehaviorTests(unittest.TestCase):
def test_default_iteration_budget_is_hardened(self) -> None:
args = baselines.parse_args(["--mart", "data/private/example.parquet"])
self.assertEqual(args.max_iter, 5000)
def _episode(
self,
prior_outcome: object,
target: int = 0,
bucket: int = 25,
episode_start: str = "2020-06-01T09:00:00",
partition: str = "train",
source_era: str = "slc",
label_source: str = "overall_result",
):
record = valid_record(
episode_start=episode_start,
partition=partition,
target=target,
bucket=bucket,
)
record["prior_first_outcome"] = prior_outcome
record["source_era"] = source_era
record["target_outcome_label_source"] = label_source
return baselines.episode_from_record(record, 2)
def test_literal_previous_outcome_and_missing_fallback(self) -> None:
rows = [
self._episode("pass"),
self._episode("fail"),
self._episode("reject"),
self._episode(""),
]
probabilities, fallback_count = baselines.literal_previous_probabilities(rows, 0.2)
self.assertEqual(probabilities, [0.0, 1.0, 1.0, 0.2])
self.assertEqual(fallback_count, 1)
def test_audit_rows_are_excluded_from_fit_cohort(self) -> None:
rows = [self._episode("pass", bucket=5), self._episode("pass", bucket=25)]
non_audit = baselines.evaluation_cohorts(rows, include_audit_breakout=False)
self.assertEqual(non_audit[0][0], "non_audit")
self.assertEqual(len(non_audit[0][1]), 1)
self.assertFalse(non_audit[0][1][0].audit_vehicle)
def test_locked_manifest_counts_withhold_labels_until_unlocked(self) -> None:
locked = self._episode(
"pass",
target=1,
bucket=25,
episode_start="2025-06-01T09:00:00",
partition="test",
)
locked_proxy = self._episode(
"pass",
target=0,
bucket=26,
episode_start="2025-06-01T09:00:00",
partition="test",
source_era="utah",
label_source="utah_obd_proxy",
)
mart = baselines.MartData(
rows_by_partition={
"train": [],
"tune": [],
"calibrate": [],
"locked_test": [locked, locked_proxy],
},
input_rows=2,
ineligible_rows=0,
eligible_rows=2,
)
closed, closed_eras, closed_labels = baselines.manifest_audit_counts(
mart, evaluate_locked=False
)
self.assertTrue(closed["locked_test"]["labels_withheld"])
self.assertNotIn("nonpass", closed["locked_test"])
self.assertNotIn("nonpass", closed_eras["locked_test"]["slc"])
self.assertEqual(
closed_labels["locked_test"],
{"overall_result": 1, "utah_obd_proxy": 1},
)
opened, opened_eras, opened_labels = baselines.manifest_audit_counts(
mart, evaluate_locked=True
)
self.assertEqual(opened["locked_test"]["nonpass"], 1)
self.assertEqual(opened_eras["locked_test"]["slc"]["nonpass"], 1)
self.assertEqual(
opened_labels["locked_test"],
{"overall_result": 1, "utah_obd_proxy": 1},
)
self.assertFalse(
baselines.TARGET_CONTRACT["utah_obd_proxy_four_class_approved"]
)
@unittest.skipUnless(SKLEARN_AVAILABLE, "numpy/scikit-learn are not installed")
class ModelingIntegrationTests(unittest.TestCase):
@staticmethod
def _partition_records(
timestamp: str,
partition: str,
start_token: int,
future_category: bool = False,
) -> List[Mapping[str, object]]:
records: List[Mapping[str, object]] = []
for index in range(40):
target = 1 if index % 4 == 0 else 0
bucket = index % 100
record = valid_record(
episode_start=timestamp,
partition=partition,
target=target,
bucket=bucket,
token_number=start_token + index,
)
record["vehicle_age"] = 2 + index % 25
record["prior_episode_count"] = 1 + index % 6
record["prior_total_attempt_count"] = 1 + index % 10
record["prior_attempt_count"] = 1 + index % 3
record["days_since_prior_episode"] = 250 + index * 4
record["days_since_prior_adverse"] = "" if index % 3 else 500 + index
record["prior_nonpass_rate"] = (index % 4) / 4
record["prior_first_outcome"] = "fail" if index % 5 == 0 else "pass"
if future_category:
record["public_county"] = "future_only_county"
records.append(record)
return records
def _synthetic_mart(self):
records: List[Mapping[str, object]] = []
records.extend(self._partition_records("2022-06-01", "train", 1))
records.extend(self._partition_records("2023-06-01", "tune", 101))
records.extend(
self._partition_records(
"2024-06-01", "calibrate", 201, future_category=True
)
)
records.extend(self._partition_records("2025-06-01", "test", 301))
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "mart.csv"
write_csv(path, records)
return baselines.load_mart(path)
def test_train_only_preprocessing_and_locked_evaluation_gate(self) -> None:
mart = self._synthetic_mart()
trained = baselines.train_baselines(
mart,
c_grid=[0.1],
min_category_count=1,
seed=7,
max_iter=2000,
)
feature_names = set(trained.encoder.vectorizer.get_feature_names_out())
self.assertNotIn("cat__public_county=future_only_county", feature_names)
self.assertFalse(any("source_era" in name for name in feature_names))
self.assertFalse(
any("target_outcome_label_source" in name for name in feature_names)
)
locked_metrics, _ = baselines.evaluate_models(
mart, trained, evaluate_locked=False, bin_count=5
)
self.assertNotIn("locked_test", {row["partition"] for row in locked_metrics})
unlocked_metrics, _ = baselines.evaluate_models(
mart, trained, evaluate_locked=True, bin_count=5
)
locked_rows = [
row for row in unlocked_metrics if row["partition"] == "locked_test"
]
self.assertTrue(locked_rows)
self.assertIn("never_fit_audit", {row["cohort"] for row in locked_rows})
self.assertIn("logistic_platt", {row["model"] for row in locked_rows})
tuning = trained.tuning_results[0]
self.assertTrue(tuning["converged"])
self.assertTrue(tuning["finite_coefficients"])
self.assertTrue(tuning["finite_probabilities"])
self.assertTrue(tuning["eligible_for_selection"])
self.assertLess(tuning["n_iter"], trained.max_iter)
self.assertLess(trained.selected_n_iter, trained.max_iter)
self.assertLess(trained.platt_n_iter, trained.max_iter)
self.assertEqual(trained.platt_model.solver, "liblinear")
self.assertEqual(trained.platt_model.penalty, "l2")
self.assertEqual(trained.platt_model.C, 1000.0)
self.assertEqual(
baselines.PLATT_CALIBRATION_CONFIG,
{
"method": "sigmoid_platt_scaling",
"solver": "liblinear",
"penalty": "l2",
"c": 1000.0,
},
)
def test_deliberately_tiny_max_iter_rejects_all_candidates(self) -> None:
mart = self._synthetic_mart()
with self.assertRaisesRegex(
baselines.DataValidationError, "No logistic candidate converged"
):
baselines.train_baselines(
mart,
c_grid=[0.03, 0.1],
min_category_count=1,
seed=7,
max_iter=1,
)
def test_nonconverged_and_nonfinite_candidates_are_excluded(self) -> None:
import numpy as np
import sklearn
from sklearn.exceptions import ConvergenceWarning
from sklearn.feature_extraction import DictVectorizer
class ControlledLogistic:
def __init__(self, C=1.0, solver="lbfgs", max_iter=100, **_kwargs):
self.C = C
self.solver = solver
self.max_iter = max_iter
def fit(self, features, _targets):
self.n_iter_ = np.asarray([2])
self.coef_ = np.zeros((1, features.shape[1]))
self.intercept_ = np.zeros(1)
if self.solver == "saga" and self.C == 0.03:
warnings.warn("controlled nonconvergence", ConvergenceWarning)
if self.solver == "saga" and self.C == 0.1:
self.coef_[0, 0] = np.nan
return self
def predict_proba(self, features):
probability = np.full(features.shape[0], 0.25)
return np.column_stack((1.0 - probability, probability))
def decision_function(self, features):
return np.zeros(features.shape[0])
dependencies = (
np,
sklearn,
DictVectorizer,
ControlledLogistic,
ConvergenceWarning,
)
mart = self._synthetic_mart()
with mock.patch.object(
baselines, "require_ml_dependencies", return_value=dependencies
):
trained = baselines.train_baselines(
mart,
c_grid=[0.03, 0.1, 1.0],
min_category_count=1,
seed=7,
max_iter=100,
)
self.assertEqual(trained.selected_c, 1.0)
by_c = {result["c"]: result for result in trained.tuning_results}
self.assertFalse(by_c[0.03]["converged"])
self.assertFalse(by_c[0.03]["eligible_for_selection"])
self.assertFalse(by_c[0.1]["finite_coefficients"])
self.assertFalse(by_c[0.1]["eligible_for_selection"])
self.assertTrue(by_c[1.0]["eligible_for_selection"])
def test_convergence_capture_replays_unrelated_warnings(self) -> None:
from sklearn.exceptions import ConvergenceWarning
class WarningModel:
def fit(self, _features, _targets):
warnings.warn("unrelated diagnostic", RuntimeWarning)
warnings.warn("did not converge", ConvergenceWarning)
with self.assertWarnsRegex(RuntimeWarning, "unrelated diagnostic"):
messages = baselines.fit_with_convergence_capture(
WarningModel(), None, None, ConvergenceWarning
)
self.assertEqual(messages, ["did not converge"])
if __name__ == "__main__":
unittest.main()