From 4b117975f0fa66c9548aae974bec509d3a83feff Mon Sep 17 00:00:00 2001
From: Alexsander Hamir <alexsanderhamirgomesbaptista@gmail.com>
Date: Mon, 26 Jan 2026 10:11:56 -0800
Subject: [PATCH] Make test_get_users_key_count deterministic by creating
 dedicated test user

- Create a test user with auto_create_key=False to ensure known starting state
- Filter get_users by user_ids to target only the test user
- Verify initial key count is 0 before creating a key
- Clean up test user after test completes
- This ensures consistent behavior across CI and local environments
---
 .../test_key_management.py                    | 56 ++++++++++++++-----
 1 file changed, 43 insertions(+), 13 deletions(-)

diff --git a/tests/proxy_admin_ui_tests/test_key_management.py b/tests/proxy_admin_ui_tests/test_key_management.py
index a196080eada..fa39a05a27d 100644
--- a/tests/proxy_admin_ui_tests/test_key_management.py
+++ b/tests/proxy_admin_ui_tests/test_key_management.py
@@ -473,18 +473,41 @@ async def test_get_users_key_count(prisma_client):
     setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
     await litellm.proxy.proxy_server.prisma_client.connect()
 
-    # Get initial user list and select the first user
-    initial_users = await get_users(role=None, page=1, page_size=20)
+    # Create a test user with no initial keys to ensure deterministic behavior
+    test_user_id = f"test_user_key_count-{uuid.uuid4()}"
+    test_user_request = NewUserRequest(
+        user_id=test_user_id,
+        user_role=LitellmUserRoles.INTERNAL_USER.value,
+        auto_create_key=False,  # Ensure we start with 0 keys
+    )
+
+    await new_user(
+        test_user_request,
+        UserAPIKeyAuth(
+            user_role=LitellmUserRoles.PROXY_ADMIN,
+            api_key="sk-1234",
+            user_id="admin",
+        ),
+    )
+
+    # Get initial key count for the test user
+    initial_users = await get_users(
+        user_ids=test_user_id,
+        role=None,
+        page=1,
+        page_size=20,
+    )
     print("initial_users", initial_users)
-    assert len(initial_users["users"]) > 0, "No users found to test with"
-
+    assert len(initial_users["users"]) == 1, "Test user should be found"
     test_user = initial_users["users"][0]
+    assert test_user.user_id == test_user_id
     initial_key_count = test_user.key_count
+    assert initial_key_count == 0, f"Expected initial key count to be 0, but got {initial_key_count}"
 
-    # Create a new key for the selected user
+    # Create a new key for the test user
     new_key = await generate_key_fn(
         data=GenerateKeyRequest(
-            user_id=test_user.user_id,
+            user_id=test_user_id,
             key_alias=f"test_key_{uuid.uuid4()}",
             models=["fake-model"],
         ),
@@ -496,19 +519,26 @@ async def test_get_users_key_count(prisma_client):
     )
 
     # Get updated user list and check key count
-    updated_users = await get_users(role=None, page=1, page_size=20)
+    updated_users = await get_users(
+        user_ids=test_user_id,
+        role=None,
+        page=1,
+        page_size=20,
+    )
     print("updated_users", updated_users)
-    updated_key_count = None
-    for user in updated_users["users"]:
-        if user.user_id == test_user.user_id:
-            updated_key_count = user.key_count
-            break
+    assert len(updated_users["users"]) == 1, "Test user should still be found"
+    updated_user = updated_users["users"][0]
+    updated_key_count = updated_user.key_count
 
-    assert updated_key_count is not None, "Test user not found in updated users list"
     assert (
         updated_key_count == initial_key_count + 1
     ), f"Expected key count to increase by 1, but got {updated_key_count} (was {initial_key_count})"
 
+    # Clean up test user and keys
+    await prisma_client.db.litellm_usertable.delete(
+        where={"user_id": test_user_id}
+    )
+
 
 async def cleanup_existing_teams(prisma_client):
     all_teams = await prisma_client.db.litellm_teamtable.find_many()