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

test_enforce_unique_key_alias

Browse source
This commit is contained in:
Ishaan Jaff 2024-11-27 18:02:33 -08:00
parent 9c3acfbb5d
commit 818a118bda
2 changed files with 82 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

11
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -593,7 +593,7 @@ async def update_key_fn(
await _enforce_unique_key_alias(
key_alias=non_default_values.get("key_alias", None),
prisma_client=prisma_client,
existing_key_id=existing_key_row.token,
existing_key_token=existing_key_row.token,
)
response = await prisma_client.update_data(
@ -1899,7 +1899,7 @@ async def test_key_logging(
async def _enforce_unique_key_alias(
key_alias: Optional[str],
prisma_client: Any,
existing_key_id: Optional[str] = None,
existing_key_token: Optional[str] = None,
) -> None:
"""
Helper to enforce unique key aliases across all keys.
@ -1907,16 +1907,17 @@ async def _enforce_unique_key_alias(
Args:
key_alias (Optional[str]): The key alias to check
prisma_client (Any): Prisma client instance
existing_key_id (Optional[str]): ID of existing key being updated, to exclude from uniqueness check
existing_key_token (Optional[str]): ID of existing key being updated, to exclude from uniqueness check
(The Admin UI passes key_alias, in all Edit key requests. So we need to be sure that if we find a key with the same alias, it's not the same key we're updating)
Raises:
HTTPException: If key alias already exists on a different key
"""
if key_alias is not None and prisma_client is not None:
where_clause: dict[str, Any] = {"key_alias": key_alias}
if existing_key_id:
if existing_key_token:
# Exclude the current key from the uniqueness check
where_clause["NOT"] = {"token": existing_key_id}
where_clause["NOT"] = {"token": existing_key_token}
existing_key = await prisma_client.db.litellm_verificationtoken.find_first(
where=where_clause

76
tests/proxy_unit_tests/test_key_generate_prisma.py
View file

@ -3623,3 +3623,79 @@ async def test_key_alias_uniqueness(prisma_client):
print("got exceptions, e=", e)
print("vars(e)=", vars(e))
pytest.fail(f"An unexpected error occurred: {str(e)}")
@pytest.mark.asyncio
async def test_enforce_unique_key_alias(prisma_client):
"""
Unit test the _enforce_unique_key_alias function:
1. Test it allows unique aliases
2. Test it blocks duplicate aliases for new keys
3. Test it allows updating a key with its own existing alias
4. Test it blocks updating a key with another key's alias
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
_enforce_unique_key_alias,
)
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
await litellm.proxy.proxy_server.prisma_client.connect()
try:
# Test 1: Allow unique alias
unique_alias = f"test-alias-{uuid.uuid4()}"
await _enforce_unique_key_alias(
key_alias=unique_alias,
prisma_client=prisma_client,
) # Should pass
# Create a key with this alias in the database
key1 = await generate_key_fn(
data=GenerateKeyRequest(key_alias=unique_alias),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-1234",
user_id="1234",
),
)
# Test 2: Block duplicate alias for new key
try:
await _enforce_unique_key_alias(
key_alias=unique_alias,
prisma_client=prisma_client,
)
pytest.fail("Should not allow duplicate alias")
except Exception as e:
assert "Unique key aliases across all keys are required" in str(e.message)
# Test 3: Allow updating key with its own alias
await _enforce_unique_key_alias(
key_alias=unique_alias,
existing_key_token=hash_token(key1.key),
prisma_client=prisma_client,
) # Should pass
# Test 4: Block updating with another key's alias
another_key = await generate_key_fn(
data=GenerateKeyRequest(key_alias=f"test-alias-{uuid.uuid4()}"),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-1234",
user_id="1234",
),
)
try:
await _enforce_unique_key_alias(
key_alias=unique_alias,
existing_key_token=another_key.key,
prisma_client=prisma_client,
)
pytest.fail("Should not allow using another key's alias")
except Exception as e:
assert "Unique key aliases across all keys are required" in str(e.message)
except Exception as e:
print("Unexpected error:", e)
pytest.fail(f"An unexpected error occurred: {str(e)}")