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

32
litellm/tests/test_key_generate_dynamodb.py
View file

@ -32,6 +32,17 @@ from litellm.proxy.utils import DBClient
from starlette.datastructures import URL
request_data = {
"model": "azure-gpt-3.5",
"messages": [
{"role": "user", "content": "this is my new test. respond in 50 lines"}
],
}
@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",
@ -41,16 +52,13 @@ 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
request_data = {
"model": "azure-gpt-3.5",
"messages": [
{"role": "user", "content": "this is my new test. respond in 50 lines"}
],
}
return custom_db_client
def test_generate_and_call_with_valid_key():
def test_generate_and_call_with_valid_key(custom_db_client):
# 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, "master_key", "sk-1234")
@ -76,7 +84,7 @@ def test_generate_and_call_with_valid_key():
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
setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
@ -101,7 +109,7 @@ def test_call_with_invalid_key():
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
setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
@ -136,7 +144,7 @@ def test_call_with_invalid_model():
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
setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
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)}")
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
setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
@ -233,7 +241,7 @@ def test_call_with_key_over_budget():
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
setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")