From a56fa567ddf60dd984ae7a0ce8afe7718186248a Mon Sep 17 00:00:00 2001 From: "Paul Gauthier (aider)" Date: Sat, 7 Dec 2024 13:44:08 -0800 Subject: [PATCH] test: update print call count assertions in sendchat tests --- tests/basic/test_sendchat.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/basic/test_sendchat.py b/tests/basic/test_sendchat.py index 6b2684a40..efce33f51 100644 --- a/tests/basic/test_sendchat.py +++ b/tests/basic/test_sendchat.py @@ -37,8 +37,9 @@ class TestSendChat(unittest.TestCase): ] # Call the simple_send_with_retries method - simple_send_with_retries("model", ["message"]) - assert mock_print.call_count == 3 + simple_send_with_retries(self.mock_model, self.mock_messages) + # Should print: error message, description (if any), and retry message + assert mock_print.call_count == 2 # Error message + retry message @patch("litellm.completion") def test_send_completion_basic(self, mock_completion): @@ -90,4 +91,5 @@ class TestSendChat(unittest.TestCase): result = simple_send_with_retries(self.mock_model, self.mock_messages) assert result is None + # Should only print the error message assert mock_print.call_count == 1