test_get_users_includes_timestamps

This commit is contained in:
Ishaan Jaff 2025-04-15 17:03:48 -07:00
parent 725b726aed
commit 0fda9da174

View file

@ -1,6 +1,7 @@
import json
import os
import sys
from datetime import datetime, timezone
import pytest
from fastapi.testclient import TestClient
@ -10,7 +11,12 @@ sys.path.insert(
) # Adds the parent directory to the system path
from litellm.proxy._types import LiteLLM_UserTableFiltered, UserAPIKeyAuth
from litellm.proxy.management_endpoints.internal_user_endpoints import ui_view_users
from litellm.proxy.management_endpoints.internal_user_endpoints import (
LiteLLM_UserTableWithKeyCount,
get_user_key_counts,
get_users,
ui_view_users,
)
from litellm.proxy.proxy_server import app
client = TestClient(app)
@ -82,3 +88,68 @@ def test_user_daily_activity_types():
assert not hasattr(
daily_spend_metadata, field
), f"Field {field} is reported in DailySpendMetadata"
@pytest.mark.asyncio
async def test_get_users_includes_timestamps(mocker):
"""
Test that /user/list endpoint returns users with created_at and updated_at fields.
"""
# Mock the prisma client
mock_prisma_client = mocker.MagicMock()
# Create mock user data with timestamps
mock_user_data = {
"user_id": "test-user-timestamps",
"user_email": "timestamps@example.com",
"user_role": "internal_user",
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
}
mock_user_row = mocker.MagicMock()
mock_user_row.model_dump.return_value = mock_user_data
# Setup the mock find_many response as an async function
async def mock_find_many(*args, **kwargs):
return [mock_user_row]
# Setup the mock count response as an async function
async def mock_count(*args, **kwargs):
return 1
mock_prisma_client.db.litellm_usertable.find_many = mock_find_many
mock_prisma_client.db.litellm_usertable.count = mock_count
# Patch the prisma client import in the endpoint
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
# Mock the helper function get_user_key_counts
async def mock_get_user_key_counts(*args, **kwargs):
return {"test-user-timestamps": 0}
mocker.patch(
"litellm.proxy.management_endpoints.internal_user_endpoints.get_user_key_counts",
mock_get_user_key_counts,
)
# Call get_users function directly
response = await get_users(page=1, page_size=1)
print("user /list response: ", response)
# Assertions
assert response is not None
assert "users" in response
assert "total" in response
assert response["total"] == 1
assert len(response["users"]) == 1
user_response = response["users"][0]
assert user_response.user_id == "test-user-timestamps"
assert user_response.created_at is not None
assert isinstance(user_response.created_at, datetime)
assert user_response.updated_at is not None
assert isinstance(user_response.updated_at, datetime)
assert user_response.created_at == mock_user_data["created_at"]
assert user_response.updated_at == mock_user_data["updated_at"]
assert user_response.key_count == 0