From 50695da13273f61e163ba73759d50d58dbe93860 Mon Sep 17 00:00:00 2001 From: 7000pctAUTO Date: Wed, 4 Feb 2026 11:42:02 +0000 Subject: [PATCH] Fix test files to match implementation --- tests/test_ollama_client.py | 54 ++++++++++++++++++++++--------------- 1 file changed, 32 insertions(+), 22 deletions(-) diff --git a/tests/test_ollama_client.py b/tests/test_ollama_client.py index 4bbc1a0..01fcff1 100644 --- a/tests/test_ollama_client.py +++ b/tests/test_ollama_client.py @@ -1,7 +1,6 @@ """Tests for Ollama client module.""" -import pytest -from unittest.mock import Mock, patch +from unittest.mock import patch, MagicMock from shellgenius.ollama_client import OllamaClient, get_ollama_client @@ -9,10 +8,13 @@ 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" - + 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" @@ -20,21 +22,27 @@ class TestOllamaClient: 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" - + 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() == True + assert client.is_available() is 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" - + 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 = {"models": [{"name": "codellama"}, {"name": "llama2"}]} @@ -47,10 +55,13 @@ class TestOllamaClient: 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" - + 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"} @@ -58,7 +69,7 @@ class TestOllamaClient: with patch.object(client.client, 'generate', return_value=mock_response): result = client.generate("test prompt") - assert result["success"] == True + assert result["success"] is True assert "Generated text" in str(result["response"]) @@ -66,8 +77,7 @@ 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" + mock_config.return_value = {} client = get_ollama_client(host="custom:9999", model="custom-model")