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

(test) /key/gen

Browse source
This commit is contained in:
ishaan-jaff 2024-01-16 17:36:33 -08:00
parent 57af57d4e9
commit aa676cc961
1 changed files with 32 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

45
litellm/tests/test_key_generate_dynamodb.py
View file

@ -32,15 +32,14 @@ from litellm.proxy.proxy_server import new_user, user_api_key_auth
from litellm.proxy._types import NewUserRequest, DynamoDBArgs from litellm.proxy._types import NewUserRequest, DynamoDBArgs
from litellm.proxy.utils import DBClient from litellm.proxy.utils import DBClient
db_args = DynamoDBArgs( db_args = {
ssl_verify=False, "ssl_verify": False,
billing_mode="PAY_PER_REQUEST", "billing_mode": "PAY_PER_REQUEST",
region_name="us-west-2", "region_name": "us-west-2",
) }
db_args_dict = db_args.model_dump()
custom_db_client = DBClient( custom_db_client = DBClient(
custom_db_type="dynamo_db", custom_db_type="dynamo_db",
custom_db_args=db_args_dict, custom_db_args=db_args,
) )
request_data = { request_data = {
@ -54,6 +53,7 @@ request_data = {
def test_generate_and_call_with_valid_key(): def test_generate_and_call_with_valid_key():
# 1. Generate a Key, and use it to make a call # 1. Generate a Key, and use it to make a call
setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client) setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
try: try:
async def test(): async def test():
@ -63,6 +63,31 @@ def test_generate_and_call_with_valid_key():
generated_key = key.key generated_key = key.key
bearer_token = "Bearer " + generated_key bearer_token = "Bearer " + generated_key
from starlette.datastructures import URL
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
# 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)
asyncio.run(test())
# result = user_auth(ValidRequest(key))
# assert result is True
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
def test_call_with_invalid_key():
# 2. Make a call with invalid key, expect it to fail
setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
try:
async def test():
generated_key = "bad-key"
bearer_token = "Bearer " + generated_key
request = Request(scope={"type": "http"}) request = Request(scope={"type": "http"})
@ -77,12 +102,6 @@ def test_generate_and_call_with_valid_key():
pytest.fail(f"An exception occurred - {str(e)}") pytest.fail(f"An exception occurred - {str(e)}")
# def test_call_with_invalid_key():
# # 2. Make a call with invalid key, expect it to fail
# result = user_auth(InvalidKeyRequest())
# assert result is False
# def test_call_with_invalid_model(): # def test_call_with_invalid_model():
# # 3. Make a call to a key with an invalid model - expect to fail # # 3. Make a call to a key with an invalid model - expect to fail
# key = new_user(ValidNewUserRequest()) # key = new_user(ValidNewUserRequest())