Function test_split_stratified

tests/test_data_split.py:47–54

()

def test_split_stratified():
    rows = _make_rows()
    train, cal, holdout = stratified_3way_split(rows, seed=42)
    train_bks = {r["benchmark"] for r in train}
    holdout_bks = {r["benchmark"] for r in holdout}
    assert "bk1" in train_bks
    assert "bk2" in train_bks
    assert "bk1" in holdout_bks