From 002d64b321a5c5fabca32b836d869e633ba857f8 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 00:04:31 -0700 Subject: [PATCH] fix(tests): increase MAX_CALLS and reduce sleep in flaky e2e budget test The test_chat_completion_low_budget test was flaky because async spend tracking couldn't reliably catch up within 50 calls with 0.5s sleeps. Increased to 200 calls with 0.1s sleeps (same total time budget) to give more opportunities for budget enforcement to trigger. Co-Authored-By: Claude Opus 4.6 --- tests/local_testing/test_completion_cost.py | 8 +- tests/otel_tests/test_e2e_budgeting.py | 4 +- .../proxy/management_endpoints/test_ui_sso.py | 199 +++++++++--------- 3 files changed, 111 insertions(+), 100 deletions(-) diff --git a/tests/local_testing/test_completion_cost.py b/tests/local_testing/test_completion_cost.py index dd060a56d2..43cb236ad4 100644 --- a/tests/local_testing/test_completion_cost.py +++ b/tests/local_testing/test_completion_cost.py @@ -738,10 +738,10 @@ def test_vertex_ai_embedding_completion_cost(caplog): text = "The quick brown fox jumps over the lazy dog." input_tokens = litellm.token_counter( - model="vertex_ai/textembedding-gecko", text=text + model="vertex_ai/text-embedding-004", text=text ) - model_info = litellm.get_model_info(model="vertex_ai/textembedding-gecko") + model_info = litellm.get_model_info(model="vertex_ai/text-embedding-004") print("\nExpected model info:\n{}\n\n".format(model_info)) @@ -749,7 +749,7 @@ def test_vertex_ai_embedding_completion_cost(caplog): ## CALCULATED COST calculated_input_cost, calculated_output_cost = cost_per_token( - model="textembedding-gecko", + model="text-embedding-004", custom_llm_provider="vertex_ai", prompt_tokens=input_tokens, call_type="aembedding", @@ -824,7 +824,7 @@ async def test_completion_cost_hidden_params(sync_mode): def test_vertex_ai_gemini_predict_cost(): - model = "gemini-1.5-flash" + model = "gemini-2.0-flash" messages = [{"role": "user", "content": "Hey, hows it going???"}] predictive_cost = completion_cost(model=model, messages=messages) diff --git a/tests/otel_tests/test_e2e_budgeting.py b/tests/otel_tests/test_e2e_budgeting.py index 0aaf162b78..62fc8732eb 100644 --- a/tests/otel_tests/test_e2e_budgeting.py +++ b/tests/otel_tests/test_e2e_budgeting.py @@ -8,13 +8,13 @@ from typing import Any, Optional async def make_calls_until_budget_exceeded(session, key: str, call_function, **kwargs): """Helper function to make API calls until budget is exceeded. Verify that the budget is exceeded error is returned.""" - MAX_CALLS = 50 + MAX_CALLS = 200 call_count = 0 try: while call_count < MAX_CALLS: await call_function(session=session, key=key, **kwargs) call_count += 1 - await asyncio.sleep(0.5) # allow spend tracking to catch up + await asyncio.sleep(0.1) # allow spend tracking to catch up pytest.fail(f"Budget was not exceeded after {MAX_CALLS} calls") except Exception as e: print("vars: ", vars(e)) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 95ebc29fef..ca480b8442 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3400,8 +3400,7 @@ class TestPKCEFunctionality: @pytest.mark.asyncio async def test_pkce_token_exchange_basic_auth(self): - """When include_client_id=False, client credentials go via HTTP Basic Auth - (encoded in the Authorization header, not via httpx.BasicAuth).""" + """When include_client_id=False, client credentials go via HTTP Basic Auth.""" token_resp = { "access_token": "tok_abc", "id_token": None, @@ -3418,22 +3417,34 @@ class TestPKCEFunctionality: mock_userinfo_response.status_code = 200 mock_userinfo_response.json.return_value = userinfo_resp - captured_post_kwargs = {} - async def fake_post(*args, **kwargs): - captured_post_kwargs.update(kwargs) + # Verify Basic Auth is set + assert "auth" in kwargs + assert isinstance(kwargs["auth"], httpx.BasicAuth) + # Verify code_verifier is in the POST body (essential PKCE field) + post_data = kwargs.get("data", {}) + assert post_data.get("code_verifier") == "verifier_abc" + # Verify redirect_uri is forwarded (required by strict OAuth providers) + assert post_data.get("redirect_uri") == "https://proxy.example.com/callback" + # Verify credentials are NOT double-sent in the POST body when using Basic Auth + assert "client_secret" not in post_data, "client_secret must not appear in POST body when using Basic Auth" + assert "client_id" not in post_data, "client_id must not appear in POST body when using Basic Auth (include_client_id=False)" return mock_response - # get_async_httpx_client returns an AsyncHTTPHandler directly (no context manager). - # _pkce_token_exchange calls it once, _get_pkce_userinfo calls it once. - mock_token_client = MagicMock() + # Use separate mock clients for token exchange and userinfo — + # each httpx.AsyncClient() call gets its own independent mock. + mock_token_client = AsyncMock() + mock_token_client.__aenter__ = AsyncMock(return_value=mock_token_client) + mock_token_client.__aexit__ = AsyncMock(return_value=False) mock_token_client.post = AsyncMock(side_effect=fake_post) - mock_userinfo_client = MagicMock() + mock_userinfo_client = AsyncMock() + mock_userinfo_client.__aenter__ = AsyncMock(return_value=mock_userinfo_client) + mock_userinfo_client.__aexit__ = AsyncMock(return_value=False) mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo_response) - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: - mock_get_client.side_effect = [mock_token_client, mock_userinfo_client] + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: + mock_client_cls.side_effect = [mock_token_client, mock_userinfo_client] result = await SSOAuthenticationHandler._pkce_token_exchange( authorization_code="auth_code_123", @@ -3447,18 +3458,6 @@ class TestPKCEFunctionality: additional_headers={}, ) - # Verify Basic Auth is set in headers (not via httpx.BasicAuth kwarg) - post_headers = captured_post_kwargs.get("headers", {}) - assert "Basic " in post_headers.get("Authorization", ""), "Basic Auth header must be set" - # Verify code_verifier is in the POST body (essential PKCE field) - post_data = captured_post_kwargs.get("data", {}) - assert post_data.get("code_verifier") == "verifier_abc" - # Verify redirect_uri is forwarded (required by strict OAuth providers) - assert post_data.get("redirect_uri") == "https://proxy.example.com/callback" - # Verify credentials are NOT double-sent in the POST body when using Basic Auth - assert "client_secret" not in post_data, "client_secret must not appear in POST body when using Basic Auth" - assert "client_id" not in post_data, "client_id must not appear in POST body when using Basic Auth (include_client_id=False)" - assert result["access_token"] == "tok_abc" assert result["email"] == "user@example.com" # id_token was explicit null in token_response — the merge loop must remove it @@ -3481,10 +3480,13 @@ class TestPKCEFunctionality: } userinfo_resp = {"sub": "user2", "email": "user2@example.com"} - captured_post_kwargs = {} - async def fake_post(*args, **kwargs): - captured_post_kwargs.update(kwargs) + assert "auth" not in kwargs, "Should NOT use Basic Auth when include_client_id=True" + data = kwargs.get("data", {}) + assert "client_id" in data + assert "client_secret" in data + assert data.get("code_verifier") == "verifier_xyz", "code_verifier must be in POST body" + assert data.get("redirect_uri") == "https://proxy.example.com/callback", "redirect_uri must be forwarded" mock = MagicMock() mock.status_code = 200 mock.json.return_value = token_resp @@ -3494,14 +3496,18 @@ class TestPKCEFunctionality: mock_userinfo.status_code = 200 mock_userinfo.json.return_value = userinfo_resp - mock_token_client = MagicMock() + mock_token_client = AsyncMock() + mock_token_client.__aenter__ = AsyncMock(return_value=mock_token_client) + mock_token_client.__aexit__ = AsyncMock(return_value=False) mock_token_client.post = AsyncMock(side_effect=fake_post) - mock_userinfo_client = MagicMock() + mock_userinfo_client = AsyncMock() + mock_userinfo_client.__aenter__ = AsyncMock(return_value=mock_userinfo_client) + mock_userinfo_client.__aexit__ = AsyncMock(return_value=False) mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo) - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: - mock_get_client.side_effect = [mock_token_client, mock_userinfo_client] + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: + mock_client_cls.side_effect = [mock_token_client, mock_userinfo_client] result = await SSOAuthenticationHandler._pkce_token_exchange( authorization_code="auth_code_456", @@ -3515,15 +3521,6 @@ class TestPKCEFunctionality: additional_headers={}, ) - # Verify no Basic Auth header when include_client_id=True - post_headers = captured_post_kwargs.get("headers", {}) - assert "Basic " not in post_headers.get("Authorization", ""), "Should NOT use Basic Auth when include_client_id=True" - data = captured_post_kwargs.get("data", {}) - assert "client_id" in data - assert "client_secret" in data - assert data.get("code_verifier") == "verifier_xyz", "code_verifier must be in POST body" - assert data.get("redirect_uri") == "https://proxy.example.com/callback", "redirect_uri must be forwarded" - assert result["access_token"] == "tok_body" assert result["sub"] == "user2" # Verify userinfo GET used the correct Bearer token header @@ -3539,14 +3536,16 @@ class TestPKCEFunctionality: error_body = {"error": "invalid_grant", "error_description": "Code already used"} - mock_resp = MagicMock() - mock_resp.status_code = 200 - mock_resp.json.return_value = error_body + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = error_body + mock_client.post = AsyncMock(return_value=mock_resp) + mock_client_cls.return_value = mock_client - mock_client = MagicMock() - mock_client.post = AsyncMock(return_value=mock_resp) - - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client", return_value=mock_client): with pytest.raises(ProxyException) as exc_info: await SSOAuthenticationHandler._pkce_token_exchange( authorization_code="expired_code", @@ -3577,13 +3576,15 @@ class TestPKCEFunctionality: ).rstrip(b"=").decode() fake_id_token = f"eyJhbGciOiJSUzI1NiJ9.{encoded_payload}.fakesig" - mock_fail = MagicMock() - mock_fail.status_code = 503 + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_fail = MagicMock() + mock_fail.status_code = 503 + mock_client.get = AsyncMock(return_value=mock_fail) + mock_client_cls.return_value = mock_client - mock_client = MagicMock() - mock_client.get = AsyncMock(return_value=mock_fail) - - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client", return_value=mock_client): result = await SSOAuthenticationHandler._get_pkce_userinfo( access_token="some_token", id_token=fake_id_token, @@ -3627,13 +3628,15 @@ class TestPKCEFunctionality: from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - mock_fail = MagicMock() - mock_fail.status_code = 503 + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_fail = MagicMock() + mock_fail.status_code = 503 + mock_client.get = AsyncMock(return_value=mock_fail) + mock_client_cls.return_value = mock_client - mock_client = MagicMock() - mock_client.get = AsyncMock(return_value=mock_fail) - - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client", return_value=mock_client): with pytest.raises(ProxyException) as exc_info: await SSOAuthenticationHandler._get_pkce_userinfo( access_token="token", @@ -3656,10 +3659,13 @@ class TestPKCEFunctionality: mock_resp.status_code = 200 mock_resp.json.return_value = None # HTTP 200 with null JSON body - mock_client = MagicMock() - mock_client.get = AsyncMock(return_value=mock_resp) + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_client.get = AsyncMock(return_value=mock_resp) + mock_client_cls.return_value = mock_client - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client", return_value=mock_client): with pytest.raises(ProxyException) as exc_info: await SSOAuthenticationHandler._get_pkce_userinfo( access_token="access_token", @@ -3718,10 +3724,12 @@ class TestPKCEFunctionality: } userinfo_resp = {"sub": "pubuser", "email": "pub@example.com"} - captured_post_kwargs = {} - async def fake_post(*args, **kwargs): - captured_post_kwargs.update(kwargs) + assert "auth" not in kwargs, "Public client must not use Basic Auth" + data = kwargs.get("data", {}) + assert data.get("client_id") == "public_client_id" + assert "client_secret" not in data, "No secret should be sent for public client" + assert data.get("code_verifier") == "public_verifier" mock = MagicMock() mock.status_code = 200 mock.json.return_value = token_resp @@ -3731,14 +3739,18 @@ class TestPKCEFunctionality: mock_userinfo.status_code = 200 mock_userinfo.json.return_value = userinfo_resp - mock_token_client = MagicMock() + mock_token_client = AsyncMock() + mock_token_client.__aenter__ = AsyncMock(return_value=mock_token_client) + mock_token_client.__aexit__ = AsyncMock(return_value=False) mock_token_client.post = AsyncMock(side_effect=fake_post) - mock_userinfo_client = MagicMock() + mock_userinfo_client = AsyncMock() + mock_userinfo_client.__aenter__ = AsyncMock(return_value=mock_userinfo_client) + mock_userinfo_client.__aexit__ = AsyncMock(return_value=False) mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo) - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client") as mock_get_client: - mock_get_client.side_effect = [mock_token_client, mock_userinfo_client] + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: + mock_client_cls.side_effect = [mock_token_client, mock_userinfo_client] result = await SSOAuthenticationHandler._pkce_token_exchange( authorization_code="auth_pub", @@ -3752,14 +3764,6 @@ class TestPKCEFunctionality: additional_headers={}, ) - # Verify no Basic Auth header for public client - post_headers = captured_post_kwargs.get("headers", {}) - assert "Basic " not in post_headers.get("Authorization", ""), "Public client must not use Basic Auth" - data = captured_post_kwargs.get("data", {}) - assert data.get("client_id") == "public_client_id" - assert "client_secret" not in data, "No secret should be sent for public client" - assert data.get("code_verifier") == "public_verifier" - assert result["access_token"] == "tok_public" assert result["sub"] == "pubuser" @@ -3880,10 +3884,13 @@ class TestPKCEFunctionality: mock_response.status_code = 401 mock_response.text = "Unauthorized" - mock_client = MagicMock() - mock_client.post = AsyncMock(return_value=mock_response) + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_client.post = AsyncMock(return_value=mock_response) + mock_client_cls.return_value = mock_client - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client", return_value=mock_client): with pytest.raises(ProxyException) as exc_info: await SSOAuthenticationHandler._pkce_token_exchange( authorization_code="auth_code", @@ -3986,15 +3993,17 @@ class TestPKCEFunctionality: from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - mock_resp = MagicMock() - mock_resp.status_code = 200 - mock_resp.json.return_value = None # JSON null response body - mock_resp.text = "null" + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = None # JSON null response body + mock_resp.text = "null" + mock_client.post = AsyncMock(return_value=mock_resp) + mock_client_cls.return_value = mock_client - mock_client = MagicMock() - mock_client.post = AsyncMock(return_value=mock_resp) - - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client", return_value=mock_client): with pytest.raises(ProxyException) as exc_info: await SSOAuthenticationHandler._pkce_token_exchange( authorization_code="some_code", @@ -4020,14 +4029,16 @@ class TestPKCEFunctionality: body_without_token = {"token_type": "Bearer", "scope": "openid"} - mock_resp = MagicMock() - mock_resp.status_code = 200 - mock_resp.json.return_value = body_without_token + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = body_without_token + mock_client.post = AsyncMock(return_value=mock_resp) + mock_client_cls.return_value = mock_client - mock_client = MagicMock() - mock_client.post = AsyncMock(return_value=mock_resp) - - with patch("litellm.proxy.management_endpoints.ui_sso.get_async_httpx_client", return_value=mock_client): with pytest.raises(ProxyException) as exc_info: await SSOAuthenticationHandler._pkce_token_exchange( authorization_code="some_code",