From 42ae279b91544ebcaafd4cc595d35ee2a9f460eb Mon Sep 17 00:00:00 2001 From: Paul Gauthier Date: Sat, 7 Dec 2024 15:09:04 -0800 Subject: [PATCH] fix: update test cases to use Model wrapper class --- tests/basic/test_sendchat.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/basic/test_sendchat.py b/tests/basic/test_sendchat.py index aa2b4a254..fb2b6fb5a 100644 --- a/tests/basic/test_sendchat.py +++ b/tests/basic/test_sendchat.py @@ -75,7 +75,7 @@ class TestSendChat(unittest.TestCase): mock_completion.return_value.choices = None # Should return None on AttributeError - result = simple_send_with_retries(self.mock_model, self.mock_messages) + result = simple_send_with_retries(Model(self.mock_model), self.mock_messages) assert result is None @patch("litellm.completion") @@ -89,7 +89,7 @@ class TestSendChat(unittest.TestCase): message="Invalid request", llm_provider="test_provider", model="test_model" ) - result = simple_send_with_retries(self.mock_model, self.mock_messages) + result = simple_send_with_retries(Model(self.mock_model), self.mock_messages) assert result is None # Should only print the error message assert mock_print.call_count == 1