fix(test_router.py): catch timeouts

This commit is contained in:
Krrish Dholakia 2023-11-17 15:55:59 -08:00
parent bb9e7c65e9
commit a99efc544e

View file

@ -246,13 +246,15 @@ def test_acompletion_on_router():
] ]
async def get_response(): async def get_response():
router = Router(model_list=model_list, redis_host=os.environ["REDIS_HOST"], redis_password=os.environ["REDIS_PASSWORD"], redis_port=os.environ["REDIS_PORT"], cache_responses=True, timeout=0.1) router = Router(model_list=model_list, redis_host=os.environ["REDIS_HOST"], redis_password=os.environ["REDIS_PASSWORD"], redis_port=os.environ["REDIS_PORT"], cache_responses=True, timeout=10)
response1 = await router.acompletion(model="gpt-3.5-turbo", messages=messages) response1 = await router.acompletion(model="gpt-3.5-turbo", messages=messages)
print(f"response1: {response1}") print(f"response1: {response1}")
response2 = await router.acompletion(model="gpt-3.5-turbo", messages=messages) response2 = await router.acompletion(model="gpt-3.5-turbo", messages=messages)
print(f"response2: {response2}") print(f"response2: {response2}")
assert response1["choices"][0]["message"]["content"] == response2["choices"][0]["message"]["content"] assert response1["choices"][0]["message"]["content"] == response2["choices"][0]["message"]["content"]
asyncio.run(get_response()) asyncio.run(get_response())
except litellm.Timeout as e:
pass
except Exception as e: except Exception as e:
traceback.print_exc() traceback.print_exc()
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")