From 0fda9da174c1784dc6ded243b5b7bf3c475eb74f Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 15 Apr 2025 17:03:48 -0700
Subject: [PATCH] test_get_users_includes_timestamps

---
 .../test_internal_user_endpoints.py           | 73 ++++++++++++++++++-
 1 file changed, 72 insertions(+), 1 deletion(-)

diff --git a/tests/litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/litellm/proxy/management_endpoints/test_internal_user_endpoints.py
index 54fda943eb..73ed8e1a30 100644
--- a/tests/litellm/proxy/management_endpoints/test_internal_user_endpoints.py
+++ b/tests/litellm/proxy/management_endpoints/test_internal_user_endpoints.py
@@ -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