diff --git a/tests/test_coder.py b/tests/test_coder.py index 037049022..61b60cc79 100644 --- a/tests/test_coder.py +++ b/tests/test_coder.py @@ -2,6 +2,8 @@ import os import unittest from unittest.mock import MagicMock, patch +import openai + from aider.coder import Coder @@ -117,7 +119,6 @@ class TestCoder(unittest.TestCase): # Assert that the returned message is the expected one self.assertEqual(result, 'a good "commit message"') - @patch("aider.coder.openai.ChatCompletion.create") @patch("aider.coder.time.sleep") def test_send_with_retries_rate_limit_error(self, mock_sleep, mock_chat_completion_create): @@ -127,8 +128,12 @@ class TestCoder(unittest.TestCase): # Initialize the Coder object with the mocked IO and mocked repo coder = Coder(io=mock_io, openai_api_key="fake_key") - # Set up the mock to raise RateLimitError on the first call and return None on the second call - mock_chat_completion_create.side_effect = [openai.error.RateLimitError("Rate limit exceeded"), None] + # Set up the mock to raise RateLimitError on + # the first call and return None on the second call + mock_chat_completion_create.side_effect = [ + openai.error.RateLimitError("Rate limit exceeded"), + None, + ] # Call the send_with_retries method coder.send_with_retries("model", ["message"]) @@ -136,5 +141,6 @@ class TestCoder(unittest.TestCase): # Assert that time.sleep was called once mock_sleep.assert_called_once() + if __name__ == "__main__": unittest.main()