6krrt/tests/test_seed_sweep.py

"""End-to-end test of the reference sweep, with the provider stubbed.

The sweep is what makes `eco` and the per-request energy ceiling real numbers
instead of the neutral 0.5, and `llm-router-seed.timer` runs it every six
hours. It had been dead for weeks: `log_observation` grew `request_id`,
`session_key` and `session_dir` in the middle of its signature and made the
trailing three keyword-only, while this call still passed six positionals. The
result was a TypeError -- not a RequestException, so the per-sample `except`
never caught it -- raised AFTER the first billed completion. Every scheduled
run spent money and wrote nothing.

So this is a smoke test of `main()` rather than an assertion about a
signature: it catches this class of drift wherever it next appears, and it
costs nothing to run because `sample_once` never leaves the process.
"""

import sqlite3
import sys
from pathlib import Path

import pytest

import dispatcher
import seed_energy

ROOT = Path(__file__).resolve().parent.parent
SCHEMA_SQL = (ROOT / "schema.sql").read_text()

MODEL = "gemma-4-31b"


def _payload(model_id):
    return {
        "id": "chatcmpl-seed-1",
        "model": model_id,
        "choices": [{"message": {"content": "A B-tree is..."},
                     "finish_reason": "length"}],
        "usage": {"prompt_tokens": 31, "completion_tokens": 400},
        "energy": {"energy_kwh": 4.75e-05, "avg_power_watts": 420.0,
                   "duration_seconds": 1.6, "attribution_ratio": 0.25,
                   "carbon_g_co2eq": 2.32e-04, "carbon_source": "agent_cache",
                   "grid_id": "FI"},
        "cost": {"request_cost_usd": 3.8e-04, "allowance_remaining_usd": 49.5},
    }


@pytest.fixture
def sweep(tmp_path, monkeypatch):
    db_path = tmp_path / "test.db"
    conn = sqlite3.connect(db_path)
    conn.executescript(SCHEMA_SQL)
    conn.execute(
        """
        INSERT INTO models (
            model_id, provider, tier, latency_class, access_level,
            availability, cost_per_1m_completion, last_updated
        ) VALUES (?, 'neuralwatt', 1, 'standard', 'public', 'active', 0.42,
                  '2026-08-22T00:00:00+00:00')
        """,
        (MODEL,),
    )
    conn.commit()
    conn.close()

    # main() loads its own config; log_observation reads dispatcher's. They must
    # agree on the database or the sweep writes somewhere the test cannot see.
    monkeypatch.setattr(dispatcher.cfg.database, "path", str(db_path))
    monkeypatch.setattr(seed_energy, "load_config", lambda path: dispatcher.cfg)
    monkeypatch.setenv("NEURALWATT_API_KEY", "test-key")
    monkeypatch.setattr(seed_energy.time, "sleep", lambda _: None)

    calls = []

    def fake_sample(base_url, api_key, model_id, max_tokens=400, timeout=300):
        calls.append({"model_id": model_id, "max_tokens": max_tokens})
        return _payload(model_id)

    monkeypatch.setattr(seed_energy, "sample_once", fake_sample)
    yield calls, db_path


def _run(monkeypatch, *argv):
    monkeypatch.setattr(sys, "argv", ["seed_energy.py", *argv])
    return seed_energy.main()


def test_a_sweep_writes_the_observation_it_paid_for(sweep, monkeypatch):
    """The regression: this raised TypeError after the first billed call."""
    calls, db_path = sweep

    assert _run(monkeypatch, "--samples", "2") == 0
    assert len(calls) == 2

    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    rows = conn.execute(
        "SELECT * FROM energy_observations WHERE task_category = ?",
        (seed_energy.SEED_CATEGORY,),
    ).fetchall()
    conn.close()

    assert len(rows) == 2
    row = rows[0]
    assert row["model_id"] == MODEL
    assert row["prompt_tokens"] == 31
    assert row["completion_tokens"] == 400
    assert row["energy_kwh"] == pytest.approx(4.75e-05)
    assert row["carbon_g_co2eq"] == pytest.approx(2.32e-04)
    assert row["cost_usd"] == pytest.approx(3.8e-04)


def test_the_sweep_records_the_completion_id(sweep, monkeypatch):
    """Positional drift silently dropped request_id; it is the join key for /outcome."""
    _, db_path = sweep
    _run(monkeypatch, "--samples", "1")

    conn = sqlite3.connect(db_path)
    request_id = conn.execute(
        "SELECT request_id FROM energy_observations"
    ).fetchone()[0]
    conn.close()

    assert request_id == "chatcmpl-seed-1"


def test_max_tokens_reaches_the_request(sweep, monkeypatch):
    """The flag was parsed, printed in the banner, and never sent."""
    calls, _ = sweep
    _run(monkeypatch, "--samples", "1", "--max-tokens", "800")

    assert calls[0]["max_tokens"] == 800


def test_a_dry_run_spends_nothing(sweep, monkeypatch):
    calls, db_path = sweep

    assert _run(monkeypatch, "--dry-run") == 0
    assert not calls

    conn = sqlite3.connect(db_path)
    assert conn.execute("SELECT COUNT(*) FROM energy_observations").fetchone()[0] == 0
    conn.close()