phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(test) prisma key gen

Browse source
This commit is contained in:
ishaan-jaff 2024-01-17 11:58:14 -08:00
parent ec68c3a29f
commit 4036c64bdf
1 changed files with 16 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

21
litellm/tests/test_key_generate_prisma.py
View file

@ -27,6 +27,9 @@ import pytest, logging, asyncio
import litellm, asyncio import litellm, asyncio
from litellm.proxy.proxy_server import new_user, user_api_key_auth, user_update from litellm.proxy.proxy_server import new_user, user_api_key_auth, user_update
from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm._logging import verbose_proxy_logger
verbose_proxy_logger.setLevel(level=logging.DEBUG)
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
@ -35,10 +38,6 @@ from litellm.caching import DualCache
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
prisma_client = PrismaClient(
database_url=os.environ["PROXY_DATABASE_URL"], proxy_logging_obj=proxy_logging_obj
)
request_data = { request_data = {
"model": "azure-gpt-3.5", "model": "azure-gpt-3.5",
@ -46,16 +45,22 @@ request_data = {
{"role": "user", "content": "this is my new test. respond in 50 lines"} {"role": "user", "content": "this is my new test. respond in 50 lines"}
], ],
} }
prisma_client = PrismaClient(
database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj
)
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
prisma_client.connect()
print("prisma client=", prisma_client)
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(): async def test():
await litellm.litellm.proxy.proxy_server.prisma_client.connect()
request = NewUserRequest() request = NewUserRequest()
key = await new_user(request) key = await new_user(request)
print(key) print(key)
@ -82,6 +87,7 @@ def test_call_with_invalid_key():
try: try:
async def test(): async def test():
await litellm.litellm.proxy.proxy_server.prisma_client.connect()
generated_key = "bad-key" generated_key = "bad-key"
bearer_token = "Bearer " + generated_key bearer_token = "Bearer " + generated_key
@ -90,6 +96,7 @@ def test_call_with_invalid_key():
# use generated key to auth in # use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token) result = await user_api_key_auth(request=request, api_key=bearer_token)
print("got result", result)
pytest.fail(f"This should have failed!. IT's an invalid key") pytest.fail(f"This should have failed!. IT's an invalid key")
asyncio.run(test()) asyncio.run(test())
@ -107,6 +114,7 @@ def test_call_with_invalid_model():
try: try:
async def test(): async def test():
await litellm.litellm.proxy.proxy_server.prisma_client.connect()
request = NewUserRequest(models=["mistral"]) request = NewUserRequest(models=["mistral"])
key = await new_user(request) key = await new_user(request)
print(key) print(key)
@ -142,6 +150,7 @@ def test_call_with_valid_model():
try: try:
async def test(): async def test():
await litellm.litellm.proxy.proxy_server.prisma_client.connect()
request = NewUserRequest(models=["mistral"]) request = NewUserRequest(models=["mistral"])
key = await new_user(request) key = await new_user(request)
print(key) print(key)
@ -173,6 +182,7 @@ def test_call_with_key_over_budget():
try: try:
async def test(): async def test():
await litellm.litellm.proxy.proxy_server.prisma_client.connect()
request = NewUserRequest(max_budget=0.00001) request = NewUserRequest(max_budget=0.00001)
key = await new_user(request) key = await new_user(request)
print(key) print(key)
@ -244,6 +254,7 @@ def test_call_with_key_over_budget_stream():
try: try:
async def test(): async def test():
await litellm.litellm.proxy.proxy_server.prisma_client.connect()
request = NewUserRequest(max_budget=0.00001) request = NewUserRequest(max_budget=0.00001)
key = await new_user(request) key = await new_user(request)
print(key) print(key)