chore: import upstream snapshot with attribution

servers/fastapi/tests/unit/test_model_availability.py

@@ -0,0 +1,98 @@
+import asyncio
+
+import pytest
+
+from models.ollama_model_status import OllamaModelStatus
+from utils import model_availability
+from utils.model_availability import check_llm_and_image_provider_api_or_model_availability
+
+
+def _set_litellm_env(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("CAN_CHANGE_KEYS", "false")
+    monkeypatch.setenv("LLM", "litellm")
+    monkeypatch.setenv("LITELLM_BASE_URL", "http://localhost:4000")
+    monkeypatch.setenv("LITELLM_MODEL", "test-model")
+
+
+def test_skips_image_provider_when_generation_disabled(monkeypatch):
+    _set_litellm_env(monkeypatch)
+    monkeypatch.setenv("DISABLE_IMAGE_GENERATION", "true")
+    monkeypatch.delenv("IMAGE_PROVIDER", raising=False)
+
+    asyncio.run(check_llm_and_image_provider_api_or_model_availability())
+
+
+def test_skips_invalid_image_provider_when_generation_disabled(monkeypatch):
+    _set_litellm_env(monkeypatch)
+    monkeypatch.setenv("DISABLE_IMAGE_GENERATION", "true")
+    monkeypatch.setenv("IMAGE_PROVIDER", "not-a-real-provider")
+
+    asyncio.run(check_llm_and_image_provider_api_or_model_availability())
+
+
+def test_requires_image_provider_when_generation_enabled(monkeypatch):
+    _set_litellm_env(monkeypatch)
+    monkeypatch.delenv("DISABLE_IMAGE_GENERATION", raising=False)
+    monkeypatch.delenv("IMAGE_PROVIDER", raising=False)
+
+    with pytest.raises(Exception, match="IMAGE_PROVIDER must be provided"):
+        asyncio.run(check_llm_and_image_provider_api_or_model_availability())
+
+
+def test_ollama_requires_selected_model_to_already_be_available(monkeypatch):
+    monkeypatch.setenv("CAN_CHANGE_KEYS", "false")
+    monkeypatch.setenv("LLM", "ollama")
+    monkeypatch.setenv("OLLAMA_MODEL", "llama3:8b")
+    monkeypatch.setenv("DISABLE_IMAGE_GENERATION", "true")
+
+    async def list_models():
+        return []
+
+    monkeypatch.setattr(model_availability, "list_available_ollama_models", list_models)
+
+    with pytest.raises(Exception, match="not available in Ollama"):
+        asyncio.run(check_llm_and_image_provider_api_or_model_availability())
+
+
+def test_ollama_accepts_selected_available_model(monkeypatch):
+    monkeypatch.setenv("CAN_CHANGE_KEYS", "false")
+    monkeypatch.setenv("LLM", "ollama")
+    monkeypatch.setenv("OLLAMA_MODEL", "llama3:8b")
+    monkeypatch.setenv("DISABLE_IMAGE_GENERATION", "true")
+
+    async def list_models():
+        return [
+            OllamaModelStatus(
+                name="llama3:8b",
+                size=1,
+                downloaded=1,
+                status="pulled",
+                done=True,
+            )
+        ]
+
+    monkeypatch.setattr(model_availability, "list_available_ollama_models", list_models)
+
+    asyncio.run(check_llm_and_image_provider_api_or_model_availability())
+
+
+def test_ollama_accepts_selected_available_experimental_model(monkeypatch):
+    monkeypatch.setenv("CAN_CHANGE_KEYS", "false")
+    monkeypatch.setenv("LLM", "ollama")
+    monkeypatch.setenv("OLLAMA_MODEL", "custom-local-model:latest")
+    monkeypatch.setenv("DISABLE_IMAGE_GENERATION", "true")
+
+    async def list_models():
+        return [
+            OllamaModelStatus(
+                name="custom-local-model:latest",
+                size=1,
+                downloaded=1,
+                status="pulled",
+                done=True,
+            )
+        ]
+
+    monkeypatch.setattr(model_availability, "list_available_ollama_models", list_models)
+
+    asyncio.run(check_llm_and_image_provider_api_or_model_availability())