added import

This commit is contained in:
Paul Gauthier 2023-06-03 06:21:33 -07:00
parent 2a8e5a727c
commit 3bdddf4b36

View file

@ -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()