wehub-resource-sync
99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
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())
|