177 lines
6.8 KiB
Python
177 lines
6.8 KiB
Python
"""Integration tests for tier.py — apply_tiering DB upsert.
|
|
|
|
Seeds a temp sqlite DB from schema.sql, inserts synthetic model rows with
|
|
known expected tiers, and asserts the tiering pass writes the right tier
|
|
for every row (no NULLs), honors the override map, matches the heuristic,
|
|
is idempotent, and warns when reasoning_default_enabled is uniformly False.
|
|
"""
|
|
|
|
import sqlite3
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from config import RouterConfig, TieringConfig, load_config
|
|
from tier import apply_tiering
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
SCHEMA_SQL = (ROOT / "config" / "schema.sql").read_text()
|
|
|
|
|
|
def _make_config(model_tiers: dict[str, int] | None = None) -> RouterConfig:
|
|
"""Real config.yaml with a synthetic tiering block (override map)."""
|
|
cfg = load_config(ROOT / "config" / "config.yaml")
|
|
return cfg.model_copy(
|
|
update={
|
|
"tiering": TieringConfig(
|
|
cheap_completion_max=1.00,
|
|
model_tiers=model_tiers or {},
|
|
)
|
|
}
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def db(tmp_path):
|
|
"""Empty temp sqlite DB created from schema.sql."""
|
|
conn = sqlite3.connect(tmp_path / "test.db")
|
|
conn.executescript(SCHEMA_SQL)
|
|
yield conn
|
|
conn.close()
|
|
|
|
|
|
def _insert(conn: sqlite3.Connection, rows: list[dict]) -> None:
|
|
"""Insert synthetic model rows (only the columns the tiering pass reads).
|
|
|
|
``supports_reasoning`` defaults to mirroring ``reasoning_default_enabled``
|
|
so callers only specify the signal under test; the fallback path is
|
|
exercised by setting them independently.
|
|
"""
|
|
for r in rows:
|
|
default_enabled = r["reasoning_default_enabled"]
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO models (
|
|
model_id, provider, supports_reasoning, reasoning_default_enabled,
|
|
reasoning_mode, cost_per_1m_completion, pricing_tbd, last_updated
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
r["model_id"],
|
|
r["provider"],
|
|
int(r.get("supports_reasoning", default_enabled)),
|
|
int(default_enabled),
|
|
r.get("reasoning_mode", "default"),
|
|
r["cost_per_1m_completion"],
|
|
int(r["pricing_tbd"]),
|
|
"2026-08-07T00:00:00+00:00",
|
|
),
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def _tiers(conn: sqlite3.Connection) -> dict[tuple[str, str], int | None]:
|
|
return {
|
|
(model_id, provider): tier
|
|
for model_id, provider, tier in conn.execute(
|
|
"SELECT model_id, provider, tier FROM models"
|
|
)
|
|
}
|
|
|
|
|
|
# --- required test 1: no NULLs, every tier in {1, 2, 3} -------------------
|
|
|
|
def test_every_row_gets_a_tier_in_1_2_3(db):
|
|
# Given: a mixed catalog covering every heuristic branch
|
|
_insert(db, [
|
|
{"model_id": "reasoner", "provider": "p", "reasoning_default_enabled": True,
|
|
"cost_per_1m_completion": 5.0, "pricing_tbd": False},
|
|
{"model_id": "cheap", "provider": "p", "reasoning_default_enabled": False,
|
|
"cost_per_1m_completion": 0.25, "pricing_tbd": False},
|
|
{"model_id": "mid", "provider": "p", "reasoning_default_enabled": False,
|
|
"cost_per_1m_completion": 5.0, "pricing_tbd": False},
|
|
{"model_id": "null-cost", "provider": "p", "reasoning_default_enabled": False,
|
|
"cost_per_1m_completion": None, "pricing_tbd": False},
|
|
{"model_id": "tbd", "provider": "p", "reasoning_default_enabled": False,
|
|
"cost_per_1m_completion": 0.10, "pricing_tbd": True},
|
|
])
|
|
# When: applying the tiering pass
|
|
apply_tiering(db, _make_config())
|
|
# Then: every row has a tier in {1, 2, 3} and none are NULL
|
|
tiers = _tiers(db)
|
|
assert len(tiers) == 5
|
|
assert all(tier in (1, 2, 3) for tier in tiers.values())
|
|
|
|
|
|
# --- required test 2: override-map rows match the override ----------------
|
|
|
|
def test_override_map_row_matches_override(db):
|
|
# Given: a reasoning model (heuristic would say 3) overridden to tier 1
|
|
_insert(db, [
|
|
{"model_id": "deep-reasoner", "provider": "p", "reasoning_default_enabled": True,
|
|
"cost_per_1m_completion": 5.0, "pricing_tbd": False},
|
|
])
|
|
# When: applying tiering with an override map
|
|
apply_tiering(db, _make_config(model_tiers={"deep-reasoner": 1}))
|
|
# Then: the override wins over the heuristic
|
|
assert _tiers(db)[("deep-reasoner", "p")] == 1
|
|
|
|
|
|
# --- required test 3: heuristic rows match expected -----------------------
|
|
|
|
def test_heuristic_rows_match_expected_tiers(db):
|
|
# Given: reasoning -> 3, cheap non-reasoning -> 1, everything else -> 2
|
|
_insert(db, [
|
|
{"model_id": "reasoner", "provider": "p", "reasoning_default_enabled": True,
|
|
"cost_per_1m_completion": 5.0, "pricing_tbd": False},
|
|
{"model_id": "cheap", "provider": "p", "reasoning_default_enabled": False,
|
|
"cost_per_1m_completion": 0.25, "pricing_tbd": False},
|
|
{"model_id": "mid", "provider": "p", "reasoning_default_enabled": False,
|
|
"cost_per_1m_completion": 5.0, "pricing_tbd": False},
|
|
])
|
|
# When: applying the tiering pass
|
|
apply_tiering(db, _make_config())
|
|
# Then: each row matches its expected heuristic tier
|
|
tiers = _tiers(db)
|
|
assert tiers[("reasoner", "p")] == 3
|
|
assert tiers[("cheap", "p")] == 1
|
|
assert tiers[("mid", "p")] == 2
|
|
|
|
|
|
# --- required test 4: re-running is idempotent ----------------------------
|
|
|
|
def test_rerun_is_idempotent(db):
|
|
# Given: a mixed catalog
|
|
_insert(db, [
|
|
{"model_id": "reasoner", "provider": "p", "reasoning_default_enabled": True,
|
|
"cost_per_1m_completion": 5.0, "pricing_tbd": False},
|
|
{"model_id": "cheap", "provider": "p", "reasoning_default_enabled": False,
|
|
"cost_per_1m_completion": 0.25, "pricing_tbd": False},
|
|
{"model_id": "mid", "provider": "p", "reasoning_default_enabled": False,
|
|
"cost_per_1m_completion": 5.0, "pricing_tbd": False},
|
|
])
|
|
# When: applying the tiering pass twice
|
|
apply_tiering(db, _make_config())
|
|
first = _tiers(db)
|
|
apply_tiering(db, _make_config())
|
|
second = _tiers(db)
|
|
# Then: tiers are identical across runs
|
|
assert first == second
|
|
|
|
|
|
# --- required test 5: sanity-guard warning path ---------------------------
|
|
|
|
def test_warns_when_reasoning_default_enabled_uniformly_false(db, capsys):
|
|
# Given: a catalog where no row reasons by default
|
|
_insert(db, [
|
|
{"model_id": "cheap", "provider": "p", "reasoning_default_enabled": False,
|
|
"cost_per_1m_completion": 0.25, "pricing_tbd": False},
|
|
{"model_id": "mid", "provider": "p", "reasoning_default_enabled": False,
|
|
"cost_per_1m_completion": 5.0, "pricing_tbd": False},
|
|
])
|
|
# When: applying the tiering pass
|
|
apply_tiering(db, _make_config())
|
|
# Then: a clear warning is emitted (not a hard failure)
|
|
captured = capsys.readouterr()
|
|
assert "reasoning_default_enabled is uniformly False" in captured.err
|