phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(user_api_key_auth.py): update valid token cache with updated team object cache

Browse source
This commit is contained in:
Krrish Dholakia 2024-07-19 17:06:49 -07:00
parent fa7037e48a
commit 35e640076b
8 changed files with 94 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

37
litellm/tests/test_user_api_key_auth.py
View file

@ -44,3 +44,40 @@ def test_check_valid_ip(
request = Request(client_ip)
assert _check_valid_ip(allowed_ips, request) == expected_result # type: ignore
@pytest.mark.asyncio
async def test_check_blocked_team():
"""
cached valid_token obj has team_blocked = true
cached team obj has team_blocked = false
assert team is not blocked
"""
from fastapi import Request
from starlette.datastructures import URL
from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.proxy_server import hash_token, user_api_key_cache
_team_id = "1234"
team_obj = LiteLLM_TeamTable(team_id=_team_id, blocked=False)
user_key = "sk-12345678"
valid_token = UserAPIKeyAuth(
team_id=_team_id, team_blocked=True, token=hash_token(user_key)
)
user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
setattr(litellm.proxy.proxy_server, "prisma_client", "hello-world")
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
await user_api_key_auth(request=request, api_key="Bearer " + user_key)