test(test_router.py): handle image gen timeouts

This commit is contained in:
Krrish Dholakia 2024-01-10 15:49:32 +05:30
parent fe632c08a4
commit 990c32a5d6

View file

@ -458,6 +458,8 @@ async def test_aimg_gen_on_router():
except Exception as e:
if "Your task failed as a result of our safety system." in str(e):
pass
elif "Operation polling timed out" in str(e):
pass
else:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
@ -533,6 +535,11 @@ def test_aembedding_on_router():
)
router.reset()
except Exception as e:
if "Your task failed as a result of our safety system." in str(e):
pass
elif "Operation polling timed out" in str(e):
pass
else:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")