Spaces:
Running
Running
File size: 3,713 Bytes
c8365f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | import asyncio
import unittest
from unittest.mock import patch
from models.ai_client import AIClient, ProviderConfig, ProviderUnavailableError
class _FailingCompletions:
def create(self, **_kwargs):
raise RuntimeError("quota exhausted")
class _FailingChat:
completions = _FailingCompletions()
class _FailingClient:
chat = _FailingChat()
class _ClientWithFailingProviders(AIClient):
def __init__(self):
self.providers = [
ProviderConfig(name="primary", api_key="x", base_url="https://example.invalid", default_model="model-a"),
ProviderConfig(name="fallback", api_key="y", base_url="https://example.invalid", default_model="model-b"),
]
self._client_cache = {}
self._rr_indices = {}
# Stato minimo richiesto dai percorsi chat/stream dopo l’introduzione
# del circuit breaker per profilo. Non chiama AIClient.__init__ e non
# carica provider o segreti dall’ambiente.
self._breaker = {}
self._breaker_threshold = 2
self._breaker_cooldown_s = 60.0
def _client_for(self, _provider):
return _FailingClient()
class ProviderUnavailableTests(unittest.IsolatedAsyncioTestCase):
async def test_chat_raises_structured_error_when_every_provider_fails(self):
client = _ClientWithFailingProviders()
with self.assertRaises(ProviderUnavailableError) as raised:
await client.chat([{"role": "user", "content": "hello"}], max_tokens=8)
self.assertCountEqual(raised.exception.providers, ("primary", "fallback"))
self.assertNotIn("api_key", str(raised.exception).lower())
async def test_stream_chat_raises_structured_error_when_every_provider_fails(self):
client = _ClientWithFailingProviders()
with self.assertRaises(ProviderUnavailableError) as raised:
async for _ in client.stream_chat([{"role": "user", "content": "hello"}], max_tokens=8):
pass
self.assertCountEqual(raised.exception.providers, ("primary", "fallback"))
self.assertNotIn("api_key", str(raised.exception).lower())
async def test_stream_chat_expands_a_role_specific_provider_pool(self):
client = _ClientWithFailingProviders()
client.providers = [
ProviderConfig(name="gemini-role", api_key="x", base_url="https://example.invalid", default_model="gemini")
]
runtime_fallback = ProviderConfig(
name="nvidia", api_key="y", base_url="https://fallback.invalid", default_model="nemotron"
)
with patch.object(client, "_load_providers", return_value=[runtime_fallback]):
with self.assertRaises(ProviderUnavailableError) as raised:
async for _ in client.stream_chat([{"role": "user", "content": "hello"}], max_tokens=8):
pass
self.assertCountEqual(raised.exception.providers, ("gemini-role", "nvidia"))
class RuntimeModelOverrideTests(unittest.TestCase):
def test_groq_runtime_model_overrides_database_model(self):
row = {
"base_url": "https://api.groq.com/openai/v1",
"default_model": "llama-3.3-70b-versatile",
}
with patch.dict("os.environ", {"GROQ_MODEL": "openai/gpt-oss-120b"}, clear=False):
self.assertEqual(
AIClient._runtime_model_override(row),
"openai/gpt-oss-120b",
)
def test_unknown_provider_keeps_database_model(self):
row = {"base_url": "https://example.invalid/v1", "default_model": "custom-model"}
self.assertEqual(AIClient._runtime_model_override(row), "custom-model")
if __name__ == "__main__":
unittest.main()
|