phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

simplify test

Browse source
This commit is contained in:
swiftdevil 2025-03-14 11:42:26 -07:00
parent 0c98a05c8f
commit b6b1aff76f
1 changed files with 8 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

20
tests/local_testing/test_completion_with_retries.py
View file

@ -155,22 +155,18 @@ async def test_acompletion_with_retries_retries():
from unittest.mock import patch from unittest.mock import patch
retry_count = 0
with patch.object(litellm, "acompletion") as mock_completion: with patch.object(litellm, "acompletion") as mock_completion:
async def timeout_fnc(*args, **kwargs): async def timeout_fnc(*args, **kwargs):
nonlocal retry_count
retry_count += 1
await mock_completion(*args, **kwargs) await mock_completion(*args, **kwargs)
if retry_count < 3: raise litellm.Timeout(message='test', model='gpt-3.5-turbo', llm_provider='mock')
raise litellm.Timeout(message='test', model='gpt-3.5-turbo', llm_provider='mock')
await acompletion_with_retries( with pytest.raises(litellm.Timeout):
model="gpt-3.5-turbo", await acompletion_with_retries(
messages=[{"gm": "vibe", "role": "user"}], model="gpt-3.5-turbo",
num_retries=3, messages=[{"gm": "vibe", "role": "user"}],
original_function=timeout_fnc, num_retries=3,
) original_function=timeout_fnc,
)
assert mock_completion.call_count == 3 assert mock_completion.call_count == 3