jpmschweitzerandClaude Opus 4.5
a22e168666
Add comprehensive test suites for: - NPM client (27 tests) - Ollama client (16 tests) - AI client and controller (34 tests) - Static controller (8 tests) - Tools controller DNS lookup (9 tests) - OIDC authentication (10 tests) - Housekeeping endpoints (28 tests) - Infrastructure endpoints (15 tests) - Health endpoints (12 tests) - Portainer client (12 tests) - Home Assistant client (24 tests) Total: 285 tests passing with 65% code coverage. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
308 lines
11 KiB
Python
308 lines
11 KiB
Python
"""Tests for Ollama client."""
|
|
import pytest
|
|
from unittest.mock import patch, AsyncMock, MagicMock
|
|
import json
|
|
|
|
from src.models.ollama_client import OllamaClient, get_ollama_client, close_ollama_client
|
|
|
|
|
|
class TestOllamaClientInit:
|
|
"""Test OllamaClient initialization."""
|
|
|
|
@patch("src.models.ollama_client.settings")
|
|
def test_uses_settings_defaults(self, mock_settings):
|
|
"""Client should use settings for defaults."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 60
|
|
|
|
client = OllamaClient()
|
|
|
|
assert client.base_url == "http://ollama:11434"
|
|
assert client.timeout == 60
|
|
|
|
@patch("src.models.ollama_client.settings")
|
|
def test_creates_http_client(self, mock_settings):
|
|
"""Client should create httpx AsyncClient."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 30
|
|
|
|
client = OllamaClient()
|
|
|
|
assert client.client is not None
|
|
|
|
|
|
class TestOllamaClientClose:
|
|
"""Test client close functionality."""
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("src.models.ollama_client.settings")
|
|
async def test_close_closes_client(self, mock_settings):
|
|
"""close should close the HTTP client."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 30
|
|
|
|
client = OllamaClient()
|
|
|
|
with patch.object(client.client, "aclose", new_callable=AsyncMock) as mock_close:
|
|
await client.close()
|
|
mock_close.assert_called_once()
|
|
|
|
|
|
class TestOllamaClientResolveModel:
|
|
"""Test model resolution."""
|
|
|
|
@patch("src.models.ollama_client.settings")
|
|
def test_resolves_aliased_model(self, mock_settings):
|
|
"""resolve_model should map alias to actual model."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 30
|
|
mock_settings.model_aliases = {"gpt-3.5-turbo": "gemma:7b"}
|
|
|
|
client = OllamaClient()
|
|
result = client.resolve_model("gpt-3.5-turbo")
|
|
|
|
assert result == "gemma:7b"
|
|
|
|
@patch("src.models.ollama_client.settings")
|
|
def test_returns_original_if_no_alias(self, mock_settings):
|
|
"""resolve_model should return original if no alias found."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 30
|
|
mock_settings.model_aliases = {}
|
|
|
|
client = OllamaClient()
|
|
result = client.resolve_model("llama2")
|
|
|
|
assert result == "llama2"
|
|
|
|
|
|
class TestOllamaClientHealthCheck:
|
|
"""Test health check functionality."""
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("src.models.ollama_client.settings")
|
|
async def test_health_check_returns_true_on_200(self, mock_settings):
|
|
"""Health check should return True when Ollama responds 200."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 30
|
|
|
|
client = OllamaClient()
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
|
|
with patch.object(client.client, "get", new_callable=AsyncMock) as mock_get:
|
|
mock_get.return_value = mock_response
|
|
result = await client.health_check()
|
|
|
|
assert result is True
|
|
mock_get.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("src.models.ollama_client.settings")
|
|
async def test_health_check_returns_false_on_error(self, mock_settings):
|
|
"""Health check should return False on connection error."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 30
|
|
|
|
client = OllamaClient()
|
|
|
|
with patch.object(client.client, "get", new_callable=AsyncMock) as mock_get:
|
|
mock_get.side_effect = Exception("Connection refused")
|
|
result = await client.health_check()
|
|
|
|
assert result is False
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("src.models.ollama_client.settings")
|
|
async def test_health_check_returns_false_on_non_200(self, mock_settings):
|
|
"""Health check should return False on non-200 status."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 30
|
|
|
|
client = OllamaClient()
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 500
|
|
|
|
with patch.object(client.client, "get", new_callable=AsyncMock) as mock_get:
|
|
mock_get.return_value = mock_response
|
|
result = await client.health_check()
|
|
|
|
assert result is False
|
|
|
|
|
|
class TestOllamaClientListModels:
|
|
"""Test list models functionality."""
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("src.models.ollama_client.settings")
|
|
async def test_list_models_returns_dict(self, mock_settings):
|
|
"""list_models should return dict with models."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 30
|
|
|
|
client = OllamaClient()
|
|
|
|
models_data = {
|
|
"models": [
|
|
{"name": "llama2", "size": 1000000},
|
|
{"name": "gemma:7b", "size": 2000000}
|
|
]
|
|
}
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = models_data
|
|
mock_response.raise_for_status = MagicMock()
|
|
|
|
with patch.object(client.client, "get", new_callable=AsyncMock) as mock_get:
|
|
mock_get.return_value = mock_response
|
|
result = await client.list_models()
|
|
|
|
assert result == models_data
|
|
assert len(result["models"]) == 2
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("src.models.ollama_client.settings")
|
|
async def test_list_models_raises_on_error(self, mock_settings):
|
|
"""list_models should raise on error."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 30
|
|
|
|
client = OllamaClient()
|
|
|
|
with patch.object(client.client, "get", new_callable=AsyncMock) as mock_get:
|
|
mock_get.side_effect = Exception("Connection error")
|
|
|
|
with pytest.raises(Exception):
|
|
await client.list_models()
|
|
|
|
|
|
class TestOllamaClientGenerateNonStreaming:
|
|
"""Test non-streaming generation."""
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("src.models.ollama_client.settings")
|
|
async def test_generate_non_streaming_returns_response(self, mock_settings):
|
|
"""generate_non_streaming should return response dict."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 30
|
|
mock_settings.model_aliases = {}
|
|
|
|
client = OllamaClient()
|
|
|
|
response_data = {
|
|
"message": {"content": "Hello! How can I help?"},
|
|
"prompt_eval_count": 10,
|
|
"eval_count": 20
|
|
}
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = response_data
|
|
mock_response.raise_for_status = MagicMock()
|
|
|
|
with patch.object(client.client, "post", new_callable=AsyncMock) as mock_post:
|
|
mock_post.return_value = mock_response
|
|
result = await client.generate_non_streaming("llama2", "Hello")
|
|
|
|
assert result["response"] == "Hello! How can I help?"
|
|
assert result["tokens"]["prompt"] == 10
|
|
assert result["tokens"]["completion"] == 20
|
|
assert result["tokens"]["total"] == 30
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("src.models.ollama_client.settings")
|
|
async def test_generate_non_streaming_includes_max_tokens(self, mock_settings):
|
|
"""generate_non_streaming should include max_tokens in payload."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 30
|
|
mock_settings.model_aliases = {}
|
|
|
|
client = OllamaClient()
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {"message": {"content": "Hi"}}
|
|
mock_response.raise_for_status = MagicMock()
|
|
|
|
with patch.object(client.client, "post", new_callable=AsyncMock) as mock_post:
|
|
mock_post.return_value = mock_response
|
|
await client.generate_non_streaming("llama2", "Hello", max_tokens=100)
|
|
|
|
call_args = mock_post.call_args
|
|
assert call_args[1]["json"]["options"]["num_predict"] == 100
|
|
|
|
|
|
class TestOllamaClientGenerateStreaming:
|
|
"""Test streaming generation."""
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("src.models.ollama_client.settings")
|
|
async def test_generate_streaming_yields_content(self, mock_settings):
|
|
"""generate_streaming should yield content chunks."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 30
|
|
mock_settings.model_aliases = {}
|
|
|
|
client = OllamaClient()
|
|
|
|
# Create mock streaming response
|
|
async def mock_aiter_lines():
|
|
yield json.dumps({"message": {"content": "Hello"}})
|
|
yield json.dumps({"message": {"content": " world"}})
|
|
yield json.dumps({"done": True})
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.raise_for_status = MagicMock()
|
|
mock_response.aiter_lines = mock_aiter_lines
|
|
|
|
mock_stream_context = AsyncMock()
|
|
mock_stream_context.__aenter__.return_value = mock_response
|
|
mock_stream_context.__aexit__.return_value = None
|
|
|
|
with patch.object(client.client, "stream", return_value=mock_stream_context):
|
|
chunks = []
|
|
async for chunk in client.generate_streaming("llama2", "Hi"):
|
|
chunks.append(chunk)
|
|
|
|
assert "Hello" in chunks
|
|
assert " world" in chunks
|
|
|
|
|
|
class TestOllamaClientSingleton:
|
|
"""Test singleton pattern."""
|
|
|
|
@patch("src.models.ollama_client.settings")
|
|
def test_get_ollama_client_returns_same_instance(self, mock_settings):
|
|
"""get_ollama_client should return singleton."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 30
|
|
|
|
import src.models.ollama_client as module
|
|
module._ollama_client = None
|
|
|
|
client1 = get_ollama_client()
|
|
client2 = get_ollama_client()
|
|
|
|
assert client1 is client2
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("src.models.ollama_client.settings")
|
|
async def test_close_ollama_client_clears_singleton(self, mock_settings):
|
|
"""close_ollama_client should clear the singleton."""
|
|
mock_settings.ollama_base_url = "http://ollama:11434"
|
|
mock_settings.ollama_timeout = 30
|
|
|
|
import src.models.ollama_client as module
|
|
module._ollama_client = None
|
|
|
|
client = get_ollama_client()
|
|
|
|
with patch.object(client.client, "aclose", new_callable=AsyncMock):
|
|
await close_ollama_client()
|
|
|
|
assert module._ollama_client is None
|