fixed test_sendchat

This commit is contained in:
Paul Gauthier 2024-04-18 16:13:08 -07:00
parent 363b3202ab
commit 957026d3cc

View file

@ -12,12 +12,11 @@ class PrintCalled(Exception):
class TestSendChat(unittest.TestCase): class TestSendChat(unittest.TestCase):
@patch("litellm.completion")
@patch("builtins.print") @patch("builtins.print")
def test_send_with_retries_rate_limit_error(self, mock_print): def test_send_with_retries_rate_limit_error(self, mock_print, mock_completion):
mock_client = MagicMock()
# Set up the mock to raise # Set up the mock to raise
mock_client.chat.completions.create.side_effect = [ mock_completion.side_effect = [
openai.RateLimitError( openai.RateLimitError(
"rate limit exceeded", "rate limit exceeded",
response=MagicMock(), response=MagicMock(),
@ -27,20 +26,18 @@ class TestSendChat(unittest.TestCase):
] ]
# Call the send_with_retries method # Call the send_with_retries method
send_with_retries(mock_client, "model", ["message"], None, False) send_with_retries("model", ["message"], None, False)
mock_print.assert_called_once() mock_print.assert_called_once()
@patch("aider.sendchat.openai.ChatCompletion.create") @patch("litellm.completion")
@patch("builtins.print") @patch("builtins.print")
def test_send_with_retries_connection_error(self, mock_print, mock_chat_completion_create): def test_send_with_retries_connection_error(self, mock_print, mock_completion):
mock_client = MagicMock()
# Set up the mock to raise # Set up the mock to raise
mock_client.chat.completions.create.side_effect = [ mock_completion.side_effect = [
httpx.ConnectError("Connection error"), httpx.ConnectError("Connection error"),
None, None,
] ]
# Call the send_with_retries method # Call the send_with_retries method
send_with_retries(mock_client, "model", ["message"], None, False) send_with_retries("model", ["message"], None, False)
mock_print.assert_called_once() mock_print.assert_called_once()