engineering/opencompletion.com
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix test failures in guarded_ai test files

Browse source 
- Update test_initialize_model_map to mock models.list() response properly
- Update test_get_openai_client_and_model_default to match new MODEL_X behavior
- Fix test_initialize_model_map_with_env_vars in functional tests

Tests now properly mock the OpenAI client's models.list() response, which
returns model IDs that are used as keys in MODEL_CLIENT_MAP, not endpoint names.
This commit is contained in:
Claude 2025-11-10 19:50:06 +00:00
parent 39e39e80f1
commit ad47efd31d
No known key found for this signature in database
2 changed files with 44 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

22
tests/functional/test_guarded_ai.py
View file

@ -440,10 +440,28 @@ class TestGuardedAIClientAndErrorHandling(unittest.TestCase):
mock_client2 = MagicMock()
mock_get_client.side_effect = [mock_client1, mock_client2]
# Mock the models.list() response for both clients
mock_model1 = MagicMock()
mock_model1.id = "test-model-1"
mock_client1.models.list.return_value.data = [mock_model1]
mock_model2 = MagicMock()
mock_model2.id = "test-model-2"
mock_client2.models.list.return_value.data = [mock_model2]
guarded_ai.initialize_model_map()
self.assertIn("endpoint_1", guarded_ai.MODEL_CLIENT_MAP)
self.assertIn("endpoint_2", guarded_ai.MODEL_CLIENT_MAP)
# Check that models were added to the map
self.assertIn("test-model-1", guarded_ai.MODEL_CLIENT_MAP)
self.assertIn("test-model-2", guarded_ai.MODEL_CLIENT_MAP)
self.assertEqual(
guarded_ai.MODEL_CLIENT_MAP["test-model-1"][1],
"https://api.test1.com",
)
self.assertEqual(
guarded_ai.MODEL_CLIENT_MAP["test-model-2"][1],
"https://api.test2.com",
)
def test_initialize_model_map_with_errors(self):
"""Test error handling in model map initialization"""

28
tests/unit/test_guarded_ai_functions.py
View file

@ -297,17 +297,32 @@ class TestGuardedAI(unittest.TestCase):
mock_client = MagicMock()
mock_get_client.return_value = mock_client
# Mock the models.list() response
mock_model = MagicMock()
mock_model.id = "test-model-id"
mock_client.models.list.return_value.data = [mock_model]
# Clear and reinitialize
import guarded_ai
guarded_ai.MODEL_CLIENT_MAP = {}
initialize_model_map()
# Verify client was created and stored
# Verify client was created and stored with actual model ID
mock_get_client.assert_called_with("http://test.com", "test-key")
self.assertIn("endpoint_0", guarded_ai.MODEL_CLIENT_MAP)
self.assertEqual(guarded_ai.MODEL_CLIENT_MAP["endpoint_0"][0], mock_client)
self.assertIn("test-model-id", guarded_ai.MODEL_CLIENT_MAP)
self.assertEqual(guarded_ai.MODEL_CLIENT_MAP["test-model-id"][0], mock_client)
self.assertEqual(
guarded_ai.MODEL_CLIENT_MAP["test-model-id"][1], "http://test.com"
)
@patch.dict(
"os.environ",
{
"MODEL_ENDPOINT_1": "http://hermes.test",
"MODEL_API_KEY_1": "hermes-key",
},
)
def test_get_openai_client_and_model_default(self):
"""Test getting OpenAI client with default model"""
with patch("guarded_ai.MODEL_CLIENT_MAP", {}):
@ -315,9 +330,14 @@ class TestGuardedAI(unittest.TestCase):
mock_client = MagicMock()
mock_get_client.return_value = mock_client
# Mock the models.list() response for MODEL_1
mock_model = MagicMock()
mock_model.id = "adamo1139/Hermes-3-Llama-3.1-8B-FP8-Dynamic"
mock_client.models.list.return_value.data = [mock_model]
client, model = get_openai_client_and_model()
# Should return default model name
# Should return MODEL_1's first model
self.assertEqual(model, "adamo1139/Hermes-3-Llama-3.1-8B-FP8-Dynamic")
self.assertEqual(client, mock_client)