tests/unit/test_llm_manager.py

from __future__ import annotations

import asyncio
import pytest
import pytest_asyncio

from llm.manager import (
    AllProvidersFailedError,
    FallbackPolicy,
    LLMManager,
    ProviderAlreadyRegisteredError,
    ProviderNotFoundError,
)
from llm.providers.base_provider import (
    BaseProvider,
    LLMMessage,
    LLMProviderError,
    LLMRequest,
    LLMResponse,
)


# ---------------------------------------------------------------------------
# Stub providers for testing
# ---------------------------------------------------------------------------

class _SuccessProvider(BaseProvider):
    """A provider that always succeeds."""

    def __init__(self, name: str = "success") -> None:
        super().__init__(name=name)

    async def generate(self, request: LLMRequest) -> LLMResponse:
        return LLMResponse(
            content=f"response from {self.name}",
            model=request.model,
            provider=self.name,
        )


class _FailProvider(BaseProvider):
    """A provider that always fails with LLMProviderError."""

    def __init__(self, name: str = "fail") -> None:
        super().__init__(name=name)
        self.call_count: int = 0

    async def generate(self, request: LLMRequest) -> LLMResponse:
        self.call_count += 1
        raise LLMProviderError(
            f"provider {self.name} failed (attempt {self.call_count})",
            provider_name=self.name,
        )


class _FlakyProvider(BaseProvider):
    """A provider that fails N times then succeeds."""

    def __init__(self, fail_count: int = 1, name: str = "flaky") -> None:
        super().__init__(name=name)
        self._fail_count = fail_count
        self.call_count: int = 0

    async def generate(self, request: LLMRequest) -> LLMResponse:
        self.call_count += 1
        if self.call_count <= self._fail_count:
            raise LLMProviderError(
                f"transient failure (attempt {self.call_count})",
                provider_name=self.name,
            )
        return LLMResponse(
            content=f"success after {self.call_count} attempts",
            model=request.model,
            provider=self.name,
        )


# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------

def _make_request(model: str = "test-model") -> LLMRequest:
    return LLMRequest(
        messages=[LLMMessage(role="user", content="Hello")],
        model=model,
    )


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------

class TestLLMManagerRegistration:
    @pytest.mark.asyncio
    async def test_register_and_list(self) -> None:
        manager = LLMManager(default_provider="p1")
        await manager.register_provider("p1", _SuccessProvider("p1"))
        assert manager.list_providers() == ["p1"]

    @pytest.mark.asyncio
    async def test_duplicate_registration_raises(self) -> None:
        manager = LLMManager(default_provider="p1")
        await manager.register_provider("p1", _SuccessProvider("p1"))
        with pytest.raises(ProviderAlreadyRegisteredError):
            await manager.register_provider("p1", _SuccessProvider("p1"))

    @pytest.mark.asyncio
    async def test_duplicate_with_overwrite(self) -> None:
        manager = LLMManager(default_provider="p1")
        await manager.register_provider("p1", _SuccessProvider("p1"))
        await manager.register_provider("p1", _SuccessProvider("p1"), overwrite=True)
        assert manager.list_providers() == ["p1"]

    @pytest.mark.asyncio
    async def test_has_providers(self) -> None:
        manager = LLMManager()
        assert not manager.has_providers()
        await manager.register_provider("p1", _SuccessProvider("p1"))
        assert manager.has_providers()


class TestLLMManagerGetProvider:
    def test_get_default(self) -> None:
        manager = LLMManager(default_provider="p1")
        provider = _SuccessProvider("p1")
        manager._providers["p1"] = provider
        assert manager.get_provider() is provider

    def test_get_by_name(self) -> None:
        manager = LLMManager(default_provider="p1")
        provider = _SuccessProvider("p2")
        manager._providers["p2"] = provider
        assert manager.get_provider("p2") is provider

    def test_not_found_raises(self) -> None:
        manager = LLMManager(default_provider="missing")
        with pytest.raises(ProviderNotFoundError):
            manager.get_provider()


class TestLLMManagerGenerate:
    @pytest.mark.asyncio
    async def test_successful_generate(self) -> None:
        manager = LLMManager(default_provider="p1")
        await manager.register_provider("p1", _SuccessProvider("p1"))
        response = await manager.generate(_make_request())
        assert response.provider == "p1"
        assert "response" in response.content

    @pytest.mark.asyncio
    async def test_no_providers_raises(self) -> None:
        manager = LLMManager(default_provider="missing")
        with pytest.raises(ProviderNotFoundError):
            await manager.generate(_make_request())

    @pytest.mark.asyncio
    async def test_fallback_to_second_provider(self) -> None:
        policy = FallbackPolicy(
            providers=["fail", "success"],
            max_retries_per_provider=0,
        )
        manager = LLMManager(default_provider="fail", fallback_policy=policy)
        fail_provider = _FailProvider("fail")
        await manager.register_provider("fail", fail_provider)
        await manager.register_provider("success", _SuccessProvider("success"))

        response = await manager.generate(_make_request())
        assert response.provider == "success"

    @pytest.mark.asyncio
    async def test_all_providers_fail_raises(self) -> None:
        policy = FallbackPolicy(
            providers=["fail1", "fail2"],
            max_retries_per_provider=0,
        )
        manager = LLMManager(default_provider="fail1", fallback_policy=policy)
        await manager.register_provider("fail1", _FailProvider("fail1"))
        await manager.register_provider("fail2", _FailProvider("fail2"))

        with pytest.raises(AllProvidersFailedError) as exc_info:
            await manager.generate(_make_request())

        assert len(exc_info.value.attempts) == 2

    @pytest.mark.asyncio
    async def test_retry_then_success(self) -> None:
        policy = FallbackPolicy(
            providers=["flaky"],
            max_retries_per_provider=2,
            base_delay_seconds=0.01,
            max_delay_seconds=0.05,
        )
        manager = LLMManager(default_provider="flaky", fallback_policy=policy)
        flaky = _FlakyProvider(fail_count=1, name="flaky")
        await manager.register_provider("flaky", flaky)

        response = await manager.generate(_make_request())
        assert response.provider == "flaky"
        assert flaky.call_count == 2  # 1 fail + 1 success

    @pytest.mark.asyncio
    async def test_retry_exhausted_then_fallback(self) -> None:
        policy = FallbackPolicy(
            providers=["flaky", "success"],
            max_retries_per_provider=1,
            base_delay_seconds=0.01,
            max_delay_seconds=0.05,
        )
        manager = LLMManager(default_provider="flaky", fallback_policy=policy)
        flaky = _FlakyProvider(fail_count=3, name="flaky")  # Will never succeed within 1 retry
        await manager.register_provider("flaky", flaky)
        await manager.register_provider("success", _SuccessProvider("success"))

        response = await manager.generate(_make_request())
        assert response.provider == "success"

    @pytest.mark.asyncio
    async def test_explicit_provider_overrides_default(self) -> None:
        manager = LLMManager(default_provider="p1")
        await manager.register_provider("p1", _SuccessProvider("p1"))
        await manager.register_provider("p2", _SuccessProvider("p2"))

        response = await manager.generate(_make_request(), provider_name="p2")
        assert response.provider == "p2"


class TestFallbackPolicy:
    def test_default_values(self) -> None:
        policy = FallbackPolicy()
        assert policy.providers == []
        assert policy.max_retries_per_provider == 2
        assert policy.base_delay_seconds == 1.0
        assert policy.max_delay_seconds == 30.0

    def test_custom_values(self) -> None:
        policy = FallbackPolicy(
            providers=["openai", "claude"],
            max_retries_per_provider=3,
            base_delay_seconds=0.5,
            max_delay_seconds=10.0,
        )
        assert policy.providers == ["openai", "claude"]
        assert policy.max_retries_per_provider == 3