mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
test_end_user_jwt_auth
This commit is contained in:
parent
54095a256a
commit
9c29f35c4b
2 changed files with 67 additions and 13 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue