76 lines
2.9 KiB
Python
76 lines
2.9 KiB
Python
"""Tests for Ollama client module."""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock, patch
|
|
|
|
from shellgenius.ollama_client import OllamaClient, get_ollama_client
|
|
|
|
|
|
class TestOllamaClient:
|
|
def test_init(self):
|
|
"""Test client initialization."""
|
|
with patch('shellgenius.ollama_client.get_config') as mock_config:
|
|
mock_config.return_value.ollama_host = "localhost:11434"
|
|
mock_config.return_value.ollama_model = "codellama"
|
|
|
|
client = OllamaClient()
|
|
|
|
assert client.host == "localhost:11434"
|
|
assert client.model == "codellama"
|
|
|
|
def test_is_available(self):
|
|
"""Test availability check."""
|
|
with patch('shellgenius.ollama_client.get_config') as mock_config:
|
|
mock_config.return_value.ollama_host = "localhost:11434"
|
|
mock_config.return_value.ollama_model = "codellama"
|
|
|
|
client = OllamaClient()
|
|
|
|
with patch.object(client, 'list_models', return_value=["codellama"]):
|
|
assert client.is_available() == True
|
|
|
|
def test_list_models(self):
|
|
"""Test listing models."""
|
|
with patch('shellgenius.ollama_client.get_config') as mock_config:
|
|
mock_config.return_value.ollama_host = "localhost:11434"
|
|
mock_config.return_value.ollama_model = "codellama"
|
|
|
|
client = OllamaClient()
|
|
|
|
mock_response = {"models": [{"name": "codellama"}, {"name": "llama2"}]}
|
|
|
|
with patch.object(client.client, 'list', return_value=mock_response):
|
|
models = client.list_models()
|
|
|
|
assert len(models) == 2
|
|
assert "codellama" in models
|
|
|
|
def test_generate(self):
|
|
"""Test text generation."""
|
|
with patch('shellgenius.ollama_client.get_config') as mock_config:
|
|
mock_config.return_value.ollama_host = "localhost:11434"
|
|
mock_config.return_value.ollama_model = "codellama"
|
|
|
|
client = OllamaClient()
|
|
|
|
mock_response = {"response": "Generated text"}
|
|
|
|
with patch.object(client.client, 'generate', return_value=mock_response):
|
|
result = client.generate("test prompt")
|
|
|
|
assert result["success"] == True
|
|
assert "Generated text" in str(result["response"])
|
|
|
|
|
|
class TestGetOllamaClient:
|
|
def test_convenience_function(self):
|
|
"""Test the convenience function for getting client."""
|
|
with patch('shellgenius.ollama_client.get_config') as mock_config:
|
|
mock_config.return_value.ollama_host = "localhost:11434"
|
|
mock_config.return_value.ollama_model = "custom-model"
|
|
|
|
client = get_ollama_client(host="custom:9999", model="custom-model")
|
|
|
|
assert client.host == "custom:9999"
|
|
assert client.model == "custom-model"
|