pragent/tests/pilot/test_cost_model.py

"""Unit tests for the per-review cost model. No network."""
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(os.path.join(HERE, "..", ".."))
sys.path.insert(0, os.path.join(ROOT, "pilot"))

import cost_model as cm  # noqa: E402

FACTORY = cm.measure_factory(ROOT)


# ---------------------------------------------------------------------------
# measure_factory — reads the real files
# ---------------------------------------------------------------------------


def test_measure_factory_finds_agent_and_skills():
    assert FACTORY["agent"] > 500
    for skill in cm.ALWAYS_SKILLS:
        assert FACTORY[f"skill:{skill}"] > 100, skill
    for lens in ("security", "tests", "perf"):
        assert FACTORY[f"subagent:{lens}"] > 100, lens


def test_measure_factory_missing_root_is_empty_not_an_error():
    f = cm.measure_factory("/nonexistent-path-for-test")
    assert f == {"agent": 0, "subagent:security": 0, "subagent:tests": 0, "subagent:perf": 0}


# ---------------------------------------------------------------------------
# prefix_tokens — tiers load different skill sets
# ---------------------------------------------------------------------------


def test_prefix_grows_with_tier():
    sizes = [cm.prefix_tokens(t, FACTORY) for t in ("trivial", "lite", "full", "oversized")]
    assert sizes == sorted(sizes)
    assert sizes[0] < sizes[-1]


def test_prefix_includes_harness_and_agent():
    assert cm.prefix_tokens("trivial", FACTORY) > cm.HARNESS_TOKENS + FACTORY["agent"]


def test_unknown_tier_still_returns_the_always_skills():
    assert cm.prefix_tokens("nope", FACTORY) == cm.prefix_tokens("trivial", FACTORY)


# ---------------------------------------------------------------------------
# tier_usage — the loop's resend behaviour is what costs money
# ---------------------------------------------------------------------------


def _tier(name):
    return next(t for t in cm.DEFAULT_TIERS if t.name == name)


def test_caching_moves_the_stable_prefix_out_of_uncached_input():
    t = _tier("full")
    cached = cm.tier_usage(t, FACTORY, caching=True)
    uncached = cm.tier_usage(t, FACTORY, caching=False)
    assert cached.uncached_input < uncached.uncached_input
    assert cached.cached_input > 0
    assert uncached.cached_input == 0
    assert uncached.cache_writes == 0


def test_total_input_is_the_same_work_either_way():
    # Caching changes the *price* of the tokens, not how many are sent.
    t = _tier("full")
    a = cm.tier_usage(t, FACTORY, caching=True)
    b = cm.tier_usage(t, FACTORY, caching=False)
    # With caching the step-1 stable block is billed as a cache write rather
    # than as input, so it moves columns — the grand total of tokens sent is
    # identical.
    assert a.total_input + a.cache_writes == b.total_input


def test_more_steps_cost_more_input():
    base = _tier("lite")
    more = cm.Tier(
        "lite-long", base.diff_tokens, base.steps * 2, base.file_reads,
        base.tokens_per_read, base.output_tokens,
    )
    assert cm.tier_usage(more, FACTORY).total_input > cm.tier_usage(base, FACTORY).total_input


def test_trivial_tier_does_no_tool_work():
    u = cm.tier_usage(_tier("trivial"), FACTORY)
    assert u.uncached_input == 0  # no tool-result tail at all
    assert u.output > 0


def test_subagents_add_input_and_output():
    t = _tier("full")
    with_subs = cm.Tier(
        t.name, t.diff_tokens, t.steps, t.file_reads, t.tokens_per_read,
        t.output_tokens, subagents=2,
    )
    a, b = cm.tier_usage(t, FACTORY), cm.tier_usage(with_subs, FACTORY)
    assert b.total_input > a.total_input
    assert b.output > a.output


# ---------------------------------------------------------------------------
# cost — prices and discounts
# ---------------------------------------------------------------------------


def test_cost_is_ordered_by_model_price():
    u = cm.tier_usage(_tier("full"), FACTORY)
    opus = cm.cost(u, cm.PRICES["claude-opus-5"])
    sonnet = cm.cost(u, cm.PRICES["claude-sonnet-5"])
    haiku = cm.cost(u, cm.PRICES["claude-haiku-4-5"])
    assert opus > sonnet > haiku > 0


def test_batch_is_exactly_half():
    u = cm.tier_usage(_tier("full"), FACTORY)
    p = cm.PRICES["claude-opus-5"]
    assert abs(cm.cost(u, p, batch=True) * 2 - cm.cost(u, p)) < 1e-9


def test_cost_matches_a_hand_calculation():
    u = cm.Usage(uncached_input=1_000_000, cached_input=1_000_000,
                 cache_writes=1_000_000, output=1_000_000)
    p = cm.PRICES["claude-opus-5"]  # 5 / 25 / 6.25 / 0.50
    assert abs(cm.cost(u, p) - (5.00 + 0.50 + 6.25 + 25.00)) < 1e-9


def test_caching_is_cheaper_than_not_caching():
    for name in ("lite", "full", "oversized"):
        t = _tier(name)
        p = cm.PRICES["claude-opus-5"]
        assert cm.cost(cm.tier_usage(t, FACTORY, True), p) < \
               cm.cost(cm.tier_usage(t, FACTORY, False), p), name


def test_cost_rises_monotonically_with_tier():
    p = cm.PRICES["claude-sonnet-5"]
    costs = [cm.cost(cm.tier_usage(_tier(n), FACTORY), p)
             for n in ("trivial", "lite", "full", "oversized")]
    assert costs == sorted(costs)


# ---------------------------------------------------------------------------
# blended + CLI
# ---------------------------------------------------------------------------


def test_blended_sits_between_the_cheapest_and_priciest_tier():
    p = cm.PRICES["claude-opus-5"]
    blended = cm.blended_cost(cm.DEFAULT_TIERS, FACTORY, p, True)
    per_tier = [cm.cost(cm.tier_usage(t, FACTORY), p) for t in cm.DEFAULT_TIERS]
    assert min(per_tier) < blended < max(per_tier)


def test_shares_that_do_not_sum_to_one_are_normalised():
    p = cm.PRICES["claude-opus-5"]
    tiers = [cm.Tier(t.name, t.diff_tokens, t.steps, t.file_reads,
                     t.tokens_per_read, t.output_tokens, t.subagents, share=t.share * 2)
             for t in cm.DEFAULT_TIERS]
    doubled = cm.blended_cost(tiers, FACTORY, p, True)
    normal = cm.blended_cost(cm.DEFAULT_TIERS, FACTORY, p, True)
    assert abs(doubled - normal) < 1e-9


def test_report_renders_every_requested_model():
    text = cm.report(cm.DEFAULT_TIERS, 350, True, ["claude-opus-5", "gpt-5.6-luna"])
    assert "Claude Opus 5" in text
    assert "GPT-5.6 Luna" in text
    assert "Claude Sonnet 5" not in text
    assert "350 PRs/month" in text


def test_main_rejects_unknown_model(capsys):
    try:
        cm.main(["--models", "gpt-9"])
    except SystemExit as e:
        assert e.code != 0
    else:
        raise AssertionError("expected SystemExit")


def test_main_rejects_bad_mix():
    try:
        cm.main(["--mix", "50,50"])
    except SystemExit as e:
        assert e.code != 0
    else:
        raise AssertionError("expected SystemExit")


def test_main_runs(capsys):
    assert cm.main(["--models", "claude-sonnet-5", "--prs-per-month", "10"]) == 0
    assert "per month" in capsys.readouterr().out