dev-puppeteer/my-deepagent/tests/unit/test_pricing.py

"""Unit tests for src/my_deepagent/monitoring/pricing.py."""

from __future__ import annotations

import httpx
import pytest
import respx

from my_deepagent.errors import MyDeepAgentError
from my_deepagent.monitoring.pricing import (
    ModelPrice,
    PricingCache,
    _parse_pricing_payload,
    fetch_openrouter_pricing,
)

# ---------------------------------------------------------------------------
# _parse_pricing_payload
# ---------------------------------------------------------------------------


def test_parse_valid_payload_returns_model_prices() -> None:
    data: dict[str, object] = {
        "data": [
            {
                "id": "deepseek/deepseek-chat",
                "pricing": {"prompt": "0.000001", "completion": "0.000002"},
                "context_length": 32768,
            },
            {
                "id": "anthropic/claude-sonnet",
                "pricing": {"prompt": "0.000003", "completion": "0.000015"},
                "context_length": 200000,
            },
        ]
    }
    result = _parse_pricing_payload(data)
    assert len(result) == 2
    assert result[0].model == "deepseek/deepseek-chat"
    assert result[0].input_per_1k_usd == pytest.approx(0.001)
    assert result[0].output_per_1k_usd == pytest.approx(0.002)
    assert result[0].context_length == 32768
    assert result[1].model == "anthropic/claude-sonnet"


def test_parse_empty_data_list_returns_empty() -> None:
    result = _parse_pricing_payload({"data": []})
    assert result == []


def test_parse_data_is_not_list_returns_empty() -> None:
    # data is a dict instead of list — malformed response
    result = _parse_pricing_payload({"data": {"id": "bad"}})
    assert result == []


def test_parse_missing_data_key_returns_empty() -> None:
    result = _parse_pricing_payload({})
    assert result == []


def test_parse_skips_entries_without_id() -> None:
    data: dict[str, object] = {
        "data": [
            {"pricing": {"prompt": "0.000001", "completion": "0.000002"}, "context_length": 1000},
        ]
    }
    result = _parse_pricing_payload(data)
    assert result == []


def test_parse_skips_entries_with_invalid_pricing_values() -> None:
    data: dict[str, object] = {
        "data": [
            {
                "id": "model/x",
                "pricing": {"prompt": "not-a-number", "completion": "also-bad"},
                "context_length": 1000,
            }
        ]
    }
    result = _parse_pricing_payload(data)
    assert result == []


def test_parse_handles_null_pricing_gracefully() -> None:
    data: dict[str, object] = {
        "data": [
            {"id": "model/y", "pricing": None, "context_length": 0},
        ]
    }
    result = _parse_pricing_payload(data)
    # pricing=None -> {} -> prompt/completion default to "0"
    assert len(result) == 1
    assert result[0].input_per_1k_usd == 0.0
    assert result[0].output_per_1k_usd == 0.0


def test_parse_handles_missing_context_length() -> None:
    data: dict[str, object] = {
        "data": [
            {"id": "model/z", "pricing": {"prompt": "0.000001", "completion": "0.000002"}},
        ]
    }
    result = _parse_pricing_payload(data)
    assert len(result) == 1
    assert result[0].context_length == 0


def test_parse_non_dict_entry_is_skipped() -> None:
    data: dict[str, object] = {"data": ["not-a-dict", None]}
    result = _parse_pricing_payload(data)
    assert result == []


# ---------------------------------------------------------------------------
# PricingCache.compute_cost
# ---------------------------------------------------------------------------


def test_compute_cost_known_model() -> None:
    cache = PricingCache()
    cache.set(
        [
            ModelPrice(
                model="deepseek/deepseek-chat",
                input_per_1k_usd=0.001,
                output_per_1k_usd=0.002,
                context_length=32768,
            )
        ]
    )
    cost = cache.compute_cost("deepseek/deepseek-chat", input_tokens=1000, output_tokens=500)
    assert cost == pytest.approx(0.001 * 1.0 + 0.002 * 0.5)


def test_compute_cost_openrouter_prefix_stripped() -> None:
    cache = PricingCache()
    cache.set(
        [
            ModelPrice(
                model="deepseek/deepseek-chat",
                input_per_1k_usd=0.001,
                output_per_1k_usd=0.002,
                context_length=32768,
            )
        ]
    )
    # Should strip "openrouter:" prefix when looking up
    cost = cache.compute_cost(
        "openrouter:deepseek/deepseek-chat", input_tokens=1000, output_tokens=0
    )
    assert cost == pytest.approx(0.001)


def test_compute_cost_unknown_model_returns_zero() -> None:
    cache = PricingCache()
    cost = cache.compute_cost("unknown/model", input_tokens=1000, output_tokens=1000)
    assert cost == 0.0


def test_compute_cost_zero_tokens_returns_zero() -> None:
    cache = PricingCache()
    cache.set(
        [ModelPrice(model="m/x", input_per_1k_usd=1.0, output_per_1k_usd=2.0, context_length=1000)]
    )
    assert cache.compute_cost("m/x", input_tokens=0, output_tokens=0) == 0.0


def test_pricing_cache_get_strips_openrouter_prefix() -> None:
    cache = PricingCache()
    cache.set(
        [ModelPrice(model="a/b", input_per_1k_usd=0.5, output_per_1k_usd=1.0, context_length=0)]
    )
    assert cache.get("openrouter:a/b") is not None
    assert cache.get("a/b") is not None


# ---------------------------------------------------------------------------
# fetch_openrouter_pricing (respx mock)
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_fetch_openrouter_pricing_success() -> None:
    payload = {
        "data": [
            {
                "id": "deepseek/deepseek-chat",
                "pricing": {"prompt": "0.000001", "completion": "0.000002"},
                "context_length": 64000,
            }
        ]
    }
    with respx.mock:
        respx.get("https://openrouter.ai/api/v1/models").mock(
            return_value=httpx.Response(200, json=payload)
        )
        result = await fetch_openrouter_pricing(
            api_key="sk-or-test", base_url="https://openrouter.ai/api/v1"
        )
    assert len(result) == 1
    assert result[0].model == "deepseek/deepseek-chat"


@pytest.mark.asyncio
async def test_fetch_openrouter_pricing_http_error_raises_recoverable() -> None:
    with respx.mock:
        respx.get("https://openrouter.ai/api/v1/models").mock(
            return_value=httpx.Response(401, json={"error": "unauthorized"})
        )
        with pytest.raises(MyDeepAgentError) as exc_info:
            await fetch_openrouter_pricing(
                api_key="bad-key", base_url="https://openrouter.ai/api/v1"
            )
    assert exc_info.value.code == "network_blip"


@pytest.mark.asyncio
async def test_fetch_openrouter_pricing_connect_error_raises_recoverable() -> None:
    with respx.mock:
        respx.get("https://openrouter.ai/api/v1/models").mock(
            side_effect=httpx.ConnectError("connection refused")
        )
        with pytest.raises(MyDeepAgentError) as exc_info:
            await fetch_openrouter_pricing(
                api_key="sk-or-test", base_url="https://openrouter.ai/api/v1"
            )
    assert exc_info.value.code == "network_blip"