refactor: Update calls to simple_send_with_retries to use model method

This commit is contained in:
Paul Gauthier (aider) 2025-02-04 12:04:25 -08:00
parent cfe9c86edd
commit 535b3ce286

View file

@ -4,7 +4,6 @@ from unittest.mock import MagicMock, patch
from aider.exceptions import LiteLLMExceptions from aider.exceptions import LiteLLMExceptions
from aider.llm import litellm from aider.llm import litellm
from aider.models import Model from aider.models import Model
from aider.sendchat import simple_send_with_retries
class PrintCalled(Exception): class PrintCalled(Exception):
@ -38,7 +37,7 @@ class TestSendChat(unittest.TestCase):
] ]
# Call the simple_send_with_retries method # Call the simple_send_with_retries method
simple_send_with_retries(Model(self.mock_model), self.mock_messages) Model(self.mock_model).simple_send_with_retries(self.mock_messages)
assert mock_print.call_count == 3 assert mock_print.call_count == 3
@patch("litellm.completion") @patch("litellm.completion")
@ -75,7 +74,7 @@ class TestSendChat(unittest.TestCase):
mock_completion.return_value.choices = None mock_completion.return_value.choices = None
# Should return None on AttributeError # Should return None on AttributeError
result = simple_send_with_retries(Model(self.mock_model), self.mock_messages) result = Model(self.mock_model).simple_send_with_retries(self.mock_messages)
assert result is None assert result is None
@patch("litellm.completion") @patch("litellm.completion")
@ -89,7 +88,7 @@ class TestSendChat(unittest.TestCase):
message="Invalid request", llm_provider="test_provider", model="test_model" message="Invalid request", llm_provider="test_provider", model="test_model"
) )
result = simple_send_with_retries(Model(self.mock_model), self.mock_messages) result = Model(self.mock_model).simple_send_with_retries(self.mock_messages)
assert result is None assert result is None
# Should only print the error message # Should only print the error message
assert mock_print.call_count == 1 assert mock_print.call_count == 1