Skip to content

Commit

Permalink
fix prompt parser test
Browse files Browse the repository at this point in the history
  • Loading branch information
saum7800 committed Apr 15, 2024
1 parent 50d20de commit 4406890
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions tests/prompt_parser_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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()


Expand Down

0 comments on commit 4406890

Please sign in to comment.