test(test_key_generate_prisma.py): fix test import

This commit is contained in:
Krrish Dholakia 2024-02-10 16:16:12 -08:00
parent 34c118e5e9
commit f4dc643b1c

View file

@ -59,7 +59,7 @@ from litellm.proxy._types import (
NewUserRequest, NewUserRequest,
GenerateKeyRequest, GenerateKeyRequest,
DynamoDBArgs, DynamoDBArgs,
DeleteKeyRequest, KeyRequest,
UpdateKeyRequest, UpdateKeyRequest,
GenerateKeyRequest, GenerateKeyRequest,
) )
@ -740,7 +740,7 @@ def test_delete_key(prisma_client):
generated_key = key.key generated_key = key.key
bearer_token = "Bearer " + generated_key bearer_token = "Bearer " + generated_key
delete_key_request = DeleteKeyRequest(keys=[generated_key]) delete_key_request = KeyRequest(keys=[generated_key])
# delete the key # delete the key
result_delete_key = await delete_key_fn(data=delete_key_request) result_delete_key = await delete_key_fn(data=delete_key_request)
@ -778,7 +778,7 @@ def test_delete_key_auth(prisma_client):
generated_key = key.key generated_key = key.key
bearer_token = "Bearer " + generated_key bearer_token = "Bearer " + generated_key
delete_key_request = DeleteKeyRequest(keys=[generated_key]) delete_key_request = KeyRequest(keys=[generated_key])
# delete the key # delete the key
result_delete_key = await delete_key_fn(data=delete_key_request) result_delete_key = await delete_key_fn(data=delete_key_request)
@ -839,7 +839,7 @@ def test_generate_and_call_key_info(prisma_client):
} }
# cleanup - delete key # cleanup - delete key
delete_key_request = DeleteKeyRequest(keys=[generated_key]) delete_key_request = KeyRequest(keys=[generated_key])
# delete the key # delete the key
await delete_key_fn(data=delete_key_request) await delete_key_fn(data=delete_key_request)
@ -908,7 +908,7 @@ def test_generate_and_update_key(prisma_client):
assert result["info"]["models"] == ["ada", "babbage", "curie", "davinci"] assert result["info"]["models"] == ["ada", "babbage", "curie", "davinci"]
# cleanup - delete key # cleanup - delete key
delete_key_request = DeleteKeyRequest(keys=[generated_key]) delete_key_request = KeyRequest(keys=[generated_key])
# delete the key # delete the key
await delete_key_fn(data=delete_key_request) await delete_key_fn(data=delete_key_request)