From 209505316f5ce74bc4ddbcb4a607c525957c5044 Mon Sep 17 00:00:00 2001 From: 7000pctAUTO Date: Wed, 4 Feb 2026 11:22:17 +0000 Subject: [PATCH] fix: remove unused imports from test files --- .tests/test_ollama_client.py | 131 +++++++++++++++++++++++++++++++++++ 1 file changed, 131 insertions(+) create mode 100644 .tests/test_ollama_client.py diff --git a/.tests/test_ollama_client.py b/.tests/test_ollama_client.py new file mode 100644 index 0000000..806a80a --- /dev/null +++ b/.tests/test_ollama_client.py @@ -0,0 +1,131 @@ +"""Tests for Ollama client module.""" + +from unittest.mock import Mock, patch, MagicMock + +from shellgenius.ollama_client import OllamaClient, get_ollama_client + + +class TestOllamaClient: + """Test Ollama client functionality.""" + + @patch('shellgenius.ollama_client.ollama.Client') + def test_client_initialization(self, mock_client_class): + """Test client initialization.""" + 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_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() + result = client.is_available() + assert result is False + + +class TestGetOllamaClient: + """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"