Create tests/unit/test_llm_manager.py

tests/unit/test_llm_manager.py

@@ -0,0 +1,244 @@
+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
+