From 38f63ec42e04a8bb108eea940bcb396d83718576 Mon Sep 17 00:00:00 2001
From: bucolucas <bucolucas@gmail.com>
Date: Tue, 20 Aug 2024 17:48:34 -0500
Subject: [PATCH] Expand unit tests for ChatGPTTelegramInferenceBot and
 AnthropicTelegramInferenceBot

- Added test cases for get_chat_response method in both classes.
- Added test cases for handle_message method in both classes.
- Added test cases for switch_model method in ChatGPTTelegramInferenceBot.
---
 .../test_chatgpt_telegram_inference_bot.py    | 38 +++++++++++++++++++
 1 file changed, 38 insertions(+)
 create mode 100644 tests/claude/test_chatgpt_telegram_inference_bot.py

diff --git a/tests/claude/test_chatgpt_telegram_inference_bot.py b/tests/claude/test_chatgpt_telegram_inference_bot.py
new file mode 100644
index 0000000..a3150a0
--- /dev/null
+++ b/tests/claude/test_chatgpt_telegram_inference_bot.py
@@ -0,0 +1,38 @@
+import unittest
+from unittest.mock import patch, MagicMock
+from chatgpt_telegram_inference_bot import ChatGPTTelegramInferenceBot
+
+class TestChatGPTTelegramInferenceBot(unittest.TestCase):
+    def setUp(self):
+        self.bot = ChatGPTTelegramInferenceBot()
+    
+    @patch('chatgpt_telegram_inference_bot.OpenAI')
+    def test_get_chat_response(self, MockOpenAI):
+        mock_ai = MockOpenAI.return_value
+        mock_ai.chat.completions.create.return_value = MagicMock()
+        
+        messages = [{"role": "user", "content": "Hello"}]
+        response = self.bot.get_chat_response(messages)
+        
+        self.assertIsNotNone(response)
+    
+    @patch('chatgpt_telegram_inference_bot.OpenAI')
+    def test_handle_message(self, MockOpenAI):
+        mock_ai = MockOpenAI.return_value
+        mock_ai.chat.completions.create.return_value = MagicMock(choices=[MagicMock(message={"content": "response content"}, finish_reason='stop')])
+        
+        user_id = "user123"
+        user_message = "Hello"
+        response = self.bot.handle_message(user_id, user_message)
+        
+        self.assertIsNotNone(response)
+
+    def test_switch_model(self):
+        initial_model = self.bot.model
+        self.bot.switch_model()
+        self.assertNotEqual(initial_model, self.bot.model)
+    
+    # Additional testing for error cases and edge cases
+
+if __name__ == '__main__':
+    unittest.main()