diff --git a/test_ai_function.py b/test_ai_function.py index bb2e19d..c8c48c7 100644 --- a/test_ai_function.py +++ b/test_ai_function.py @@ -11,15 +11,16 @@ # Run all tests, print the results, and return the number of failed tests def run_tests(model): - test_functions = [test_1, test_2, test_3, test_4, test_5, test_6] + test_functions = [test_1, test_2, test_3, test_4, test_5, test_6, test_7] test_names = [ "Generate fake people", "Generate Random Password", "Calculate area of triangle", "Calculate the nth prime number", "Encrypt text", - "Find missing numbers" -] + "Find missing numbers", + "Substitute prime numbers" + ] failed_tests = [] i = 0 @@ -167,5 +168,30 @@ def test_6(model): print("Testing if the result list contains the expected missing numbers...") assert result_list == expected_missing_numbers +def test_7(model): + function_string = "def substitute_prime_numbers_in_list(numbers: list[int]) -> list[int]:" + args = ["[5, 8, 10, 3, 29]"] + description_string = """Substitutes the prime numbers in a list with the closest greater composite and returns the list.""" + + result_string = ai_functions.ai_function(function_string, args, description_string, model) + + print(f"Output: {result_string}") + + # Assert the result can be parsed as a list + try: + result_list = ast.literal_eval(result_string) + print("Testing if result is a a list...") + assert isinstance(result_list, list) + except Exception as e: + print(e) + assert False + + # Assert the result list contains the expected composite + expected_numbers = [6, 8, 10, 4, 30] + print("Testing if the result list is the same length as the input list...") + assert len(result_list) == len(expected_numbers) + print("Testing if the result list contains the expected composite numbers...") + assert result_list == expected_numbers + run_tests("gpt-4") run_tests("gpt-3.5-turbo") \ No newline at end of file