From 990c32a5d65d5bc421ae7725d50d76ca2e87faee Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 10 Jan 2024 15:49:32 +0530 Subject: [PATCH] test(test_router.py): handle image gen timeouts --- litellm/tests/test_router.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_router.py b/litellm/tests/test_router.py index a1349d4fa7..f0f1594a83 100644 --- a/litellm/tests/test_router.py +++ b/litellm/tests/test_router.py @@ -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,8 +535,13 @@ def test_aembedding_on_router(): ) router.reset() except Exception as e: - traceback.print_exc() - pytest.fail(f"Error occurred: {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}") # test_aembedding_on_router()