546 lines
22 KiB
Python
546 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_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()
|