Merge pull request #1700 from OmriGM/tests/chat-tests

Added tests for `create_chat_message` and `generate_context` methods of the chat module

tests/unit/test_chat.py

@@ -0,0 +1,70 @@
+
+# Generated by CodiumAI
+import unittest
+import time
+from unittest.mock import patch
+
+from autogpt.chat import create_chat_message, generate_context
+
+
+class TestChat(unittest.TestCase):
+
+ # Tests that the function returns a dictionary with the correct keys and values when valid strings are provided for role and content.
+ def test_happy_path_role_content(self):
+ result = create_chat_message("system", "Hello, world!")
+ self.assertEqual(result, {"role": "system", "content": "Hello, world!"})
+
+ # Tests that the function returns a dictionary with the correct keys and values when empty strings are provided for role and content.
+ def test_empty_role_content(self):
+ result = create_chat_message("", "")
+ self.assertEqual(result, {"role": "", "content": ""})
+
+ # Tests the behavior of the generate_context function when all input parameters are empty.
+ @patch("time.strftime")
+ def test_generate_context_empty_inputs(self, mock_strftime):
+ # Mock the time.strftime function to return a fixed value
+ mock_strftime.return_value = "Sat Apr 15 00:00:00 2023"
+ # Arrange
+ prompt = ""
+ relevant_memory = ""
+ full_message_history = []
+ model = "gpt-3.5-turbo-0301"
+
+ # Act
+ result = generate_context(prompt, relevant_memory, full_message_history, model)
+
+ # Assert
+ expected_result = (-1, 47, 3, [
+ {"role": "system", "content": ""},
+ {"role": "system", "content": f"The current time and date is {time.strftime('%c')}"},
+ {"role": "system", "content": f"This reminds you of these events from your past:\n\n\n"},
+ ])
+ self.assertEqual(result, expected_result)
+
+ # Tests that the function successfully generates a current_context given valid inputs.
+ def test_generate_context_valid_inputs(self):
+ # Given
+ prompt = "What is your favorite color?"
+ relevant_memory = "You once painted your room blue."
+ full_message_history = [
+ create_chat_message("user", "Hi there!"),
+ create_chat_message("assistant", "Hello! How can I assist you today?"),
+ create_chat_message("user", "Can you tell me a joke?"),
+ create_chat_message("assistant", "Why did the tomato turn red? Because it saw the salad dressing!"),
+ create_chat_message("user", "Haha, that's funny."),
+ ]
+ model = "gpt-3.5-turbo-0301"
+
+ # When
+ result = generate_context(prompt, relevant_memory, full_message_history, model)
+
+ # Then
+ self.assertIsInstance(result[0], int)
+ self.assertIsInstance(result[1], int)
+ self.assertIsInstance(result[2], int)
+ self.assertIsInstance(result[3], list)
+ self.assertGreaterEqual(result[0], 0)
+ self.assertGreaterEqual(result[1], 0)
+ self.assertGreaterEqual(result[2], 0)
+ self.assertGreaterEqual(len(result[3]), 3) # current_context should have at least 3 messages
+ self.assertLessEqual(result[1], 2048) # token limit for GPT-3.5-turbo-0301 is 2048 tokens