From 5263367af0c17b18bbab1dff29c18e9283720f0c Mon Sep 17 00:00:00 2001 From: Paul Gauthier Date: Sat, 3 Jun 2023 06:23:07 -0700 Subject: [PATCH] added import --- tests/test_coder.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/test_coder.py b/tests/test_coder.py index e50ac756b..eee43897e 100644 --- a/tests/test_coder.py +++ b/tests/test_coder.py @@ -3,6 +3,7 @@ import unittest from unittest.mock import MagicMock, patch import openai +import requests from aider.coder import Coder @@ -141,7 +142,6 @@ class TestCoder(unittest.TestCase): # Assert that time.sleep was called once mock_sleep.assert_called_once() - @patch("aider.coder.openai.ChatCompletion.create") @patch("aider.coder.time.sleep") def test_send_with_retries_connection_error(self, mock_sleep, mock_chat_completion_create): @@ -151,8 +151,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 ConnectionError on the first call and return None on the second call - mock_chat_completion_create.side_effect = [requests.exceptions.ConnectionError("Connection error"), None] + # Set up the mock to raise ConnectionError on the first call + # and return None on the second call + mock_chat_completion_create.side_effect = [ + requests.exceptions.ConnectionError("Connection error"), + None, + ] # Call the send_with_retries method coder.send_with_retries("model", ["message"]) @@ -160,5 +164,6 @@ class TestCoder(unittest.TestCase): # Assert that time.sleep was called once mock_sleep.assert_called_once() + if __name__ == "__main__": unittest.main()