From 109f197f527c15c15f4c7129444a1853c7c1a34d Mon Sep 17 00:00:00 2001
From: Paul Gauthier <aider@paulg.org>
Date: Thu, 8 Aug 2024 15:22:58 -0300
Subject: [PATCH] feat: Add tests for simple_send_with_retries function

---
 tests/basic/test_sendchat.py | 47 ++++++++++++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)
 create mode 100644 tests/basic/test_sendchat.py

diff --git a/tests/basic/test_sendchat.py b/tests/basic/test_sendchat.py
new file mode 100644
index 000000000..f77e687ef
--- /dev/null
+++ b/tests/basic/test_sendchat.py
@@ -0,0 +1,47 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+import httpx
+
+from aider.llm import litellm
+from aider.sendchat import simple_send_with_retries
+
+
+class PrintCalled(Exception):
+    pass
+
+
+class TestSendChat(unittest.TestCase):
+    @patch("litellm.completion")
+    @patch("builtins.print")
+    def test_simple_send_with_retries_rate_limit_error(self, mock_print, mock_completion):
+        mock = MagicMock()
+        mock.status_code = 500
+
+        # Set up the mock to raise
+        mock_completion.side_effect = [
+            litellm.exceptions.RateLimitError(
+                "rate limit exceeded",
+                response=mock,
+                llm_provider="llm_provider",
+                model="model",
+            ),
+            None,
+        ]
+
+        # Call the simple_send_with_retries method
+        simple_send_with_retries("model", ["message"])
+        mock_print.assert_called_once()
+
+    @patch("litellm.completion")
+    @patch("builtins.print")
+    def test_simple_send_with_retries_connection_error(self, mock_print, mock_completion):
+        # Set up the mock to raise
+        mock_completion.side_effect = [
+            httpx.ConnectError("Connection error"),
+            None,
+        ]
+
+        # Call the simple_send_with_retries method
+        simple_send_with_retries("model", ["message"])
+        mock_print.assert_called_once()