testsAndMisc/python_pkg/token_audit/tests/test_model.py

from __future__ import annotations

import pytest

from python_pkg.token_audit.model import (
    COMPACTION_DROP_RATIO,
    WEIGHTS,
    Session,
    ToolCall,
    Turn,
    weighted,
)


def test_weights_price_output_highest() -> None:
    assert WEIGHTS["output_tokens"] > WEIGHTS["input_tokens"]
    assert WEIGHTS["cache_read_input_tokens"] < WEIGHTS["input_tokens"]
    assert 0 < COMPACTION_DROP_RATIO < 1


def test_weighted_applies_each_multiplier() -> None:
    usage = {
        "input_tokens": 10,
        "cache_creation_input_tokens": 100,
        "cache_read_input_tokens": 1000,
        "output_tokens": 1,
    }
    assert weighted(usage) == pytest.approx(10 + 125 + 100 + 5)


def test_weighted_treats_missing_kinds_as_zero() -> None:
    assert weighted({}) == 0.0


def test_turn_cost_delegates_to_weighted() -> None:
    turn = Turn(usage={"output_tokens": 2}, context=0, model="m", is_sidechain=False)
    assert turn.cost == pytest.approx(10.0)


@pytest.mark.parametrize("path", ["/a/b.png", "/a/b.PNG", "/a/b.jpeg"])
def test_tool_call_detects_images(path: str) -> None:
    assert ToolCall(name="Read", result_tokens=1, path=path).is_image


@pytest.mark.parametrize("path", ["/a/b.py", "/a/noextension", None])
def test_tool_call_rejects_non_images(path: str | None) -> None:
    assert not ToolCall(name="Read", result_tokens=1, path=path).is_image


@pytest.mark.parametrize(
    ("name", "expected"),
    [
        ("mcp__aseprite__draw_line", "aseprite"),
        ("mcp__i3wm__get_tree", "i3wm"),
        ("Read", None),
        ("mcp__broken", None),
    ],
)
def test_tool_call_mcp_server(name: str, expected: str | None) -> None:
    assert ToolCall(name=name, result_tokens=0).mcp_server == expected


def _turn(context: int, out: int = 0, read: int = 0, create: int = 0) -> Turn:
    return Turn(
        usage={
            "output_tokens": out,
            "cache_read_input_tokens": read,
            "cache_creation_input_tokens": create,
        },
        context=context,
        model="m",
        is_sidechain=False,
    )


def test_session_aggregates() -> None:
    session = Session(session_id="s", path="p")
    session.turns = [_turn(100, out=1), _turn(500, out=2)]
    assert session.turn_count == 2
    assert session.max_context == 500
    assert session.cost == pytest.approx(15.0)


def test_session_empty_has_zero_max_context() -> None:
    assert Session(session_id="s", path="p").max_context == 0


def test_cold_start_counts_only_turns_with_no_cache_read() -> None:
    session = Session(session_id="s", path="p")
    session.turns = [
        _turn(0, create=5000),  # cold: nothing to read yet
        _turn(9, read=9, create=100),  # warm: excluded
    ]
    assert session.cold_start_tokens == 5000