Test: Add unittests for 'Chat_GPT_Function.py'

Author: Alby084

projects/Chat-GPT-Discord-Bot/tests.py

@@ -0,0 +1,107 @@
+import unittest
+from unittest.mock import patch, MagicMock
+from Chat_GPT_Function import gpt, dalle3, dalle2
+from dotenv import load_dotenv
+import os
+
+load_dotenv(override=True)
+
+gpt_api_key = os.getenv("GPT_API_KEY")
+class TestGPT(unittest.TestCase):
+    @patch('Chat_GPT_Function.OpenAI')
+    def test_gpt_success(self, mock_openai):
+        mock_client = MagicMock()
+        mock_openai.return_value = mock_client
+        mock_response = MagicMock()
+        mock_response.choices = [MagicMock(message=MagicMock(content="Test output"))]
+        mock_client.chat.completions.create.return_value = mock_response
+
+        result = gpt("gpt-3.5-turbo", "prompt", "sys_prompt", 0.5)
+
+        self.assertIsInstance(result, str)
+        self.assertTrue(result.strip())
+        mock_openai.assert_called_once_with(api_key=gpt_api_key)
+        mock_client.chat.completions.create.assert_called_once_with(
+            model="gpt-3.5-turbo",
+            messages=[
+                {"role": "system", "content": "sys_prompt"},
+                {"role": "user", "content": "prompt"}
+            ],
+            temperature=0.5,
+            top_p=1
+        )
+
+    @patch('Chat_GPT_Function.OpenAI')
+    def test_gpt_failure(self, mock_openai):
+        mock_client = MagicMock()
+        mock_openai.return_value = mock_client
+        mock_client.chat.completions.create.side_effect = Exception("API error")
+
+        with self.assertRaises(Exception):
+            gpt("model", "prompt", "sys_prompt", 0.5)
+
+class TestDALLE3(unittest.TestCase):
+    @patch('Chat_GPT_Function.OpenAI')
+    def test_dalle3_success(self, mock_openai):
+        mock_client = MagicMock()
+        mock_openai.return_value = mock_client
+        mock_response = MagicMock()
+        mock_response.data = [MagicMock(url="https://example.com/image.png")]
+        mock_client.images.generate.return_value = mock_response
+
+        result = dalle3("prompt", "hd", "1792x1024", "vivid")
+
+        self.assertIsInstance(result, str)
+        self.assertTrue(result.startswith("https://"))
+        mock_openai.assert_called_once_with(api_key=gpt_api_key)
+        mock_client.images.generate.assert_called_once_with(
+            model="dall-e-3",
+            prompt="prompt",
+            size="1792x1024",
+            quality="hd",
+            style="vivid",
+            n=1
+        )
+
+
+    @patch('Chat_GPT_Function.OpenAI')
+    def test_dalle3_failure(self, mock_openai):
+        mock_client = MagicMock()
+        mock_openai.return_value = mock_client
+        mock_client.images.generate.side_effect = Exception("API error")
+
+        with self.assertRaises(Exception):
+            dalle3("prompt", "quality", "size", "style")
+
+class TestDALLE2(unittest.TestCase):
+    @patch('Chat_GPT_Function.OpenAI')
+    def test_dalle2_success(self, mock_openai):
+        mock_client = MagicMock()
+        mock_openai.return_value = mock_client
+        mock_response = MagicMock()
+        mock_response.data = [MagicMock(url="https://example.com/image.png")]
+        mock_client.images.generate.return_value = mock_response
+
+        result = dalle2("prompt", "256x256")
+
+        self.assertIsInstance(result, str)
+        self.assertTrue(result.startswith("https://"))
+        mock_openai.assert_called_once_with(api_key=gpt_api_key)
+        mock_client.images.generate.assert_called_once_with(
+            model="dall-e-2",
+            prompt="prompt",
+            size="256x256",
+            n=1
+        )
+
+    @patch('Chat_GPT_Function.OpenAI')
+    def test_dalle2_failure(self, mock_openai):
+        mock_client = MagicMock()
+        mock_openai.return_value = mock_client
+        mock_client.images.generate.side_effect = Exception("API error")
+
+        with self.assertRaises(Exception):
+            dalle2("prompt", "size")
+
+if __name__ == '__main__':
+    unittest.main()