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

test new team request

Browse source
This commit is contained in:
Ishaan Jaff 2024-04-02 20:52:16 -07:00
parent ec1ba2e216
commit 865cfd8c9f
1 changed files with 45 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

67
litellm/tests/test_key_generate_prisma.py
View file

@ -50,6 +50,7 @@ from litellm.proxy.proxy_server import (
view_spend_logs, view_spend_logs,
user_info, user_info,
info_key_fn, info_key_fn,
new_team,
) )
from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token, update_spend from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token, update_spend
from litellm._logging import verbose_proxy_logger from litellm._logging import verbose_proxy_logger
@ -63,6 +64,8 @@ from litellm.proxy._types import (
KeyRequest, KeyRequest,
UpdateKeyRequest, UpdateKeyRequest,
GenerateKeyRequest, GenerateKeyRequest,
NewTeamRequest,
UserAPIKeyAuth,
) )
from litellm.proxy.utils import DBClient from litellm.proxy.utils import DBClient
from starlette.datastructures import URL from starlette.datastructures import URL
@ -246,40 +249,60 @@ def test_call_with_valid_model(prisma_client):
pytest.fail(f"An exception occurred - {str(e)}") pytest.fail(f"An exception occurred - {str(e)}")
def test_call_with_valid_model_using_all_models(prisma_client): @pytest.mark.asyncio
async def test_call_with_valid_model_using_all_models(prisma_client):
"""
Do not delete
this is the Admin UI flow
1. Create a team with model = `all-proxy-models`
2. Create a key with model = `all-team-models`
3. Call /chat/completions with the key -> expect to pass
"""
# Make a call to a key with model = `all-proxy-models` this is an Alias from LiteLLM Admin UI # Make a call to a key with model = `all-proxy-models` this is an Alias from LiteLLM Admin UI
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
try: try:
async def test(): await litellm.proxy.proxy_server.prisma_client.connect()
await litellm.proxy.proxy_server.prisma_client.connect()
request = GenerateKeyRequest(models=["all-proxy-models"])
key = await generate_key_fn(data=request)
print(key)
generated_key = key.key team_request = NewTeamRequest(
bearer_token = "Bearer " + generated_key team_alias="testing-team",
models=["all-proxy-models"],
)
request = Request(scope={"type": "http"}) new_team_response = await new_team(
request._url = URL(url="/chat/completions") data=team_request, user_api_key_dict=UserAPIKeyAuth(user_role="proxy_admin")
)
print("new_team_response", new_team_response)
created_team_id = new_team_response["team_id"]
async def return_body(): request = GenerateKeyRequest(
return b'{"model": "mistral"}' models=["all-team-models"], team_id=created_team_id
)
key = await generate_key_fn(data=request)
print(key)
request.body = return_body generated_key = key.key
bearer_token = "Bearer " + generated_key
# use generated key to auth in request = Request(scope={"type": "http"})
result = await user_api_key_auth(request=request, api_key=bearer_token) request._url = URL(url="/chat/completions")
print("result from user auth with new key", result)
# call /key/info for key - models == "all-proxy-models" async def return_body():
key_info = await info_key_fn(key=generated_key) return b'{"model": "mistral"}'
print("key_info", key_info)
models = key_info["info"]["models"] request.body = return_body
assert models == ["all-proxy-models"]
# use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token)
print("result from user auth with new key", result)
# call /key/info for key - models == "all-proxy-models"
key_info = await info_key_fn(key=generated_key)
print("key_info", key_info)
models = key_info["info"]["models"]
assert models == ["all-team-models"]
asyncio.run(test())
except Exception as e: except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}") pytest.fail(f"An exception occurred - {str(e)}")