diff --git a/tests/unit/test_embeddings/test_ollama_dynamic_model.py b/tests/unit/test_embeddings/test_ollama_dynamic_model.py index d52a58c6..cfbc4d6e 100644 --- a/tests/unit/test_embeddings/test_ollama_dynamic_model.py +++ b/tests/unit/test_embeddings/test_ollama_dynamic_model.py @@ -14,13 +14,13 @@ from trustgraph.embeddings.ollama.processor import Processor class TestOllamaDynamicModelLoading(IsolatedAsyncioTestCase): """Test Ollama dynamic model selection""" - @patch('trustgraph.embeddings.ollama.processor.Client') + @patch('trustgraph.embeddings.ollama.processor.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__') async def test_client_initialized_with_host(self, mock_embeddings_init, mock_async_init, mock_client_class): """Test that Ollama client is initialized with correct host""" # Arrange - mock_ollama_client = Mock() + mock_ollama_client = AsyncMock() mock_response = Mock() mock_response.embeddings = [[0.1, 0.2, 0.3, 0.4, 0.5]] mock_ollama_client.embed.return_value = mock_response @@ -36,13 +36,13 @@ class TestOllamaDynamicModelLoading(IsolatedAsyncioTestCase): mock_client_class.assert_called_once_with(host="http://localhost:11434") assert processor.default_model == "test-model" - @patch('trustgraph.embeddings.ollama.processor.Client') + @patch('trustgraph.embeddings.ollama.processor.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__') async def test_on_embeddings_uses_default_model(self, mock_embeddings_init, mock_async_init, mock_client_class): """Test that on_embeddings uses default model when no model specified""" # Arrange - mock_ollama_client = Mock() + mock_ollama_client = AsyncMock() mock_response = Mock() mock_response.embeddings = [[0.1, 0.2, 0.3, 0.4, 0.5]] mock_ollama_client.embed.return_value = mock_response @@ -62,13 +62,13 @@ class TestOllamaDynamicModelLoading(IsolatedAsyncioTestCase): ) assert result == [[0.1, 0.2, 0.3, 0.4, 0.5]] - @patch('trustgraph.embeddings.ollama.processor.Client') + @patch('trustgraph.embeddings.ollama.processor.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__') async def test_on_embeddings_uses_specified_model(self, mock_embeddings_init, mock_async_init, mock_client_class): """Test that on_embeddings uses specified model when provided""" # Arrange - mock_ollama_client = Mock() + mock_ollama_client = AsyncMock() mock_response = Mock() mock_response.embeddings = [[0.1, 0.2, 0.3, 0.4, 0.5]] mock_ollama_client.embed.return_value = mock_response @@ -88,13 +88,13 @@ class TestOllamaDynamicModelLoading(IsolatedAsyncioTestCase): ) assert result == [[0.1, 0.2, 0.3, 0.4, 0.5]] - @patch('trustgraph.embeddings.ollama.processor.Client') + @patch('trustgraph.embeddings.ollama.processor.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__') async def test_multiple_model_switches(self, mock_embeddings_init, mock_async_init, mock_client_class): """Test switching between multiple models""" # Arrange - mock_ollama_client = Mock() + mock_ollama_client = AsyncMock() mock_response = Mock() mock_response.embeddings = [[0.1, 0.2, 0.3, 0.4, 0.5]] mock_ollama_client.embed.return_value = mock_response @@ -118,13 +118,13 @@ class TestOllamaDynamicModelLoading(IsolatedAsyncioTestCase): assert calls[2][1]['model'] == "model-a" assert calls[3][1]['model'] == "test-model" # Default - @patch('trustgraph.embeddings.ollama.processor.Client') + @patch('trustgraph.embeddings.ollama.processor.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__') async def test_none_model_uses_default(self, mock_embeddings_init, mock_async_init, mock_client_class): """Test that None model parameter falls back to default""" # Arrange - mock_ollama_client = Mock() + mock_ollama_client = AsyncMock() mock_response = Mock() mock_response.embeddings = [[0.1, 0.2, 0.3, 0.4, 0.5]] mock_ollama_client.embed.return_value = mock_response @@ -143,13 +143,13 @@ class TestOllamaDynamicModelLoading(IsolatedAsyncioTestCase): input=["test text"] ) - @patch('trustgraph.embeddings.ollama.processor.Client') + @patch('trustgraph.embeddings.ollama.processor.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__') async def test_initialization_without_model_uses_default(self, mock_embeddings_init, mock_async_init, mock_client_class): """Test initialization without model parameter uses module default""" # Arrange - mock_ollama_client = Mock() + mock_ollama_client = AsyncMock() mock_client_class.return_value = mock_ollama_client mock_async_init.return_value = None mock_embeddings_init.return_value = None diff --git a/tests/unit/test_text_completion/test_ollama_processor.py b/tests/unit/test_text_completion/test_ollama_processor.py index 69baf85f..35bf182a 100644 --- a/tests/unit/test_text_completion/test_ollama_processor.py +++ b/tests/unit/test_text_completion/test_ollama_processor.py @@ -15,13 +15,13 @@ from trustgraph.base import LlmResult class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): """Test Ollama processor functionality""" - @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.model.text_completion.ollama.llm.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.llm_service.LlmService.__init__') async def test_processor_initialization_basic(self, mock_llm_init, mock_async_init, mock_client_class): """Test basic processor initialization""" # Arrange - mock_client = MagicMock() + mock_client = AsyncMock() mock_client_class.return_value = mock_client # Mock the parent class initialization @@ -44,13 +44,13 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): assert hasattr(processor, 'llm') mock_client_class.assert_called_once_with(host='http://localhost:11434') - @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.model.text_completion.ollama.llm.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.llm_service.LlmService.__init__') async def test_generate_content_success(self, mock_llm_init, mock_async_init, mock_client_class): """Test successful content generation""" # Arrange - mock_client = MagicMock() + mock_client = AsyncMock() mock_response = { 'response': 'Generated response from Ollama', 'prompt_eval_count': 15, @@ -83,13 +83,13 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): assert result.model == 'llama2' mock_client.generate.assert_called_once_with('llama2', "System prompt\n\nUser prompt", options={'temperature': 0.0}) - @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.model.text_completion.ollama.llm.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.llm_service.LlmService.__init__') async def test_generate_content_generic_exception(self, mock_llm_init, mock_async_init, mock_client_class): """Test handling of generic exceptions""" # Arrange - mock_client = MagicMock() + mock_client = AsyncMock() mock_client.generate.side_effect = Exception("Connection error") mock_client_class.return_value = mock_client @@ -110,13 +110,13 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): with pytest.raises(Exception, match="Connection error"): await processor.generate_content("System prompt", "User prompt") - @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.model.text_completion.ollama.llm.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.llm_service.LlmService.__init__') async def test_processor_initialization_with_custom_parameters(self, mock_llm_init, mock_async_init, mock_client_class): """Test processor initialization with custom parameters""" # Arrange - mock_client = MagicMock() + mock_client = AsyncMock() mock_client_class.return_value = mock_client mock_async_init.return_value = None @@ -137,13 +137,13 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): assert processor.default_model == 'mistral' mock_client_class.assert_called_once_with(host='http://192.168.1.100:11434') - @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.model.text_completion.ollama.llm.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.llm_service.LlmService.__init__') async def test_processor_initialization_with_defaults(self, mock_llm_init, mock_async_init, mock_client_class): """Test processor initialization with default values""" # Arrange - mock_client = MagicMock() + mock_client = AsyncMock() mock_client_class.return_value = mock_client mock_async_init.return_value = None @@ -164,13 +164,13 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): # Should use default_ollama (http://localhost:11434 or from OLLAMA_HOST env) mock_client_class.assert_called_once() - @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.model.text_completion.ollama.llm.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.llm_service.LlmService.__init__') async def test_generate_content_empty_prompts(self, mock_llm_init, mock_async_init, mock_client_class): """Test content generation with empty prompts""" # Arrange - mock_client = MagicMock() + mock_client = AsyncMock() mock_response = { 'response': 'Default response', 'prompt_eval_count': 2, @@ -205,13 +205,13 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): # The prompt should be "" + "\n\n" + "" = "\n\n" mock_client.generate.assert_called_once_with('llama2', "\n\n", options={'temperature': 0.0}) - @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.model.text_completion.ollama.llm.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.llm_service.LlmService.__init__') async def test_generate_content_token_counting(self, mock_llm_init, mock_async_init, mock_client_class): """Test token counting from Ollama response""" # Arrange - mock_client = MagicMock() + mock_client = AsyncMock() mock_response = { 'response': 'Test response', 'prompt_eval_count': 50, @@ -243,13 +243,13 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): assert result.out_token == 25 assert result.model == 'llama2' - @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.model.text_completion.ollama.llm.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.llm_service.LlmService.__init__') async def test_ollama_client_initialization(self, mock_llm_init, mock_async_init, mock_client_class): """Test that Ollama client is initialized correctly""" # Arrange - mock_client = MagicMock() + mock_client = AsyncMock() mock_client_class.return_value = mock_client mock_async_init.return_value = None @@ -273,13 +273,13 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): # Verify processor has the client assert processor.llm == mock_client - @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.model.text_completion.ollama.llm.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.llm_service.LlmService.__init__') async def test_generate_content_prompt_construction(self, mock_llm_init, mock_async_init, mock_client_class): """Test prompt construction with system and user prompts""" # Arrange - mock_client = MagicMock() + mock_client = AsyncMock() mock_response = { 'response': 'Response with system instructions', 'prompt_eval_count': 25, @@ -312,13 +312,13 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): # Verify the combined prompt mock_client.generate.assert_called_once_with('llama2', "You are a helpful assistant\n\nWhat is AI?", options={'temperature': 0.0}) - @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.model.text_completion.ollama.llm.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.llm_service.LlmService.__init__') async def test_generate_content_temperature_override(self, mock_llm_init, mock_async_init, mock_client_class): """Test temperature parameter override functionality""" # Arrange - mock_client = MagicMock() + mock_client = AsyncMock() mock_response = { 'response': 'Response with custom temperature', 'prompt_eval_count': 20, @@ -360,13 +360,13 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): options={'temperature': 0.8} # Should use runtime override ) - @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.model.text_completion.ollama.llm.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.llm_service.LlmService.__init__') async def test_generate_content_model_override(self, mock_llm_init, mock_async_init, mock_client_class): """Test model parameter override functionality""" # Arrange - mock_client = MagicMock() + mock_client = AsyncMock() mock_response = { 'response': 'Response with custom model', 'prompt_eval_count': 18, @@ -408,13 +408,13 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): options={'temperature': 0.1} # Should use processor default ) - @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.model.text_completion.ollama.llm.AsyncClient') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.llm_service.LlmService.__init__') async def test_generate_content_both_parameters_override(self, mock_llm_init, mock_async_init, mock_client_class): """Test overriding both model and temperature parameters simultaneously""" # Arrange - mock_client = MagicMock() + mock_client = AsyncMock() mock_response = { 'response': 'Response with both overrides', 'prompt_eval_count': 22,