feat: Add tests for sanity_check_model function

This commit is contained in:
Paul Gauthier (aider) 2024-08-26 14:42:24 -07:00
parent 2e573af1d8
commit abf19a61cf

View file

@ -1,6 +1,8 @@
import unittest import unittest
from unittest.mock import patch, MagicMock
import os
from aider.models import Model, get_model_info from aider.models import Model, get_model_info, sanity_check_model
class TestModels(unittest.TestCase): class TestModels(unittest.TestCase):
@ -27,6 +29,40 @@ class TestModels(unittest.TestCase):
model = Model("gpt-4-0613") model = Model("gpt-4-0613")
self.assertEqual(model.info["max_input_tokens"], 8 * 1024) self.assertEqual(model.info["max_input_tokens"], 8 * 1024)
@patch('os.environ')
def test_sanity_check_model_all_set(self, mock_environ):
mock_environ.get.return_value = "dummy_value"
mock_io = MagicMock()
model = MagicMock()
model.name = "test-model"
model.missing_keys = ["API_KEY1", "API_KEY2"]
model.keys_in_environment = True
model.info = {"some": "info"}
sanity_check_model(mock_io, model)
mock_io.tool_error.assert_called_once_with("Model test-model: Environment variables status:")
calls = mock_io.tool_error.call_args_list
self.assertIn("- API_KEY1: ✓ Set", str(calls))
self.assertIn("- API_KEY2: ✓ Set", str(calls))
@patch('os.environ')
def test_sanity_check_model_not_set(self, mock_environ):
mock_environ.get.return_value = ""
mock_io = MagicMock()
model = MagicMock()
model.name = "test-model"
model.missing_keys = ["API_KEY1", "API_KEY2"]
model.keys_in_environment = True
model.info = {"some": "info"}
sanity_check_model(mock_io, model)
mock_io.tool_error.assert_called_once_with("Model test-model: Environment variables status:")
calls = mock_io.tool_error.call_args_list
self.assertIn("- API_KEY1: ✗ Not set", str(calls))
self.assertIn("- API_KEY2: ✗ Not set", str(calls))
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()