6krrt/tests/test_openrouter_poller.py

"""Offline tests for the OpenRouter catalog fetcher and parser in poller.py."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest
import requests

from poller import (
    OPENROUTER_VIRTUAL_ROUTERS,
    ModelRow,
    fetch_openrouter,
    parse_openrouter_model,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _fake_response(json_data: dict):
    fake = MagicMock()
    fake.json.return_value = json_data
    fake.raise_for_status.return_value = None
    fake.status_code = 200
    return fake


def _openrouter_model(
    model_id: str,
    *,
    canonical_slug: str | None = None,
    prompt: str = "0.000001",
    completion: str = "0.000002",
    context_length: int = 131072,
    max_completion_tokens: int = 8192,
    modalities: list[str] | None = None,
    supported_parameters: list[str] | None = None,
    reasoning: dict | None = None,
):
    """Build a realistic OpenRouter /v1/models entry."""
    return {
        "id": model_id,
        "name": model_id.replace("/", " ").title(),
        "canonical_slug": canonical_slug or model_id,
        "pricing": {
            "prompt": prompt,
            "completion": completion,
            "web_search": "0.000003",
            "cache_read": "0.0000005",
        },
        "top_provider": {
            "context_length": context_length,
            "max_completion_tokens": max_completion_tokens,
        },
        "architecture": {
            "input_modalities": modalities or ["text"],
        },
        "supported_parameters": supported_parameters or [],
        "reasoning": reasoning,
    }


# ---------------------------------------------------------------------------
# parse_openrouter_model
# ---------------------------------------------------------------------------

def test_parse_openrouter_model_populates_required_fields():
    raw = _openrouter_model(
        "openai/gpt-6-astra",
        canonical_slug="gpt-6-astra",
        prompt="0.0000015",
        completion="0.0000060",
        context_length=200000,
        max_completion_tokens=32768,
        modalities=["text", "image"],
        supported_parameters=["tools", "response_format", "reasoning"],
        reasoning={"mode": "medium", "default_enabled": True},
    )

    row = parse_openrouter_model(raw)

    assert row == ModelRow(
        model_id="openai/gpt-6-astra",
        provider="openrouter",
        base_model_id="gpt-6-astra",
        display_name="Openai Gpt-6-Astra",
        cost_per_1m_prompt=1.5,
        cost_per_1m_completion=6.0,
        cost_per_1m_prompt_cached=None,
        context_window=200000,
        max_output_tokens=32768,
        supports_tools=True,
        supports_json_mode=True,
        supports_vision=True,
        supports_reasoning=True,
        reasoning_default_enabled=True,
        latency_class="standard",
        reasoning_mode="medium",
        context_variant="standard",
        access_level="public",
        pricing_tbd=False,
        deprecated=False,
    )


def test_parse_openrouter_model_batch_variant_maps_to_flex():
    raw = _openrouter_model(
        "anthropic/claude-4-opus:batch",
        canonical_slug="claude-4-opus",
    )

    row = parse_openrouter_model(raw)

    assert row.latency_class == "flex"
    assert row.context_variant == "batch"
    assert row.model_id == "anthropic/claude-4-opus:batch"


def test_parse_openrouter_model_free_variant_maps_to_standard():
    raw = _openrouter_model(
        "meta-llama/llama-4-maverick:free",
        canonical_slug="llama-4-maverick",
    )

    row = parse_openrouter_model(raw)

    assert row.latency_class == "standard"
    assert row.context_variant == "free"


def test_parse_openrouter_model_reasoning_defaults():
    raw = _openrouter_model("google/gemini-2.5-flash")

    row = parse_openrouter_model(raw)

    assert row.reasoning_mode == "none"
    assert row.reasoning_default_enabled is False
    assert row.supports_reasoning is False


@pytest.mark.parametrize("router_id", sorted(OPENROUTER_VIRTUAL_ROUTERS))
def test_parse_openrouter_model_excludes_virtual_routers(router_id):
    raw = _openrouter_model(router_id)

    assert parse_openrouter_model(raw) is None


def test_parse_openrouter_model_marks_missing_prompt_price_as_tbd():
    raw = _openrouter_model("x-ai/grok-3", prompt="")

    row = parse_openrouter_model(raw)

    assert row.cost_per_1m_prompt is None
    assert row.pricing_tbd is True


def test_parse_openrouter_model_gracefully_handles_malformed_prices():
    raw = _openrouter_model("x-ai/grok-3-mini", prompt="free", completion="N/A")

    row = parse_openrouter_model(raw)

    assert row.cost_per_1m_prompt is None
    assert row.cost_per_1m_completion is None
    assert row.pricing_tbd is True


# ---------------------------------------------------------------------------
# fetch_openrouter
# ---------------------------------------------------------------------------

@patch.object(requests, "get")
def test_fetch_openrouter_returns_model_rows(mock_get):
    mock_get.return_value = _fake_response(
        {
            "data": [
                _openrouter_model(
                    "openai/gpt-6-astra",
                    canonical_slug="gpt-6-astra",
                ),
                _openrouter_model(
                    "anthropic/claude-4-opus:batch",
                    canonical_slug="claude-4-opus",
                ),
            ]
        }
    )

    rows = fetch_openrouter("openrouter")

    assert len(rows) == 2
    assert rows[0].model_id == "openai/gpt-6-astra"
    assert rows[0].base_model_id == "gpt-6-astra"
    assert rows[1].model_id == "anthropic/claude-4-opus:batch"
    assert rows[1].latency_class == "flex"
    mock_get.assert_called_once_with(
        "https://openrouter.ai/api/v1/models", timeout=20
    )


@patch.object(requests, "get")
def test_fetch_openrouter_excludes_virtual_routers(mock_get):
    mock_get.return_value = _fake_response(
        {
            "data": [
                _openrouter_model("openrouter/auto"),
                _openrouter_model("openrouter/free"),
                _openrouter_model("openrouter/fusion"),
                _openrouter_model("openai/gpt-6-astra", canonical_slug="gpt-6-astra"),
            ]
        }
    )

    rows = fetch_openrouter("openrouter")

    assert [r.model_id for r in rows] == ["openai/gpt-6-astra"]


@patch.object(requests, "get")
def test_fetch_openrouter_propagates_request_exception(mock_get):
    mock_get.side_effect = requests.RequestException("offline")

    with pytest.raises(requests.RequestException):
        fetch_openrouter("openrouter")