test_end_user_jwt_auth

This commit is contained in:
Ishaan Jaffer 2025-10-02 18:48:11 -07:00
parent 54095a256a
commit 9c29f35c4b
2 changed files with 67 additions and 13 deletions

View file

@ -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,

View file

@ -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():