diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 41aedf4b5cd..18ff9938498 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -336,8 +336,6 @@ class ProxyInitializationHelpers: api_base = f"http://{host}:{port}" if isinstance(test, str): api_base = test - else: - raise ValueError("Invalid test value") client: Final = openai.OpenAI(api_key="My API Key", base_url=api_base) response: Final = client.chat.completions.create( diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 6347c165865..27755c23b51 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -197,23 +197,27 @@ class TestProxyInitializationHelpers: mock_client.chat.completions.create.side_effect = [ mock_response, mock_stream_response, + mock_response, + mock_stream_response, ] # Execute - with pytest.raises(ValueError, match="Invalid test value"): - ProxyInitializationHelpers._run_test_chat_completion( - "localhost", 8000, "gpt-3.5-turbo", True - ) + # The --test flag passes True; the helper must fall back to http://{host}:{port} + ProxyInitializationHelpers._run_test_chat_completion( + "localhost", 8000, "gpt-3.5-turbo", True + ) - # Test with valid string test value + # Test with explicit string test value ProxyInitializationHelpers._run_test_chat_completion( "localhost", 8000, "gpt-3.5-turbo", "http://test-url" ) # Assert - mock_openai.assert_called_once_with( - api_key="My API Key", base_url="http://test-url" + assert mock_openai.call_count == 2 + mock_openai.assert_any_call( + api_key="My API Key", base_url="http://localhost:8000" ) + mock_openai.assert_any_call(api_key="My API Key", base_url="http://test-url") mock_client.chat.completions.create.assert_called() def test_get_default_unvicorn_init_args(self):