From 2a8e5a727c20d165e2066717a2879ea1894fcb07 Mon Sep 17 00:00:00 2001
From: Paul Gauthier <aider@paulg.org>
Date: Sat, 3 Jun 2023 06:20:56 -0700
Subject: [PATCH] aider: Added a test case for `coder.send_with_retries(model,
 messages)` to handle `RateLimitError` and call `time.sleep` once.

# Context:
USER: add a test case for coder.send_with_retries(model, messages). it should patch openai.ChatCompletion.create so that on the first call it raises openai.error import RateLimitError. and on the 2nd call it returns None. the test should confirm that time.sleep is called 1 time.
ASSISTANT: To add the test case for `coder.send_with_retries(model, messages)`, we will modify the `tests/test_coder.py` file. We will create a new test method called `test_send_with_retries_rate_limit_error` and use `patch` to mock the `openai.ChatCompletion.create` and `time.sleep` methods.

Here's the code change for the new test case:

```python
tests/test_coder.py
<<<<<<< ORIGINAL
if __name__ == "__main__":
    unittest.main()
=======
    @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):
        # Mock the IO object
        mock_io = MagicMock()

        # 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]

        # Call the send_with_retries method
        coder.send_with_retries("model", ["message"])

        # Assert that time.sleep was called once
        mock_sleep.assert_called_once()

if __name__ == "__main__":
    unittest.main()
>>>>>>> UPDATED
```

This test case will ensure that `send_with_retries` handles the `RateLimitError` correctly and calls `time.sleep` once before retrying.
---
 tests/test_coder.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/tests/test_coder.py b/tests/test_coder.py
index f703969e3..037049022 100644
--- a/tests/test_coder.py
+++ b/tests/test_coder.py
@@ -118,5 +118,23 @@ class TestCoder(unittest.TestCase):
         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):
+        # Mock the IO object
+        mock_io = MagicMock()
+
+        # 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]
+
+        # Call the send_with_retries method
+        coder.send_with_retries("model", ["message"])
+
+        # Assert that time.sleep was called once
+        mock_sleep.assert_called_once()
+
 if __name__ == "__main__":
     unittest.main()