From 44068908a0d9fe80147cc7eb6dd3e1bb80acead4 Mon Sep 17 00:00:00 2001 From: saum7800 Date: Mon, 15 Apr 2024 17:53:11 -0400 Subject: [PATCH] fix prompt parser test --- tests/prompt_parser_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/prompt_parser_test.py b/tests/prompt_parser_test.py index b67a61119..f33321a69 100644 --- a/tests/prompt_parser_test.py +++ b/tests/prompt_parser_test.py @@ -136,7 +136,7 @@ def test_instruction_parser_with_invalid_json(mocked_parsing_method): @patch("time.sleep") @patch( "prompt2model.utils.APIAgent.generate_one_completion", - side_effect=openai.Timeout("timeout"), + side_effect=openai.APITimeoutError("timeout"), ) def test_instruction_parser_with_timeout(mocked_parsing_method, mocked_sleep_method): """Verify that we wait and retry (a set number of times) if the API times out. @@ -165,7 +165,7 @@ def test_instruction_parser_with_timeout(mocked_parsing_method, mocked_sleep_met assert isinstance(exc_info.value, RuntimeError) # Check if the original exception (e) is present as the cause original_exception = exc_info.value.__cause__ - assert isinstance(original_exception, openai.Timeout) + assert isinstance(original_exception, openai.APITimeoutError) gc.collect()