6krrt/tests/test_load_candidates.py

"""Tests for dispatcher.load_candidates — how measured data reaches scoring.

Every behaviour here was found by a wrong routing decision rather than by
reasoning, so each gets explicit coverage:

- scoring reads the provider's ATTRIBUTED figures, because the attribution
  ratio turned out to be a stable per-model property (750x between models,
  ~1.8x within one) rather than the noise it resembles up close
- the median, not the mean, because a few models still throw 50-90x outliers
- carbon reported as a static fallback is not data and must not rank anything
"""

import sqlite3
from pathlib import Path

import pytest

from dispatcher import (
    FALLBACK_CARBON_SOURCE,
    SEED_CATEGORY,
    USD_PER_KWH,
    gross_energy_kwh,
    load_candidates,
)

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

# 1000 W for 3.6 s = 1e-3 kWh, so the arithmetic below stays readable.
WATTS = 1000.0
SECONDS = 3.6
GROSS_KWH = 1e-3


@pytest.fixture
def db(tmp_path):
    conn = sqlite3.connect(tmp_path / "test.db")
    conn.row_factory = sqlite3.Row
    conn.executescript(SCHEMA_SQL)
    conn.execute(
        """
        INSERT INTO models (model_id, provider, tier, availability, last_updated)
        VALUES ('m', 'neuralwatt', 2, 'active', '2026-08-12T00:00:00+00:00')
        """
    )
    conn.commit()
    yield conn
    conn.close()


def _observe(conn, *, cost=1.0, carbon=37.0, source="agent_cache", category=SEED_CATEGORY):
    conn.execute(
        """
        INSERT INTO energy_observations (
            model_id, provider, task_category, cost_usd,
            carbon_g_co2eq, carbon_source, observed_at
        ) VALUES ('m', 'neuralwatt', ?, ?, ?, ?, '2026-08-12T00:00:00+00:00')
        """,
        (category, cost, carbon, source),
    )
    conn.commit()


def _only(conn):
    return load_candidates(conn, "coding_general")[0]


# --- attribution is signal, not noise -------------------------------------

def test_gross_energy_identity_holds():
    # The decomposition behind the finding: billed = power x duration x
    # attribution. 1260.9 W x 1.523 s x 0.0158 = 8.43e-06 kWh, and the API
    # reported 8.438e-06. gross_energy_kwh is a diagnostic on that identity.
    assert gross_energy_kwh(1000.0, 3.6) == pytest.approx(1e-3)
    assert gross_energy_kwh(1260.9, 1.523) * 0.0158 == pytest.approx(8.43e-06, rel=1e-2)


def test_scoring_reads_billed_cost_not_pre_attribution_energy(db):
    # Given: a model whose billed cost is far below its share of pool gross,
    # because many concurrent requests share its GPUs — deepseek-v4-flash
    # bills ~1000x under gross. That discount is real money, so it must reach
    # the score rather than being normalized away.
    _observe(db, cost=1.154e-06)
    _observe(db, cost=1.154e-06)
    assert _only(db)["cost"] == pytest.approx(1.154e-06)


# --- median, not mean -----------------------------------------------------

def test_cost_uses_the_median_so_one_spike_cannot_dominate(db):
    # Given: kimi-k3's real shape — ordinary samples plus a ~90x outlier
    for c in (1.0, 1.0, 1.0, 1.0, 90.0):
        _observe(db, cost=c)
    assert _only(db)["cost"] == 1.0


def test_eco_uses_the_median_too(db):
    for carbon in (1.0, 2.0, 3.0, 4.0, 500.0):
        _observe(db, cost=1.0, carbon=carbon)
    assert _only(db)["eco"] == 3.0


# --- fabricated carbon is not data ----------------------------------------

def test_static_fallback_carbon_is_excluded(db):
    # Given: the glm-5.2-fast case. NeuralWatt could not resolve live grid
    # data and substituted 475.0 while still reporting grid_id 'FI'.
    for _ in range(3):
        _observe(db, cost=1.0, carbon=8.5e-03, source=FALLBACK_CARBON_SOURCE)
    row = _only(db)
    # Then: no eco figure at all rather than a placeholder one. scoring turns
    # that into a neutral 0.5, which is honest; ranking a model against a
    # constant nobody measured would not be.
    assert row["eco"] is None
    # Cost is unaffected — it comes from the billing block, not the estimate.
    assert row["cost"] == 1.0


def test_measured_carbon_survives_alongside_fallback_rows(db):
    _observe(db, cost=1.0, carbon=100.0, source=FALLBACK_CARBON_SOURCE)
    _observe(db, cost=1.0, carbon=2.0)
    _observe(db, cost=1.0, carbon=4.0)
    assert _only(db)["eco"] == 3.0


# --- scope ----------------------------------------------------------------

def test_only_reference_workload_observations_are_used(db):
    # Given: organic traffic, whose energy tracks request shape far more than
    # model efficiency (19x spread on one model), alongside the sweep
    _observe(db, cost=1.0)
    _observe(db, cost=1.0)
    _observe(db, cost=999.0, category="coding_general")
    assert _only(db)["cost"] == 1.0


def test_unswept_model_has_no_measurements(db):
    row = _only(db)
    assert row["cost"] is None
    assert row["eco"] is None
    assert row["samples"] == 0


# --- context estimation ---------------------------------------------------

def test_context_estimate_is_conservative_for_agent_traffic():
    # Real measurement: opencode sent a prompt that tokenized to 106,158 while
    # the chars/4 estimate put it under 94,196 — so it was routed to a model
    # that could not hold it within its effective window. Agent traffic is
    # code, JSON and tool schemas, which pack denser than prose.
    from dispatcher import estimate_prompt_tokens

    dense = "x" * 376_784  # the char count implied by that failure
    # At the old divisor of 4 this estimated 94,196 and passed the filter.
    assert estimate_prompt_tokens([{"role": "user", "content": dense}]) > 94_196


def test_context_estimate_counts_every_message():
    from dispatcher import estimate_prompt_tokens

    msgs = [{"role": "user", "content": "a" * 300}, {"role": "assistant", "content": "b" * 300}]
    assert estimate_prompt_tokens(msgs) == 200


def test_context_estimate_reads_multimodal_text_parts():
    from dispatcher import estimate_prompt_tokens

    msgs = [{"role": "user", "content": [
        {"type": "text", "text": "c" * 300},
        {"type": "image_url", "image_url": {"url": "data:..."}},
    ]}]
    assert estimate_prompt_tokens(msgs) == 100


def test_context_estimate_counts_tool_definitions():
    # upstream_body = {**body, "model": target} forwards `tools` to the
    # provider verbatim on every call, so it counts against the same context
    # window and has to count here too. Omitting it undercounts every
    # tool-carrying request -- nearly all agent traffic -- which is exactly
    # how qwen3.6-35b got a live 400 for "too long ... even after
    # compaction" while the router's own estimate said it fit.
    from dispatcher import estimate_prompt_tokens

    msgs = [{"role": "user", "content": "a" * 300}]
    tools = [{"type": "function", "function": {
        "name": "read_file",
        "parameters": {"type": "object", "properties": {"path": {"type": "string"}}},
    }}]
    without_tools = estimate_prompt_tokens(msgs)
    with_tools = estimate_prompt_tokens(msgs, tools=tools)
    assert with_tools > without_tools