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
This commit is contained in:
Alexsander Hamir 2026-01-26 10:11:56 -08:00
parent 8ece284ca3
commit 4b117975f0

View file

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