llm-connect/tests/test_cli.py

import json
from datetime import datetime, timezone
from zoneinfo import ZoneInfo

from llm_connect.cli import main
from llm_connect.quality import QualityLedger, QualityObservation
from llm_connect.usage import UsageEvent, UsageLedger


def test_rates_show_json_outputs_default_registry(capsys):
    assert main(["rates", "show", "--json"]) == 0

    payload = json.loads(capsys.readouterr().out)

    assert payload["openai/gpt-4o-mini"]["prompt_per_1k"] == 0.00015
    assert payload["moonshotai/kimi-k3"]["prompt_per_1k"] == 0.003


def test_classes_show_lists_builtins(capsys):
    assert main(["classes", "show"]) == 0

    output = capsys.readouterr().out

    assert "chunk-summarization" in output
    assert "entity-extraction" in output


def test_classes_fit_reads_quality_ledger(tmp_path, capsys):
    ledger = QualityLedger(tmp_path / "quality.jsonl")
    for _ in range(3):
        ledger.append(
            QualityObservation(
                task_type="extract",
                adapter_id="openrouter",
                model_id="openai/gpt-4o-mini",
                cost_usd=0.001,
                quality_score=0.9,
                latency_ms=100,
                tokens_in=500,
                tokens_out=350,
                recorded_at=datetime(2026, 5, 19, tzinfo=timezone.utc),
                tags={
                    "problem_class": "entity-extraction",
                    "dimensions": {
                        "chunk_words": 300,
                        "template_words": 100,
                        "expected_entities": 5,
                    },
                },
            )
        )

    assert main(["classes", "fit", str(ledger.path), "--class", "entity-extraction", "--json"]) == 0

    payload = json.loads(capsys.readouterr().out)

    assert payload["entity-extraction"]["params"]["tokens_per_entity"] == 70


def test_run_with_mock_reports_cost_and_writes_ledger(tmp_path, capsys):
    ledger = tmp_path / "usage.jsonl"
    assert (
        main(
            [
                "run",
                "hello world",
                "--provider",
                "mock",
                "--model",
                "moonshotai/kimi-k3",
                "--ledger",
                str(ledger),
                "--eur-per-usd",
                "1.0",
                "--json",
            ]
        )
        == 0
    )

    payload = json.loads(capsys.readouterr().out)
    assert payload["content"]
    assert payload["usage"]["total_tokens"] > 0
    assert payload["cost_usd"] is not None
    assert payload["cost_eur"] is not None
    events = UsageLedger(ledger).read_all()
    assert len(events) == 1
    assert events[0].source == "cli"
    assert events[0].model_id == "moonshotai/kimi-k3"


def test_cost_estimate_cli(capsys):
    assert (
        main(
            [
                "cost",
                "estimate",
                "--model",
                "moonshotai/kimi-k3",
                "--prompt-tokens",
                "1000",
                "--completion-tokens",
                "1000",
                "--eur-per-usd",
                "1",
                "--json",
            ]
        )
        == 0
    )
    payload = json.loads(capsys.readouterr().out)
    assert payload["cost_usd"] == 0.018
    assert payload["cost_eur"] == 0.018


def test_spend_week_current_and_last(tmp_path, capsys, monkeypatch):
    ledger_path = tmp_path / "usage.jsonl"
    ledger = UsageLedger(ledger_path)
    berlin = ZoneInfo("Europe/Berlin")
    current_start = datetime(2026, 8, 3, 0, 0, tzinfo=berlin)
    current_end = datetime(2026, 8, 5, 12, 0, tzinfo=berlin)
    last_start = datetime(2026, 7, 27, 0, 0, tzinfo=berlin)
    last_end = current_start

    def _fake_week_window(which="current", *, now=None, tz=None):
        if which == "last":
            return last_start, last_end
        return current_start, current_end

    monkeypatch.setattr("llm_connect.cli.week_window", _fake_week_window)

    ledger.append(
        UsageEvent(
            provider="mock",
            model_id="moonshotai/kimi-k3",
            prompt_tokens=100,
            completion_tokens=20,
            total_tokens=120,
            cost_usd=0.01,
            cost_eur=0.009,
            cost_source="rate_table:moonshotai/kimi-k3",
            source="cli",
            recorded_at=datetime(2026, 8, 4, 9, 0, tzinfo=berlin),
        )
    )
    ledger.append(
        UsageEvent(
            provider="mock",
            model_id="moonshotai/kimi-k3",
            prompt_tokens=50,
            completion_tokens=10,
            total_tokens=60,
            cost_usd=0.005,
            cost_eur=0.0045,
            cost_source="rate_table:moonshotai/kimi-k3",
            source="cli",
            recorded_at=datetime(2026, 7, 29, 9, 0, tzinfo=berlin),  # previous week
        )
    )

    assert main(["spend", "week", "--ledger", str(ledger_path), "--json"]) == 0
    current = json.loads(capsys.readouterr().out)
    assert current["event_count"] == 1
    assert current["total_tokens"] == 120
    assert current["cost_eur"] == 0.009

    assert main(["spend", "week", "--last", "--ledger", str(ledger_path), "--json"]) == 0
    last = json.loads(capsys.readouterr().out)
    assert last["event_count"] == 1
    assert last["total_tokens"] == 60