From dd3d97ace69b21d1443fec77d0b703225a22cb5e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 9 Aug 2024 17:02:29 -0700 Subject: [PATCH] test(test_optional_params.py): use num_retries instead of 'max_retries' if given --- litellm/tests/test_optional_params.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/litellm/tests/test_optional_params.py b/litellm/tests/test_optional_params.py index ad640fe295..b1ee80b59e 100644 --- a/litellm/tests/test_optional_params.py +++ b/litellm/tests/test_optional_params.py @@ -446,3 +446,20 @@ def test_bedrock_optional_params_embeddings_provider_specific_params(): wait_for_model=True, ) assert len(optional_params) == 1 + + +def test_get_optional_params_num_retries(): + """ + Relevant issue - https://github.com/BerriAI/litellm/issues/5124 + """ + with patch("litellm.main.get_optional_params", new=MagicMock()) as mock_client: + _ = litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello world"}], + num_retries=10, + ) + + mock_client.assert_called() + + print(f"mock_client.call_args: {mock_client.call_args}") + assert mock_client.call_args.kwargs["max_retries"] == 10