From 9c29f35c4b26524205770a5ed7fe779745ec90ee Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 2 Oct 2025 18:48:11 -0700 Subject: [PATCH] test_end_user_jwt_auth --- ...odel_prices_and_context_window_backup.json | 20 +++++++ tests/proxy_unit_tests/test_jwt.py | 60 +++++++++++++++---- 2 files changed, 67 insertions(+), 13 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bae7c11e66c..1b8f7f7c08f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4801,6 +4801,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 64000, @@ -4827,6 +4831,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 64000, @@ -19720,6 +19728,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 200000, "max_output_tokens": 64000, @@ -21086,6 +21098,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "input_cost_per_token_batches": 1.5e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -21108,6 +21124,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "input_cost_per_token_batches": 1.5e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/proxy_unit_tests/test_jwt.py index 2e54ade688f..ed75a56361a 100644 --- a/tests/proxy_unit_tests/test_jwt.py +++ b/tests/proxy_unit_tests/test_jwt.py @@ -1155,15 +1155,26 @@ async def test_end_user_jwt_auth(monkeypatch): ## 1. INITIAL TEAM CALL - should fail # use generated key to auth in + from litellm import Router + from litellm.types.router import RouterGeneralSettings + + # Create a router with pass_through_all_models enabled + router = Router( + model_list=[], + router_general_settings=RouterGeneralSettings( + pass_through_all_models=True + ), + ) + setattr( litellm.proxy.proxy_server, "general_settings", - {"enable_jwt_auth": True, "pass_through_all_models": True}, + {"enable_jwt_auth": True}, ) setattr( litellm.proxy.proxy_server, "llm_router", - MagicMock(), + router, ) setattr(litellm.proxy.proxy_server, "prisma_client", {}) setattr(litellm.proxy.proxy_server, "jwt_handler", jwt_handler) @@ -1171,18 +1182,39 @@ async def test_end_user_jwt_auth(monkeypatch): cost_tracking() result = await user_api_key_auth(request=request, api_key=bearer_token) - assert ( - result.end_user_id == "81b3e52a-67a6-4efb-9645-70527e101479" - ) # jwt token decoded sub value + + # Assert that end_user_id is correctly extracted from JWT token's 'sub' field + assert result.end_user_id == "81b3e52a-67a6-4efb-9645-70527e101479" temp_response = Response() from litellm.proxy.hooks.proxy_track_cost_callback import ( _should_track_cost_callback, ) - with patch.object( - litellm.proxy.hooks.proxy_track_cost_callback, "_should_track_cost_callback" - ) as mock_client: + # Mock the actual LLM completion call + mock_response = litellm.ModelResponse( + id="chatcmpl-mock", + choices=[ + litellm.Choices( + finish_reason="stop", + index=0, + message=litellm.Message( + content="Hello! I'm doing well, thank you for asking.", + role="assistant", + ), + ) + ], + created=1234567890, + model="gpt-4o", + object="chat.completion", + usage=litellm.Usage( + prompt_tokens=10, + completion_tokens=15, + total_tokens=25, + ), + ) + + with patch("litellm.acompletion", new=AsyncMock(return_value=mock_response)) as mock_completion: resp = await chat_completion( request=request, fastapi_response=temp_response, @@ -1194,11 +1226,13 @@ async def test_end_user_jwt_auth(monkeypatch): await asyncio.sleep(1) - mock_client.assert_called_once() - - mock_client.call_args.kwargs[ - "end_user_id" - ] == "81b3e52a-67a6-4efb-9645-70527e101479" + # Verify the completion was called with correct end_user_id + mock_completion.assert_called_once() + call_kwargs = mock_completion.call_args.kwargs + + # end_user_id is passed in metadata as 'user_api_key_end_user_id' + metadata = call_kwargs.get("metadata", {}) + assert metadata.get("user_api_key_end_user_id") == "81b3e52a-67a6-4efb-9645-70527e101479" def test_can_rbac_role_call_route():