Aether/tests/unit/test_provider_query_failover_endpoint.py

from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest
from pydantic import ValidationError

from src.api.admin import provider_query as provider_query_module
from src.api.admin.provider_query import (
    DEFAULT_MODEL_TEST_MESSAGE,
)
from src.api.admin.provider_query import TestModelFailoverRequest as FailoverRequestModel
from src.api.admin.provider_query import (
    _build_direct_test_candidates,
    _build_test_attempts_from_candidate_keys,
    _filter_test_candidates_by_endpoint,
    _flatten_test_candidates_for_concurrency,
    _maybe_mark_test_oauth_key_invalid,
    _require_test_endpoint_base_url,
    _resolve_test_effective_model,
    _resolve_test_message,
)
from src.services.scheduling.schemas import PoolCandidate


def _build_provider() -> tuple[SimpleNamespace, SimpleNamespace, SimpleNamespace]:
    endpoint_a = SimpleNamespace(id="ep-a", api_format="openai:chat", is_active=True)
    endpoint_b = SimpleNamespace(id="ep-b", api_format="claude:cli", is_active=True)
    key_all = SimpleNamespace(
        id="key-all", is_active=True, api_formats=["openai:chat", "claude:cli"]
    )
    key_b = SimpleNamespace(id="key-b", is_active=True, api_formats=["claude:cli"])
    provider = SimpleNamespace(
        id="provider-1", endpoints=[endpoint_a, endpoint_b], api_keys=[key_all, key_b]
    )
    return provider, endpoint_a, endpoint_b


def test_build_direct_test_candidates_respects_endpoint_id() -> None:
    provider, _endpoint_a, endpoint_b = _build_provider()

    candidates = _build_direct_test_candidates(provider, endpoint_id=endpoint_b.id)  # type: ignore[arg-type]

    assert {candidate.endpoint.id for candidate in candidates} == {endpoint_b.id}
    assert {candidate.key.id for candidate in candidates} == {"key-all", "key-b"}


def test_filter_test_candidates_by_endpoint_keeps_matching_candidates() -> None:
    provider, endpoint_a, endpoint_b = _build_provider()
    candidates = _build_direct_test_candidates(provider)  # type: ignore[arg-type]

    filtered = _filter_test_candidates_by_endpoint(candidates, endpoint_a.id)

    assert {candidate.endpoint.id for candidate in filtered} == {endpoint_a.id}
    assert all(candidate.endpoint.id != endpoint_b.id for candidate in filtered)


def test_resolve_test_effective_model_prefers_pool_key_mapping() -> None:
    provider, endpoint_a, _endpoint_b = _build_provider()
    candidate = _build_direct_test_candidates(provider, endpoint_id=endpoint_a.id)[0]  # type: ignore[arg-type]
    pool_key = SimpleNamespace(id="pool-key", _pool_mapping_matched_model="mapped-model")
    request = SimpleNamespace(mode="global", model_name="gpt-4")

    effective = _resolve_test_effective_model(
        provider=provider,  # type: ignore[arg-type]
        candidate=candidate,
        request=request,  # type: ignore[arg-type]
        gm_obj=None,
        key=pool_key,
    )

    assert effective == "mapped-model"


def test_build_test_attempts_from_candidate_keys_includes_retry_index() -> None:
    candidate_keys = [
        SimpleNamespace(
            candidate_index=2,
            retry_index=1,
            key_id="key-b",
            key_name="Key B",
            auth_type="api_key",
            status="failed",
            skip_reason=None,
            error_message="timeout",
            status_code=504,
            latency_ms=1200,
        )
    ]

    attempts = _build_test_attempts_from_candidate_keys(
        candidate_keys=candidate_keys,
        candidate_meta_by_pair={
            (2, "key-b"): {
                "endpoint_api_format": "openai:chat",
                "endpoint_base_url": "https://example.test/v1",
                "effective_model": "mapped-model",
            }
        },
        candidate_meta_by_index={},
    )

    assert len(attempts) == 1
    assert attempts[0].retry_index == 1
    assert attempts[0].effective_model == "mapped-model"
    assert attempts[0].endpoint_api_format == "openai:chat"


def test_flatten_test_candidates_for_concurrency_expands_pool_keys() -> None:
    provider, endpoint_a, _endpoint_b = _build_provider()
    pool_key_a = SimpleNamespace(
        id="pool-a",
        name="Pool A",
        auth_type="oauth",
        _pool_mapping_matched_model="mapped-a",
    )
    pool_key_b = SimpleNamespace(
        id="pool-b",
        name="Pool B",
        auth_type="oauth",
        _pool_skipped=True,
        _pool_skip_reason="quota_exhausted",
    )
    candidate = PoolCandidate(
        provider=provider,  # type: ignore[arg-type]
        endpoint=endpoint_a,  # type: ignore[arg-type]
        key=pool_key_a,  # type: ignore[arg-type]
        pool_keys=[pool_key_a, pool_key_b],  # type: ignore[list-item]
        is_cached=True,
        provider_api_format="openai:chat",
    )

    flattened = _flatten_test_candidates_for_concurrency([candidate])

    assert len(flattened) == 2
    assert flattened[0].key.id == "pool-a"
    assert flattened[0].mapping_matched_model == "mapped-a"
    assert flattened[0].is_skipped is False
    assert flattened[1].key.id == "pool-b"
    assert flattened[1].is_skipped is True
    assert flattened[1].skip_reason == "quota_exhausted"


def test_test_model_failover_request_validates_concurrency_range() -> None:
    ok = FailoverRequestModel(
        provider_id="p1",
        mode="global",
        model_name="gpt-4o-mini",
        concurrency=5,
    )
    assert ok.concurrency == 5

    with pytest.raises(ValidationError):
        FailoverRequestModel(
            provider_id="p1",
            mode="global",
            model_name="gpt-4o-mini",
            concurrency=0,
        )


def test_resolve_test_message_uses_default_for_blank_input() -> None:
    assert _resolve_test_message(None) == DEFAULT_MODEL_TEST_MESSAGE
    assert _resolve_test_message("") == DEFAULT_MODEL_TEST_MESSAGE
    assert _resolve_test_message("   ") == DEFAULT_MODEL_TEST_MESSAGE


def test_resolve_test_message_preserves_custom_input() -> None:
    assert _resolve_test_message("  custom prompt  ") == "custom prompt"


def test_require_test_endpoint_base_url_rejects_non_string() -> None:
    endpoint = SimpleNamespace(id="ep-bad", api_format="claude:chat", base_url={"url": "https://x"})

    with pytest.raises(ValueError, match="invalid base_url type"):
        _require_test_endpoint_base_url(endpoint)


def test_require_test_endpoint_base_url_trims_whitespace() -> None:
    endpoint = SimpleNamespace(
        id="ep-ok", api_format="claude:chat", base_url="  https://api.anthropic.com  "
    )

    assert _require_test_endpoint_base_url(endpoint) == "https://api.anthropic.com"


@pytest.mark.asyncio
async def test_maybe_mark_test_oauth_key_invalid_skips_account_block_when_oauth_check_fails(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    key = SimpleNamespace(id="key-1", oauth_invalid_at=None, oauth_invalid_reason=None)
    endpoint = SimpleNamespace(api_format="openai:chat")
    db = MagicMock()

    async def _fake_verify(**_: object) -> bool:
        key.oauth_invalid_reason = "[OAUTH_EXPIRED] refresh token expired"
        return False

    monkeypatch.setattr(provider_query_module, "verify_oauth_before_account_block", _fake_verify)

    await _maybe_mark_test_oauth_key_invalid(
        db=db,
        endpoint=endpoint,
        key=key,
        auth_type="oauth",
        error_payload={
            "error": {
                "code": 403,
                "message": "Please verify your account",
                "status": "PERMISSION_DENIED",
            }
        },
    )

    assert key.oauth_invalid_at is None
    assert key.oauth_invalid_reason == "[OAUTH_EXPIRED] refresh token expired"
    db.commit.assert_not_called()


@pytest.mark.asyncio
async def test_maybe_mark_test_oauth_key_invalid_marks_account_block_after_oauth_check(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    key = SimpleNamespace(id="key-2", oauth_invalid_at=None, oauth_invalid_reason=None)
    endpoint = SimpleNamespace(api_format="openai:chat")
    db = MagicMock()

    async def _fake_verify(**_: object) -> bool:
        return True

    monkeypatch.setattr(provider_query_module, "verify_oauth_before_account_block", _fake_verify)

    await _maybe_mark_test_oauth_key_invalid(
        db=db,
        endpoint=endpoint,
        key=key,
        auth_type="oauth",
        error_payload={
            "error": {
                "code": 403,
                "message": "verify your account",
                "status": "PERMISSION_DENIED",
            }
        },
    )

    assert key.oauth_invalid_at is not None
    assert str(key.oauth_invalid_reason).startswith("[ACCOUNT_BLOCK] ")
    db.commit.assert_called_once()