fix: resolve CI test failures (config access, mocks, imports)
This commit is contained in:
@@ -1,85 +1,131 @@
|
||||
"""Tests for Ollama client module."""
|
||||
|
||||
from unittest.mock import patch, MagicMock
|
||||
from unittest.mock import Mock, patch, MagicMock
|
||||
|
||||
from shellgenius.ollama_client import OllamaClient, get_ollama_client
|
||||
|
||||
|
||||
class TestOllamaClient:
|
||||
def test_init(self):
|
||||
"""Test Ollama client functionality."""
|
||||
|
||||
@patch('shellgenius.ollama_client.ollama.Client')
|
||||
def test_client_initialization(self, mock_client_class):
|
||||
"""Test client initialization."""
|
||||
mock_config = MagicMock()
|
||||
mock_config.get.side_effect = lambda key, default=None: {
|
||||
"ollama_host": "localhost:11434",
|
||||
"ollama_model": "codellama",
|
||||
}.get(key, default)
|
||||
|
||||
with patch('shellgenius.ollama_client.get_config', return_value=mock_config):
|
||||
client = OllamaClient()
|
||||
|
||||
assert client.host == "localhost:11434"
|
||||
assert client.model == "codellama"
|
||||
|
||||
def test_is_available(self):
|
||||
"""Test availability check."""
|
||||
mock_config = MagicMock()
|
||||
mock_config.get.side_effect = lambda key, default=None: {
|
||||
"ollama_host": "localhost:11434",
|
||||
"ollama_model": "codellama",
|
||||
}.get(key, default)
|
||||
|
||||
with patch('shellgenius.ollama_client.get_config', return_value=mock_config):
|
||||
client = OllamaClient()
|
||||
|
||||
with patch.object(client, 'list_models', return_value=["codellama"]):
|
||||
assert client.is_available() is True
|
||||
|
||||
def test_list_models(self):
|
||||
mock_client = Mock()
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
client = OllamaClient(host="localhost:11434", model="codellama")
|
||||
|
||||
assert client.host == "localhost:11434"
|
||||
assert client.model == "codellama"
|
||||
|
||||
@patch('shellgenius.ollama_client.ollama.Client')
|
||||
def test_list_models(self, mock_client_class):
|
||||
"""Test listing models."""
|
||||
mock_config = MagicMock()
|
||||
mock_config.get.side_effect = lambda key, default=None: {
|
||||
"ollama_host": "localhost:11434",
|
||||
"ollama_model": "codellama",
|
||||
}.get(key, default)
|
||||
|
||||
with patch('shellgenius.ollama_client.get_config', return_value=mock_config):
|
||||
mock_client = Mock()
|
||||
mock_client.list.return_value = {
|
||||
"models": [
|
||||
{"name": "codellama:latest"},
|
||||
{"name": "llama2:latest"},
|
||||
]
|
||||
}
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
client = OllamaClient()
|
||||
models = client.list_models()
|
||||
|
||||
assert len(models) == 2
|
||||
assert "codellama:latest" in models
|
||||
|
||||
@patch('shellgenius.ollama_client.ollama.Client')
|
||||
def test_generate_success(self, mock_client_class):
|
||||
"""Test successful generation."""
|
||||
mock_client = Mock()
|
||||
mock_client.generate.return_value = {
|
||||
"response": "echo hello"
|
||||
}
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
client = OllamaClient()
|
||||
result = client.generate("test prompt")
|
||||
|
||||
assert result["success"] is True
|
||||
assert result["response"]["response"] == "echo hello"
|
||||
|
||||
@patch('shellgenius.ollama_client.ollama.Client')
|
||||
def test_generate_failure(self, mock_client_class):
|
||||
"""Test failed generation."""
|
||||
mock_client = Mock()
|
||||
mock_client.generate.side_effect = Exception("Connection failed")
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
client = OllamaClient()
|
||||
result = client.generate("test prompt")
|
||||
|
||||
assert result["success"] is False
|
||||
assert "error" in result
|
||||
|
||||
@patch('shellgenius.ollama_client.ollama.Client')
|
||||
def test_chat_success(self, mock_client_class):
|
||||
"""Test successful chat."""
|
||||
mock_client = Mock()
|
||||
mock_client.chat.return_value = {
|
||||
"message": {"content": "Hello!"}
|
||||
}
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
client = OllamaClient()
|
||||
messages = [{"role": "user", "content": "Hi"}]
|
||||
result = client.chat(messages)
|
||||
|
||||
assert result["success"] is True
|
||||
|
||||
@patch('shellgenius.ollama_client.ollama.Client')
|
||||
def test_pull_model(self, mock_client_class):
|
||||
"""Test pulling a model."""
|
||||
mock_client = Mock()
|
||||
mock_client.pull.return_value = {"status": "success"}
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
client = OllamaClient()
|
||||
result = client.pull_model("llama2")
|
||||
|
||||
assert result is True
|
||||
mock_client.pull.assert_called_once_with("llama2")
|
||||
|
||||
@patch('shellgenius.ollama_client.ollama.Client')
|
||||
def test_is_available_true(self, mock_client_class):
|
||||
"""Test availability check when available."""
|
||||
mock_client = Mock()
|
||||
mock_client.list.return_value = {"models": []}
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
client = OllamaClient()
|
||||
assert client.is_available() is True
|
||||
|
||||
@patch('shellgenius.ollama_client.ollama.Client')
|
||||
def test_is_available_false(self, mock_client_class):
|
||||
"""Test availability check when not available."""
|
||||
mock_instance = MagicMock()
|
||||
mock_instance.list.side_effect = Exception("Connection refused")
|
||||
mock_client_class.return_value = mock_instance
|
||||
|
||||
with patch.object(OllamaClient, 'list_models', side_effect=Exception("Connection refused")):
|
||||
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."""
|
||||
mock_config = MagicMock()
|
||||
mock_config.get.side_effect = lambda key, default=None: {
|
||||
"ollama_host": "localhost:11434",
|
||||
"ollama_model": "codellama",
|
||||
}.get(key, default)
|
||||
|
||||
with patch('shellgenius.ollama_client.get_config', return_value=mock_config):
|
||||
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"] is True
|
||||
assert "Generated text" in str(result["response"])
|
||||
result = client.is_available()
|
||||
assert result is False
|
||||
|
||||
|
||||
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 = {}
|
||||
|
||||
client = get_ollama_client(host="custom:9999", model="custom-model")
|
||||
|
||||
assert client.host == "custom:9999"
|
||||
assert client.model == "custom-model"
|
||||
"""Test get_ollama_client convenience function."""
|
||||
|
||||
@patch('shellgenius.ollama_client.ollama.Client')
|
||||
def test_get_client(self, mock_client_class):
|
||||
"""Test getting client instance."""
|
||||
mock_client = Mock()
|
||||
mock_client_class.return_value = mock_client
|
||||
|
||||
client = get_ollama_client(host="localhost:11434")
|
||||
|
||||
assert isinstance(client, OllamaClient)
|
||||
assert client.host == "localhost:11434"
|
||||
|
||||
Reference in New Issue
Block a user