6krrt/tests/test_apply_tiering.py

"""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 / "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.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