From 27563966fbda5324e4ae88fbd7be802b5952a14a Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Fri, 26 Sep 2025 00:46:59 +0100 Subject: [PATCH] Fix tests --- .../test_base/test_flow_parameter_specs.py | 2 +- .../test_azure_processor.py | 145 ++++++++++++++++++ .../test_llamafile_processor.py | 127 +++++++++++++++ .../test_vllm_processor.py | 143 +++++++++++++++++ 4 files changed, 416 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_base/test_flow_parameter_specs.py b/tests/unit/test_base/test_flow_parameter_specs.py index d9c3e606..db811711 100644 --- a/tests/unit/test_base/test_flow_parameter_specs.py +++ b/tests/unit/test_base/test_flow_parameter_specs.py @@ -74,7 +74,7 @@ class TestFlowParameterSpecs(IsolatedAsyncioTestCase): # Create different spec types param_spec = ParameterSpec(name="model") consumer_spec = ConsumerSpec(name="input", schema=MagicMock(), handler=MagicMock()) - producer_spec = ProducerSpec(name="output") + producer_spec = ProducerSpec(name="output", schema=MagicMock()) # Act processor.register_specification(param_spec) diff --git a/tests/unit/test_text_completion/test_azure_processor.py b/tests/unit/test_text_completion/test_azure_processor.py index a1e2ba75..529a12ab 100644 --- a/tests/unit/test_text_completion/test_azure_processor.py +++ b/tests/unit/test_text_completion/test_azure_processor.py @@ -459,5 +459,150 @@ class TestAzureProcessorSimple(IsolatedAsyncioTestCase): ) + @patch('trustgraph.model.text_completion.azure.llm.requests') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_with_model_override(self, mock_llm_init, mock_async_init, mock_requests): + """Test generate_content with model parameter override""" + # Arrange + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + 'choices': [{ + 'message': { + 'content': 'Response with model override' + } + }], + 'usage': { + 'prompt_tokens': 15, + 'completion_tokens': 10 + } + } + mock_requests.post.return_value = mock_response + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'endpoint': 'https://test.inference.ai.azure.com/v1/chat/completions', + 'token': 'test-token', + 'temperature': 0.0, + 'max_output': 4192, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act - Override model + result = await processor.generate_content("System", "Prompt", model="custom-azure-model") + + # Assert + assert result.model == "custom-azure-model" # Should use overridden model + assert result.text == "Response with model override" + + @patch('trustgraph.model.text_completion.azure.llm.requests') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_with_temperature_override(self, mock_llm_init, mock_async_init, mock_requests): + """Test generate_content with temperature parameter override""" + # Arrange + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + 'choices': [{ + 'message': { + 'content': 'Response with temperature override' + } + }], + 'usage': { + 'prompt_tokens': 15, + 'completion_tokens': 10 + } + } + mock_requests.post.return_value = mock_response + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'endpoint': 'https://test.inference.ai.azure.com/v1/chat/completions', + 'token': 'test-token', + 'temperature': 0.0, # Default temperature + 'max_output': 4192, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act - Override temperature + result = await processor.generate_content("System", "Prompt", temperature=0.8) + + # Assert + assert result.text == "Response with temperature override" + + # Verify the request was made with the overridden temperature + mock_requests.post.assert_called_once() + call_args = mock_requests.post.call_args + + import json + request_body = json.loads(call_args[1]['data']) + assert request_body['temperature'] == 0.8 + + @patch('trustgraph.model.text_completion.azure.llm.requests') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_with_both_parameters_override(self, mock_llm_init, mock_async_init, mock_requests): + """Test generate_content with both model and temperature overrides""" + # Arrange + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + 'choices': [{ + 'message': { + 'content': 'Response with both parameters override' + } + }], + 'usage': { + 'prompt_tokens': 18, + 'completion_tokens': 12 + } + } + mock_requests.post.return_value = mock_response + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'endpoint': 'https://test.inference.ai.azure.com/v1/chat/completions', + 'token': 'test-token', + 'temperature': 0.0, + 'max_output': 4192, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act - Override both parameters + result = await processor.generate_content("System", "Prompt", model="override-model", temperature=0.9) + + # Assert + assert result.model == "override-model" + assert result.text == "Response with both parameters override" + + # Verify the request was made with overridden temperature + mock_requests.post.assert_called_once() + call_args = mock_requests.post.call_args + + import json + request_body = json.loads(call_args[1]['data']) + assert request_body['temperature'] == 0.9 + + if __name__ == '__main__': pytest.main([__file__]) \ No newline at end of file diff --git a/tests/unit/test_text_completion/test_llamafile_processor.py b/tests/unit/test_text_completion/test_llamafile_processor.py index 425cfce8..410b3ff2 100644 --- a/tests/unit/test_text_completion/test_llamafile_processor.py +++ b/tests/unit/test_text_completion/test_llamafile_processor.py @@ -458,5 +458,132 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase): # No specific rate limit error handling tested since SLM presumably has no rate limits + @patch('trustgraph.model.text_completion.llamafile.llm.OpenAI') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_with_model_override(self, mock_llm_init, mock_async_init, mock_openai_class): + """Test generate_content with model parameter override""" + # Arrange + mock_openai_client = MagicMock() + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message.content = "Response from overridden model" + mock_response.usage.prompt_tokens = 15 + mock_response.usage.completion_tokens = 10 + + mock_openai_client.chat.completions.create.return_value = mock_response + mock_openai_class.return_value = mock_openai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'LLaMA_CPP', + 'llamafile': 'http://localhost:8080/v1', + 'temperature': 0.0, + 'max_output': 4096, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act - Override model + result = await processor.generate_content("System", "Prompt", model="custom-llamafile-model") + + # Assert + assert result.model == "custom-llamafile-model" # Should use overridden model + assert result.text == "Response from overridden model" + + # Verify the API call was made with overridden model + call_args = mock_openai_client.chat.completions.create.call_args + assert call_args[1]['model'] == "custom-llamafile-model" + + @patch('trustgraph.model.text_completion.llamafile.llm.OpenAI') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_with_temperature_override(self, mock_llm_init, mock_async_init, mock_openai_class): + """Test generate_content with temperature parameter override""" + # Arrange + mock_openai_client = MagicMock() + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message.content = "Response with temperature override" + mock_response.usage.prompt_tokens = 18 + mock_response.usage.completion_tokens = 12 + + mock_openai_client.chat.completions.create.return_value = mock_response + mock_openai_class.return_value = mock_openai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'LLaMA_CPP', + 'llamafile': 'http://localhost:8080/v1', + 'temperature': 0.0, # Default temperature + 'max_output': 4096, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act - Override temperature + result = await processor.generate_content("System", "Prompt", temperature=0.7) + + # Assert + assert result.text == "Response with temperature override" + + # Verify the API call was made with overridden temperature + call_args = mock_openai_client.chat.completions.create.call_args + assert call_args[1]['temperature'] == 0.7 + + @patch('trustgraph.model.text_completion.llamafile.llm.OpenAI') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_with_both_parameters_override(self, mock_llm_init, mock_async_init, mock_openai_class): + """Test generate_content with both model and temperature overrides""" + # Arrange + mock_openai_client = MagicMock() + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message.content = "Response with both parameters override" + mock_response.usage.prompt_tokens = 20 + mock_response.usage.completion_tokens = 15 + + mock_openai_client.chat.completions.create.return_value = mock_response + mock_openai_class.return_value = mock_openai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'LLaMA_CPP', + 'llamafile': 'http://localhost:8080/v1', + 'temperature': 0.0, + 'max_output': 4096, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act - Override both parameters + result = await processor.generate_content("System", "Prompt", model="override-model", temperature=0.8) + + # Assert + assert result.model == "override-model" + assert result.text == "Response with both parameters override" + + # Verify the API call was made with overridden parameters + call_args = mock_openai_client.chat.completions.create.call_args + assert call_args[1]['model'] == "override-model" + assert call_args[1]['temperature'] == 0.8 + + if __name__ == '__main__': pytest.main([__file__]) \ No newline at end of file diff --git a/tests/unit/test_text_completion/test_vllm_processor.py b/tests/unit/test_text_completion/test_vllm_processor.py index 7124c229..64da8ff9 100644 --- a/tests/unit/test_text_completion/test_vllm_processor.py +++ b/tests/unit/test_text_completion/test_vllm_processor.py @@ -485,5 +485,148 @@ class TestVLLMProcessorSimple(IsolatedAsyncioTestCase): assert call_args[1]['json']['prompt'] == expected_prompt + @patch('trustgraph.model.text_completion.vllm.llm.aiohttp.ClientSession') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_with_model_override(self, mock_llm_init, mock_async_init, mock_session_class): + """Test generate_content with model parameter override""" + # Arrange + mock_session = MagicMock() + mock_response = MagicMock() + mock_response.status = 200 + mock_response.json = AsyncMock(return_value={ + 'choices': [{ + 'text': 'Response from overridden model' + }], + 'usage': { + 'prompt_tokens': 12, + 'completion_tokens': 8 + } + }) + + mock_session.post.return_value.__aenter__.return_value = mock_response + mock_session.post.return_value.__aexit__.return_value = None + mock_session_class.return_value = mock_session + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'TheBloke/Mistral-7B-v0.1-AWQ', + 'url': 'http://vllm-service:8899/v1', + 'temperature': 0.0, + 'max_output': 2048, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act - Override model + result = await processor.generate_content("System", "Prompt", model="custom-vllm-model") + + # Assert + assert result.model == "custom-vllm-model" # Should use overridden model + assert result.text == "Response from overridden model" + + @patch('trustgraph.model.text_completion.vllm.llm.aiohttp.ClientSession') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_with_temperature_override(self, mock_llm_init, mock_async_init, mock_session_class): + """Test generate_content with temperature parameter override""" + # Arrange + mock_session = MagicMock() + mock_response = MagicMock() + mock_response.status = 200 + mock_response.json = AsyncMock(return_value={ + 'choices': [{ + 'text': 'Response with temperature override' + }], + 'usage': { + 'prompt_tokens': 15, + 'completion_tokens': 10 + } + }) + + mock_session.post.return_value.__aenter__.return_value = mock_response + mock_session.post.return_value.__aexit__.return_value = None + mock_session_class.return_value = mock_session + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'TheBloke/Mistral-7B-v0.1-AWQ', + 'url': 'http://vllm-service:8899/v1', + 'temperature': 0.0, # Default temperature + 'max_output': 2048, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act - Override temperature + result = await processor.generate_content("System", "Prompt", temperature=0.7) + + # Assert + assert result.text == "Response with temperature override" + + # Verify the request was made with overridden temperature + call_args = mock_session.post.call_args + assert call_args[1]['json']['temperature'] == 0.7 + + @patch('trustgraph.model.text_completion.vllm.llm.aiohttp.ClientSession') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_with_both_parameters_override(self, mock_llm_init, mock_async_init, mock_session_class): + """Test generate_content with both model and temperature overrides""" + # Arrange + mock_session = MagicMock() + mock_response = MagicMock() + mock_response.status = 200 + mock_response.json = AsyncMock(return_value={ + 'choices': [{ + 'text': 'Response with both parameters override' + }], + 'usage': { + 'prompt_tokens': 18, + 'completion_tokens': 12 + } + }) + + mock_session.post.return_value.__aenter__.return_value = mock_response + mock_session.post.return_value.__aexit__.return_value = None + mock_session_class.return_value = mock_session + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'TheBloke/Mistral-7B-v0.1-AWQ', + 'url': 'http://vllm-service:8899/v1', + 'temperature': 0.0, + 'max_output': 2048, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act - Override both parameters + result = await processor.generate_content("System", "Prompt", model="override-model", temperature=0.8) + + # Assert + assert result.model == "override-model" + assert result.text == "Response with both parameters override" + + # Verify the request was made with overridden temperature + call_args = mock_session.post.call_args + assert call_args[1]['json']['temperature'] == 0.8 + + if __name__ == '__main__': pytest.main([__file__]) \ No newline at end of file