From c169d8f0a913a352b2b5bd1d25643b31be98193b Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Thu, 25 Sep 2025 19:18:39 +0100 Subject: [PATCH] Fix some tests, update to new model parameter usage --- .../test_azure_openai_processor.py | 6 +++--- .../test_text_completion/test_azure_processor.py | 6 +++--- .../test_text_completion/test_claude_processor.py | 8 ++++---- .../test_text_completion/test_cohere_processor.py | 8 ++++---- .../test_googleaistudio_processor.py | 8 ++++---- .../test_llamafile_processor.py | 14 +++++++------- .../test_text_completion/test_ollama_processor.py | 6 +++--- .../test_text_completion/test_openai_processor.py | 6 +++--- .../test_vertexai_processor.py | 8 ++++---- .../test_text_completion/test_vllm_processor.py | 6 +++--- 10 files changed, 38 insertions(+), 38 deletions(-) diff --git a/tests/unit/test_text_completion/test_azure_openai_processor.py b/tests/unit/test_text_completion/test_azure_openai_processor.py index b5669907..967a3893 100644 --- a/tests/unit/test_text_completion/test_azure_openai_processor.py +++ b/tests/unit/test_text_completion/test_azure_openai_processor.py @@ -44,7 +44,7 @@ class TestAzureOpenAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gpt-4' + assert processor.default_model == 'gpt-4' assert processor.temperature == 0.0 assert processor.max_output == 4192 assert hasattr(processor, 'openai') @@ -254,7 +254,7 @@ class TestAzureOpenAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gpt-35-turbo' + assert processor.default_model == 'gpt-35-turbo' assert processor.temperature == 0.7 assert processor.max_output == 2048 mock_azure_openai_class.assert_called_once_with( @@ -289,7 +289,7 @@ class TestAzureOpenAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gpt-4' + assert processor.default_model == 'gpt-4' assert processor.temperature == 0.0 # default_temperature assert processor.max_output == 4192 # default_max_output mock_azure_openai_class.assert_called_once_with( diff --git a/tests/unit/test_text_completion/test_azure_processor.py b/tests/unit/test_text_completion/test_azure_processor.py index 6ef78a2c..a1e2ba75 100644 --- a/tests/unit/test_text_completion/test_azure_processor.py +++ b/tests/unit/test_text_completion/test_azure_processor.py @@ -43,7 +43,7 @@ class TestAzureProcessorSimple(IsolatedAsyncioTestCase): assert processor.token == 'test-token' assert processor.temperature == 0.0 assert processor.max_output == 4192 - assert processor.model == 'AzureAI' + assert processor.default_model == 'AzureAI' @patch('trustgraph.model.text_completion.azure.llm.requests') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @@ -261,7 +261,7 @@ class TestAzureProcessorSimple(IsolatedAsyncioTestCase): assert processor.token == 'custom-token' assert processor.temperature == 0.7 assert processor.max_output == 2048 - assert processor.model == 'AzureAI' + assert processor.default_model == 'AzureAI' @patch('trustgraph.model.text_completion.azure.llm.requests') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @@ -289,7 +289,7 @@ class TestAzureProcessorSimple(IsolatedAsyncioTestCase): assert processor.token == 'test-token' assert processor.temperature == 0.0 # default_temperature assert processor.max_output == 4192 # default_max_output - assert processor.model == 'AzureAI' # default_model + assert processor.default_model == 'AzureAI' # default_model @patch('trustgraph.model.text_completion.azure.llm.requests') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') diff --git a/tests/unit/test_text_completion/test_claude_processor.py b/tests/unit/test_text_completion/test_claude_processor.py index 27a18b93..12951019 100644 --- a/tests/unit/test_text_completion/test_claude_processor.py +++ b/tests/unit/test_text_completion/test_claude_processor.py @@ -42,7 +42,7 @@ class TestClaudeProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'claude-3-5-sonnet-20240620' + assert processor.default_model == 'claude-3-5-sonnet-20240620' assert processor.temperature == 0.0 assert processor.max_output == 8192 assert hasattr(processor, 'claude') @@ -217,7 +217,7 @@ class TestClaudeProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'claude-3-haiku-20240307' + assert processor.default_model == 'claude-3-haiku-20240307' assert processor.temperature == 0.7 assert processor.max_output == 4096 mock_anthropic_class.assert_called_once_with(api_key='custom-api-key') @@ -246,7 +246,7 @@ class TestClaudeProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'claude-3-5-sonnet-20240620' # default_model + assert processor.default_model == 'claude-3-5-sonnet-20240620' # default_model assert processor.temperature == 0.0 # default_temperature assert processor.max_output == 8192 # default_max_output mock_anthropic_class.assert_called_once_with(api_key='test-api-key') @@ -433,7 +433,7 @@ class TestClaudeProcessorSimple(IsolatedAsyncioTestCase): # Verify processor has the client assert processor.claude == mock_claude_client - assert processor.model == 'claude-3-opus-20240229' + assert processor.default_model == 'claude-3-opus-20240229' if __name__ == '__main__': diff --git a/tests/unit/test_text_completion/test_cohere_processor.py b/tests/unit/test_text_completion/test_cohere_processor.py index ebb6b626..9e4397bc 100644 --- a/tests/unit/test_text_completion/test_cohere_processor.py +++ b/tests/unit/test_text_completion/test_cohere_processor.py @@ -41,7 +41,7 @@ class TestCohereProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'c4ai-aya-23-8b' + assert processor.default_model == 'c4ai-aya-23-8b' assert processor.temperature == 0.0 assert hasattr(processor, 'cohere') mock_cohere_class.assert_called_once_with(api_key='test-api-key') @@ -201,7 +201,7 @@ class TestCohereProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'command-light' + assert processor.default_model == 'command-light' assert processor.temperature == 0.7 mock_cohere_class.assert_called_once_with(api_key='custom-api-key') @@ -229,7 +229,7 @@ class TestCohereProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'c4ai-aya-23-8b' # default_model + assert processor.default_model == 'c4ai-aya-23-8b' # default_model assert processor.temperature == 0.0 # default_temperature mock_cohere_class.assert_called_once_with(api_key='test-api-key') @@ -395,7 +395,7 @@ class TestCohereProcessorSimple(IsolatedAsyncioTestCase): # Verify processor has the client assert processor.cohere == mock_cohere_client - assert processor.model == 'command-r' + assert processor.default_model == 'command-r' @patch('trustgraph.model.text_completion.cohere.llm.cohere.Client') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') diff --git a/tests/unit/test_text_completion/test_googleaistudio_processor.py b/tests/unit/test_text_completion/test_googleaistudio_processor.py index a3ca0057..f31715d2 100644 --- a/tests/unit/test_text_completion/test_googleaistudio_processor.py +++ b/tests/unit/test_text_completion/test_googleaistudio_processor.py @@ -42,7 +42,7 @@ class TestGoogleAIStudioProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gemini-2.0-flash-001' + assert processor.default_model == 'gemini-2.0-flash-001' assert processor.temperature == 0.0 assert processor.max_output == 8192 assert hasattr(processor, 'client') @@ -205,7 +205,7 @@ class TestGoogleAIStudioProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gemini-1.5-pro' + assert processor.default_model == 'gemini-1.5-pro' assert processor.temperature == 0.7 assert processor.max_output == 4096 mock_genai_class.assert_called_once_with(api_key='custom-api-key') @@ -234,7 +234,7 @@ class TestGoogleAIStudioProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gemini-2.0-flash-001' # default_model + assert processor.default_model == 'gemini-2.0-flash-001' # default_model assert processor.temperature == 0.0 # default_temperature assert processor.max_output == 8192 # default_max_output mock_genai_class.assert_called_once_with(api_key='test-api-key') @@ -431,7 +431,7 @@ class TestGoogleAIStudioProcessorSimple(IsolatedAsyncioTestCase): # Verify processor has the client assert processor.client == mock_genai_client - assert processor.model == 'gemini-1.5-flash' + assert processor.default_model == 'gemini-1.5-flash' @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') diff --git a/tests/unit/test_text_completion/test_llamafile_processor.py b/tests/unit/test_text_completion/test_llamafile_processor.py index bae1a4bb..9aabb0ea 100644 --- a/tests/unit/test_text_completion/test_llamafile_processor.py +++ b/tests/unit/test_text_completion/test_llamafile_processor.py @@ -42,7 +42,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'LLaMA_CPP' + assert processor.default_model == 'LLaMA_CPP' assert processor.llamafile == 'http://localhost:8080/v1' assert processor.temperature == 0.0 assert processor.max_output == 4096 @@ -91,7 +91,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase): assert result.text == "Generated response from LlamaFile" assert result.in_token == 20 assert result.out_token == 12 - assert result.model == 'llama.cpp' # Note: model in result is hardcoded to 'llama.cpp' + assert result.model == 'LLaMA_CPP' # Uses the default model name # Verify the OpenAI API call structure mock_openai_client.chat.completions.create.assert_called_once_with( @@ -157,7 +157,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'custom-llama' + assert processor.default_model == 'custom-llama' assert processor.llamafile == 'http://custom-host:8080/v1' assert processor.temperature == 0.7 assert processor.max_output == 2048 @@ -189,7 +189,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'LLaMA_CPP' # default_model + assert processor.default_model == 'LLaMA_CPP' # default_model assert processor.llamafile == 'http://localhost:8080/v1' # default_llamafile assert processor.temperature == 0.0 # default_temperature assert processor.max_output == 4096 # default_max_output @@ -237,7 +237,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase): assert result.text == "Default response" assert result.in_token == 2 assert result.out_token == 3 - assert result.model == 'llama.cpp' + assert result.model == 'LLaMA_CPP' # Verify the combined prompt is sent correctly call_args = mock_openai_client.chat.completions.create.call_args @@ -408,8 +408,8 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase): result = await processor.generate_content("System", "User") # Assert - assert result.model == 'llama.cpp' # Should always be 'llama.cpp', not 'custom-model-name' - assert processor.model == 'custom-model-name' # But processor.model should still be custom + assert result.model == 'custom-model-name' # Uses the actual model name passed to generate_content + assert processor.default_model == 'custom-model-name' # But processor.model should still be custom @patch('trustgraph.model.text_completion.llamafile.llm.OpenAI') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') diff --git a/tests/unit/test_text_completion/test_ollama_processor.py b/tests/unit/test_text_completion/test_ollama_processor.py index e846ec12..9fd5c99e 100644 --- a/tests/unit/test_text_completion/test_ollama_processor.py +++ b/tests/unit/test_text_completion/test_ollama_processor.py @@ -40,7 +40,7 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'llama2' + assert processor.default_model == 'llama2' assert hasattr(processor, 'llm') mock_client_class.assert_called_once_with(host='http://localhost:11434') @@ -134,7 +134,7 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'mistral' + 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') @@ -160,7 +160,7 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gemma2:9b' # default_model + assert processor.default_model == 'gemma2:9b' # default_model # Should use default_ollama (http://localhost:11434 or from OLLAMA_HOST env) mock_client_class.assert_called_once() diff --git a/tests/unit/test_text_completion/test_openai_processor.py b/tests/unit/test_text_completion/test_openai_processor.py index 504dad50..56425285 100644 --- a/tests/unit/test_text_completion/test_openai_processor.py +++ b/tests/unit/test_text_completion/test_openai_processor.py @@ -43,7 +43,7 @@ class TestOpenAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gpt-3.5-turbo' + assert processor.default_model == 'gpt-3.5-turbo' assert processor.temperature == 0.0 assert processor.max_output == 4096 assert hasattr(processor, 'openai') @@ -222,7 +222,7 @@ class TestOpenAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gpt-4' + assert processor.default_model == 'gpt-4' assert processor.temperature == 0.7 assert processor.max_output == 2048 mock_openai_class.assert_called_once_with(base_url='https://custom-openai-url.com/v1', api_key='custom-api-key') @@ -251,7 +251,7 @@ class TestOpenAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gpt-3.5-turbo' # default_model + assert processor.default_model == 'gpt-3.5-turbo' # default_model assert processor.temperature == 0.0 # default_temperature assert processor.max_output == 4096 # default_max_output mock_openai_class.assert_called_once_with(base_url='https://api.openai.com/v1', api_key='test-api-key') diff --git a/tests/unit/test_text_completion/test_vertexai_processor.py b/tests/unit/test_text_completion/test_vertexai_processor.py index 3910a30c..37e8ac01 100644 --- a/tests/unit/test_text_completion/test_vertexai_processor.py +++ b/tests/unit/test_text_completion/test_vertexai_processor.py @@ -47,7 +47,7 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gemini-2.0-flash-001' # It's stored as 'model', not 'model_name' + assert processor.default_model == 'gemini-2.0-flash-001' # It's stored as 'model', not 'model_name' assert hasattr(processor, 'generation_config') assert hasattr(processor, 'safety_settings') assert hasattr(processor, 'llm') @@ -223,7 +223,7 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gemini-2.0-flash-001' + assert processor.default_model == 'gemini-2.0-flash-001' mock_auth_default.assert_called_once() mock_vertexai.init.assert_called_once_with( location='us-central1', @@ -296,7 +296,7 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gemini-1.5-pro' + assert processor.default_model == 'gemini-1.5-pro' # Verify that generation_config object exists (can't easily check internal values) assert hasattr(processor, 'generation_config') @@ -440,7 +440,7 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'claude-3-sonnet@20240229' + assert processor.default_model == 'claude-3-sonnet@20240229' assert processor.is_anthropic == True # Verify service account was called with private key diff --git a/tests/unit/test_text_completion/test_vllm_processor.py b/tests/unit/test_text_completion/test_vllm_processor.py index 7d30cf74..7124c229 100644 --- a/tests/unit/test_text_completion/test_vllm_processor.py +++ b/tests/unit/test_text_completion/test_vllm_processor.py @@ -42,7 +42,7 @@ class TestVLLMProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'TheBloke/Mistral-7B-v0.1-AWQ' + assert processor.default_model == 'TheBloke/Mistral-7B-v0.1-AWQ' assert processor.base_url == 'http://vllm-service:8899/v1' assert processor.temperature == 0.0 assert processor.max_output == 2048 @@ -199,7 +199,7 @@ class TestVLLMProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'custom-model' + assert processor.default_model == 'custom-model' assert processor.base_url == 'http://custom-vllm:8080/v1' assert processor.temperature == 0.7 assert processor.max_output == 1024 @@ -228,7 +228,7 @@ class TestVLLMProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'TheBloke/Mistral-7B-v0.1-AWQ' # default_model + assert processor.default_model == 'TheBloke/Mistral-7B-v0.1-AWQ' # default_model assert processor.base_url == 'http://vllm-service:8899/v1' # default_base_url assert processor.temperature == 0.0 # default_temperature assert processor.max_output == 2048 # default_max_output