Refactor test_chat to use pytest instead of unittest (#3484)
* refactor_for_pytest * formatting --------- Co-authored-by: James Collins <collijk@uw.edu>pull/3436/head^2
parent
c1f1da27e7
commit
ee4043ae19
|
@ -1,29 +1,26 @@
|
||||||
# Generated by CodiumAI
|
# Generated by CodiumAI
|
||||||
import time
|
import time
|
||||||
import unittest
|
|
||||||
from unittest.mock import patch
|
from unittest.mock import patch
|
||||||
|
|
||||||
from autogpt.chat import create_chat_message, generate_context
|
from autogpt.chat import create_chat_message, generate_context
|
||||||
|
|
||||||
|
|
||||||
class TestChat(unittest.TestCase):
|
def test_happy_path_role_content():
|
||||||
"""Test the chat module functions."""
|
|
||||||
|
|
||||||
def test_happy_path_role_content(self):
|
|
||||||
"""Test that the function returns a dictionary with the correct keys and values when valid strings are provided for role and content."""
|
"""Test that the function returns a dictionary with the correct keys and values when valid strings are provided for role and content."""
|
||||||
result = create_chat_message("system", "Hello, world!")
|
result = create_chat_message("system", "Hello, world!")
|
||||||
self.assertEqual(result, {"role": "system", "content": "Hello, world!"})
|
assert result == {"role": "system", "content": "Hello, world!"}
|
||||||
|
|
||||||
def test_empty_role_content(self):
|
|
||||||
|
def test_empty_role_content():
|
||||||
"""Test that the function returns a dictionary with the correct keys and values when empty strings are provided for role and content."""
|
"""Test that the function returns a dictionary with the correct keys and values when empty strings are provided for role and content."""
|
||||||
result = create_chat_message("", "")
|
result = create_chat_message("", "")
|
||||||
self.assertEqual(result, {"role": "", "content": ""})
|
assert result == {"role": "", "content": ""}
|
||||||
|
|
||||||
@patch("time.strftime")
|
|
||||||
def test_generate_context_empty_inputs(self, mock_strftime):
|
def test_generate_context_empty_inputs(mocker):
|
||||||
"""Test the behavior of the generate_context function when all input parameters are empty."""
|
"""Test the behavior of the generate_context function when all input parameters are empty."""
|
||||||
# Mock the time.strftime function to return a fixed value
|
# Mock the time.strftime function to return a fixed value
|
||||||
mock_strftime.return_value = "Sat Apr 15 00:00:00 2023"
|
mocker.patch("time.strftime", return_value="Sat Apr 15 00:00:00 2023")
|
||||||
# Arrange
|
# Arrange
|
||||||
prompt = ""
|
prompt = ""
|
||||||
relevant_memory = ""
|
relevant_memory = ""
|
||||||
|
@ -50,9 +47,10 @@ class TestChat(unittest.TestCase):
|
||||||
},
|
},
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
self.assertEqual(result, expected_result)
|
assert result == expected_result
|
||||||
|
|
||||||
def test_generate_context_valid_inputs(self):
|
|
||||||
|
def test_generate_context_valid_inputs():
|
||||||
"""Test that the function successfully generates a current_context given valid inputs."""
|
"""Test that the function successfully generates a current_context given valid inputs."""
|
||||||
# Given
|
# Given
|
||||||
prompt = "What is your favorite color?"
|
prompt = "What is your favorite color?"
|
||||||
|
@ -73,16 +71,12 @@ class TestChat(unittest.TestCase):
|
||||||
result = generate_context(prompt, relevant_memory, full_message_history, model)
|
result = generate_context(prompt, relevant_memory, full_message_history, model)
|
||||||
|
|
||||||
# Then
|
# Then
|
||||||
self.assertIsInstance(result[0], int)
|
assert isinstance(result[0], int)
|
||||||
self.assertIsInstance(result[1], int)
|
assert isinstance(result[1], int)
|
||||||
self.assertIsInstance(result[2], int)
|
assert isinstance(result[2], int)
|
||||||
self.assertIsInstance(result[3], list)
|
assert isinstance(result[3], list)
|
||||||
self.assertGreaterEqual(result[0], 0)
|
assert result[0] >= 0
|
||||||
self.assertGreaterEqual(result[1], 0)
|
assert result[2] >= 0
|
||||||
self.assertGreaterEqual(result[2], 0)
|
assert result[1] >= 0
|
||||||
self.assertGreaterEqual(
|
assert len(result[3]) >= 3 # current_context should have at least 3 messages
|
||||||
len(result[3]), 3
|
assert result[1] <= 2048 # token limit for GPT-3.5-turbo-0301 is 2048 tokens
|
||||||
) # current_context should have at least 3 messages
|
|
||||||
self.assertLessEqual(
|
|
||||||
result[1], 2048
|
|
||||||
) # token limit for GPT-3.5-turbo-0301 is 2048 tokens
|
|
||||||
|
|
Loading…
Reference in New Issue