Aether/tests/services/test_format_conversion_health_buckets.py

from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from src.services.orchestration.error_classifier import ErrorClassifier
from src.services.request.executor import RequestExecutor


@asynccontextmanager
async def _noop_async_cm() -> AsyncGenerator[None]:
    yield


@pytest.mark.asyncio
async def test_executor_records_health_by_provider_format() -> None:
    db = MagicMock()

    concurrency_manager = MagicMock()
    concurrency_manager.get_current_concurrency = AsyncMock(return_value=(0, 0))
    concurrency_manager.get_key_rpm_count = AsyncMock(return_value=1)
    concurrency_manager.rpm_guard = MagicMock(return_value=_noop_async_cm())

    adaptive_manager = MagicMock()

    provider = MagicMock()
    provider.id = "p1"
    provider.name = "p1"

    endpoint = MagicMock()
    endpoint.id = "e1"
    endpoint.api_format = "openai:chat"
    endpoint.api_family = "openai"
    endpoint.endpoint_kind = "chat"

    key = MagicMock()
    key.id = "k1"
    key.api_key = "encrypted"
    key.rpm_limit = 10
    key.learned_rpm_limit = None
    key.cache_ttl_minutes = 0

    candidate = MagicMock()
    candidate.provider = provider
    candidate.endpoint = endpoint
    candidate.key = key
    candidate.is_cached = False

    async def request_func(
        _provider: Any, _endpoint: Any, _key: Any, _candidate: Any
    ) -> dict[str, bool]:
        return {"ok": True}

    with (
        patch("src.services.request.executor.RequestCandidateService.mark_candidate_started"),
        patch("src.services.request.executor.RequestCandidateService.mark_candidate_success"),
        patch("src.services.request.executor.get_adaptive_reservation_manager") as mock_res_mgr,
        patch("src.services.request.executor.get_health_monitor") as mock_get_health_monitor,
    ):
        mock_res_mgr.return_value.calculate_reservation.return_value = MagicMock(
            ratio=0.0, phase="stable", confidence=1.0
        )
        record_success = mock_get_health_monitor.return_value.record_success

        executor = RequestExecutor(
            db=db, concurrency_manager=concurrency_manager, adaptive_manager=adaptive_manager
        )
        await executor.execute(
            candidate=candidate,
            candidate_id="c1",
            candidate_index=0,
            user_api_key=MagicMock(user_id="u1", id="ak1"),
            request_func=request_func,
            request_id="r1",
            api_format="claude:chat",  # client_format
            model_name="m",
            is_stream=False,
        )

        record_success.assert_called()
        assert record_success.call_args.kwargs["api_format"] == "openai:chat"


@pytest.mark.asyncio
async def test_error_classifier_records_failure_by_provider_format() -> None:
    db = MagicMock()
    classifier = ErrorClassifier(db=db, cache_scheduler=None, adaptive_manager=MagicMock())

    provider = MagicMock()
    provider.name = "p1"

    endpoint = MagicMock()
    endpoint.id = "e1"
    endpoint.api_format = "openai:chat"
    endpoint.api_family = "openai"
    endpoint.endpoint_kind = "chat"

    key = MagicMock()
    key.id = "k1"

    with patch(
        "src.services.orchestration.error_handler.get_health_monitor"
    ) as mock_get_health_monitor:
        record_failure = mock_get_health_monitor.return_value.record_failure
        await classifier.handle_retriable_error(
            error=RuntimeError("boom"),
            provider=provider,
            endpoint=endpoint,
            key=key,
            affinity_key="aff",
            api_format="claude:chat",  # client_format
            global_model_id="gm1",
            captured_key_concurrent=None,
            elapsed_ms=None,
            request_id="r1",
            attempt=1,
            max_attempts=2,
        )

        record_failure.assert_called()
        assert record_failure.call_args.kwargs["api_format"] == "openai:chat"