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
Media 2023-04-28 20:27:52 +02:00 committed by GitHub
parent c1f1da27e7
commit ee4043ae19
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 67 additions and 73 deletions

View File

@ -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