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

(test) add test fixture for dynamodb testing

Browse source
This commit is contained in:
ishaan-jaff 2024-01-17 12:23:04 -08:00
parent 4036c64bdf
commit 2273f8055c
1 changed files with 24 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

40
litellm/tests/test_key_generate_dynamodb.py
View file

@ -32,16 +32,6 @@ from litellm.proxy.utils import DBClient
from starlette.datastructures import URL from starlette.datastructures import URL
db_args = {
"ssl_verify": False,
"billing_mode": "PAY_PER_REQUEST",
"region_name": "us-west-2",
}
custom_db_client = DBClient(
custom_db_type="dynamo_db",
custom_db_args=db_args,
)
request_data = { request_data = {
"model": "azure-gpt-3.5", "model": "azure-gpt-3.5",
"messages": [ "messages": [
@ -50,7 +40,25 @@ request_data = {
} }
def test_generate_and_call_with_valid_key(): @pytest.fixture
def custom_db_client():
# Assuming DBClient is a class that needs to be instantiated
db_args = {
"ssl_verify": False,
"billing_mode": "PAY_PER_REQUEST",
"region_name": "us-west-2",
}
custom_db_client = DBClient(
custom_db_type="dynamo_db",
custom_db_args=db_args,
)
# Reset litellm.proxy.proxy_server.prisma_client to None
litellm.proxy.proxy_server.prisma_client = None
return custom_db_client
def test_generate_and_call_with_valid_key(custom_db_client):
# 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") setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
@ -76,7 +84,7 @@ 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(): def test_call_with_invalid_key(custom_db_client):
# 2. Make a call with invalid key, expect it to fail # 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, "custom_db_client", custom_db_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
@ -101,7 +109,7 @@ def test_call_with_invalid_key():
pass pass
def test_call_with_invalid_model(): def test_call_with_invalid_model(custom_db_client):
# 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
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") setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
@ -136,7 +144,7 @@ def test_call_with_invalid_model():
pass pass
def test_call_with_valid_model(): def test_call_with_valid_model(custom_db_client):
# 4. Make a call to a key with a valid model - expect to pass # 4. Make a call to a key with a valid model - expect to pass
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") setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
@ -167,7 +175,7 @@ def test_call_with_valid_model():
pytest.fail(f"An exception occurred - {str(e)}") pytest.fail(f"An exception occurred - {str(e)}")
def test_call_with_key_over_budget(): def test_call_with_key_over_budget(custom_db_client):
# 5. Make a call with a key over budget, expect to fail # 5. Make a call with a key over budget, expect to fail
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") setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
@ -233,7 +241,7 @@ def test_call_with_key_over_budget():
print(vars(e)) print(vars(e))
def test_call_with_key_over_budget_stream(): def test_call_with_key_over_budget_stream(custom_db_client):
# 6. Make a call with a key over budget, expect to fail # 6. Make a call with a key over budget, expect to fail
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") setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")