tests/test_cost_log.py

"""tests/test_cost_log.py — pricing table and usage-shape parsing (agent/cost_log.py)."""
from __future__ import annotations

from agent.cost_log import _PRICING, compute_run_cost


class _FakeMessage:
    def __init__(self, usage_metadata=None, response_metadata=None):
        self.usage_metadata = usage_metadata
        self.response_metadata = response_metadata or {}


def test_compute_run_cost_reads_usage_metadata_shape():
    msg = _FakeMessage(usage_metadata={"input_tokens": 100, "output_tokens": 40})
    result = compute_run_cost([msg], "claude-haiku-4-5-20251001")
    assert result["input_tokens"] == 100
    assert result["output_tokens"] == 40


def test_compute_run_cost_reads_anthropic_legacy_shape():
    msg = _FakeMessage(response_metadata={"usage": {"input_tokens": 200, "output_tokens": 80}})
    result = compute_run_cost([msg], "claude-haiku-4-5-20251001")
    assert result["input_tokens"] == 200
    assert result["output_tokens"] == 80


def test_compute_run_cost_reads_openai_legacy_shape():
    msg = _FakeMessage(
        response_metadata={"token_usage": {"prompt_tokens": 300, "completion_tokens": 120}}
    )
    result = compute_run_cost([msg], "gpt-5-mini")
    assert result["input_tokens"] == 300
    assert result["output_tokens"] == 120


def test_compute_run_cost_does_not_double_count_when_multiple_shapes_present():
    # usage_metadata takes priority; legacy shapes on the same message must be ignored.
    msg = _FakeMessage(
        usage_metadata={"input_tokens": 100, "output_tokens": 40},
        response_metadata={
            "usage": {"input_tokens": 999, "output_tokens": 999},
            "token_usage": {"prompt_tokens": 999, "completion_tokens": 999},
        },
    )
    result = compute_run_cost([msg], "claude-haiku-4-5-20251001")
    assert result["input_tokens"] == 100
    assert result["output_tokens"] == 40


def test_compute_run_cost_sums_across_messages():
    msgs = [
        _FakeMessage(usage_metadata={"input_tokens": 10, "output_tokens": 5}),
        _FakeMessage(response_metadata={"usage": {"input_tokens": 20, "output_tokens": 8}}),
    ]
    result = compute_run_cost(msgs, "claude-haiku-4-5-20251001")
    assert result["input_tokens"] == 30
    assert result["output_tokens"] == 13


def test_pricing_table_has_a_row_for_every_catalog_model():
    from agent.llm import MODEL_CATALOG
    for models in MODEL_CATALOG.values():
        for model_id, _label in models:
            assert model_id in _PRICING, f"missing pricing row for {model_id}"


def test_compute_run_cost_unknown_model_falls_back_to_default_pricing():
    msg = _FakeMessage(usage_metadata={"input_tokens": 1_000_000, "output_tokens": 1_000_000})
    result = compute_run_cost([msg], "some-unlisted-model")
    assert result["cost_usd"] > 0