test_enforce_unique_key_alias

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -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

tests/proxy_unit_tests/test_key_generate_prisma.py

@@ -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)}")