From ae44022a407fe25dac5312dff3f973b276f84b55 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 12 Feb 2026 21:37:19 -0800 Subject: [PATCH] Fixing spend tests --- .../test_spend_management_endpoints.py | 728 ++++++------------ 1 file changed, 231 insertions(+), 497 deletions(-) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 0d4a711812f..aaa14ebd1e9 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -15,7 +15,86 @@ sys.path.insert( from unittest.mock import AsyncMock, MagicMock, patch import litellm + import litellm.proxy.proxy_server as ps + + +def _default_date_range(): + """Return (start_date, end_date) for the common 7-day range used in UI spend tests.""" + now = datetime.datetime.now(timezone.utc) + return ( + (now - datetime.timedelta(days=7)).strftime("%Y-%m-%d %H:%M:%S"), + now.strftime("%Y-%m-%d %H:%M:%S"), + ) + + +def _filter_logs_by_date_range(logs, where): + """Filter logs by startTime gte/lte from where conditions.""" + if "startTime" not in where: + return logs + date_filters = where["startTime"] + filtered = [] + for log in logs: + log_date = datetime.datetime.fromisoformat( + log["startTime"].replace("Z", "+00:00") + ) + if "gte" in date_filters: + fd = date_filters["gte"] + filter_date = ( + datetime.datetime.fromisoformat(fd.replace("Z", "+00:00")) + if "T" in fd + else datetime.datetime.strptime(fd, "%Y-%m-%d %H:%M:%S") + ) + if log_date < filter_date: + continue + if "lte" in date_filters: + fd = date_filters["lte"] + filter_date = ( + datetime.datetime.fromisoformat(fd.replace("Z", "+00:00")) + if "T" in fd + else datetime.datetime.strptime(fd, "%Y-%m-%d %H:%M:%S") + ) + if log_date > filter_date: + continue + filtered.append(log) + return filtered + + +def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=None): + """ + Create a MockPrismaClient for /spend/logs/ui endpoint tests. + + Args: + mock_spend_logs: List of mock spend log dicts. + filter_fn: Callable[[dict], list] - receives where_conditions from count(), + returns the filtered list of logs for that query. + team_lookup_fn: Optional async callable for team RBAC (find_unique). + If provided, adds litellm_teamtable to db. + """ + filtered_holder = [] + + class MockDB: + async def count(self, *args, **kwargs): + where = kwargs.get("where", {}) + filtered = filter_fn(where) + filtered_holder.clear() + filtered_holder.extend(filtered) + return len(filtered) + + async def query_raw(self, sql_query, *params): + page_size = params[-2] if len(params) >= 2 else 50 + skip = params[-1] if len(params) >= 1 else 0 + return filtered_holder[skip : skip + page_size] + + class MockPrismaClient: + def __init__(self): + self.db = MockDB() + self.db.litellm_spendlogs = self.db + if team_lookup_fn is not None: + self.db.litellm_teamtable = self + self.find_unique = team_lookup_fn + + return MockPrismaClient() from litellm.proxy._types import ( LitellmUserRoles, Member, @@ -257,7 +336,6 @@ def reset_router_callbacks(): @pytest.mark.asyncio async def test_ui_view_spend_logs_with_user_id(client, monkeypatch): - # Mock data for the test mock_spend_logs = [ { "id": "log1", @@ -281,43 +359,17 @@ async def test_ui_view_spend_logs_with_user_id(client, monkeypatch): }, ] - # Create a mock prisma client - class MockDB: - async def find_many(self, *args, **kwargs): - # Filter based on user_id in the where conditions - print("kwargs to find_many", json.dumps(kwargs, indent=4)) - if ( - "where" in kwargs - and "user" in kwargs["where"] - and kwargs["where"]["user"] == "test_user_1" - ): - return [mock_spend_logs[0]] - return mock_spend_logs + def filter_by_user(where): + if "user" in where and where["user"] == "test_user_1": + return [mock_spend_logs[0]] + return mock_spend_logs - async def count(self, *args, **kwargs): - # Return count based on user_id filter - if ( - "where" in kwargs - and "user" in kwargs["where"] - and kwargs["where"]["user"] == "test_user_1" - ): - return 1 - return len(mock_spend_logs) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_user), + ) - class MockPrismaClient: - def __init__(self): - self.db = MockDB() - self.db.litellm_spendlogs = self.db - - # Apply the monkeypatch to replace the prisma_client - mock_prisma_client = MockPrismaClient() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - # Set up test dates - start_date = ( - datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7) - ).strftime("%Y-%m-%d %H:%M:%S") - end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + start_date, end_date = _default_date_range() # Make the request with user_id filter response = client.get( @@ -431,19 +483,23 @@ async def test_ui_view_spend_logs_sort_by_and_sort_order( """Test that spend logs are returned in the correct order for each sort_by/sort_order.""" base_logs = list(_SORT_TEST_LOGS) - async def mock_find_many(*args, **kwargs): - order = kwargs.get("order", {}) - return _sort_logs(base_logs, order) - async def mock_count(*args, **kwargs): return len(base_logs) + async def mock_query_raw(sql_query, *params): + # Endpoint uses raw SQL with ORDER BY startTime DESC; mock returns sorted data + order = {"startTime": "desc"} if sort_by is None else {sort_by: sort_order or "desc"} + sorted_logs = _sort_logs(base_logs, order) + page_size = params[-2] if len(params) >= 2 else 50 + skip = params[-1] if len(params) >= 1 else 0 + return sorted_logs[skip : skip + page_size] + class MockPrismaClient: def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() - self.db.litellm_spendlogs.find_many = AsyncMock(side_effect=mock_find_many) self.db.litellm_spendlogs.count = AsyncMock(side_effect=mock_count) + self.db.query_raw = AsyncMock(side_effect=mock_query_raw) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient()) monkeypatch.setattr( @@ -539,7 +595,6 @@ async def test_ui_view_spend_logs_sort_validation_errors( @pytest.mark.asyncio async def test_ui_view_spend_logs_with_team_id(client, monkeypatch): - # Mock data for the test mock_spend_logs = [ { "id": "log1", @@ -563,54 +618,25 @@ async def test_ui_view_spend_logs_with_team_id(client, monkeypatch): }, ] - # Create a mock prisma client - class MockDB: - async def find_many(self, *args, **kwargs): - # Filter based on team_id in the where conditions - if ( - "where" in kwargs - and "team_id" in kwargs["where"] - and kwargs["where"]["team_id"] == "team1" - ): - return [mock_spend_logs[0]] - return mock_spend_logs + def filter_by_team(where): + if "team_id" in where and where["team_id"] == "team1": + return [mock_spend_logs[0]] + return mock_spend_logs - async def count(self, *args, **kwargs): - # Return count based on team_id filter - if ( - "where" in kwargs - and "team_id" in kwargs["where"] - and kwargs["where"]["team_id"] == "team1" - ): - return 1 - return len(mock_spend_logs) - - class MockPrismaClient: - def __init__(self): - self.db = MockDB() - self.db.litellm_spendlogs = self.db - - # Apply the monkeypatch - mock_prisma_client = MockPrismaClient() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - # Mock _is_admin_view_safe to return True to bypass permission checks + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_team), + ) monkeypatch.setattr( "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", - lambda user_api_key_dict: True + lambda user_api_key_dict: True, ) - - # Override auth dependency to return PROXY_ADMIN app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" ) try: - # Set up test dates - start_date = ( - datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7) - ).strftime("%Y-%m-%d %H:%M:%S") - end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + start_date, end_date = _default_date_range() # Make the request with team_id filter response = client.get( @@ -640,43 +666,26 @@ async def test_ui_view_spend_logs_internal_user_scoped_without_user_id(client, m """ Internal users should only be able to view their own spend even if user_id is not provided. """ - # Mock spend logs for 2 users mock_spend_logs = [ {"id": "log1", "request_id": "req1", "api_key": "sk-test-key", "user": "internal_user_1", "team_id": "team1", "spend": 0.05, "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-3.5-turbo"}, {"id": "log2", "request_id": "req2", "api_key": "sk-test-key", "user": "internal_user_2", "team_id": "team1", "spend": 0.10, "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-4"}, ] - # Prisma client mock that filters by "user" where condition - class MockDB: - async def find_many(self, *args, **kwargs): - where = kwargs.get("where", {}) - if "user" in where and where["user"] == "internal_user_1": - return [mock_spend_logs[0]] - return mock_spend_logs + def filter_by_user(where): + if "user" in where and where["user"] == "internal_user_1": + return [mock_spend_logs[0]] + return mock_spend_logs - async def count(self, *args, **kwargs): - where = kwargs.get("where", {}) - if "user" in where and where["user"] == "internal_user_1": - return 1 - return len(mock_spend_logs) - - class MockPrismaClient: - def __init__(self): - self.db = MockDB() - self.db.litellm_spendlogs = self.db - - mock_prisma_client = MockPrismaClient() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - # Override auth dependency to return INTERNAL_USER with specific user_id - # Override using the function reference attached to the running app module + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_user), + ) app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user_1" ) try: - start_date = (datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7)).strftime("%Y-%m-%d %H:%M:%S") - end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + start_date, end_date = _default_date_range() # No user_id provided; should auto-scope to authenticated internal user's own id response = client.get( @@ -699,55 +708,32 @@ async def test_ui_view_spend_logs_team_admin_can_view_team_spend(client, monkeyp """ Team admins should be able to view team-wide spend when team_id is provided. """ - # Mock spend logs for two teams mock_spend_logs = [ {"id": "log1", "request_id": "req1", "api_key": "sk-test-key", "user": "member1", "team_id": "team_admin_team", "spend": 0.05, "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-3.5-turbo"}, {"id": "log2", "request_id": "req2", "api_key": "sk-test-key", "user": "member2", "team_id": "team_other", "spend": 0.10, "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-4"}, ] - class MockDB: - async def find_many(self, *args, **kwargs): - where = kwargs.get("where", {}) - if "team_id" in where and where["team_id"] == "team_admin_team": - return [mock_spend_logs[0]] - return mock_spend_logs + def filter_by_team(where): + if "team_id" in where and where["team_id"] == "team_admin_team": + return [mock_spend_logs[0]] + return mock_spend_logs - async def count(self, *args, **kwargs): - where = kwargs.get("where", {}) - if "team_id" in where and where["team_id"] == "team_admin_team": - return 1 - return len(mock_spend_logs) + class TeamTable: + members_with_roles = [Member(user_id="admin_user", role="admin")] - class MockPrismaClient: - def __init__(self): - self.db = MockDB() - self.db.litellm_spendlogs = self.db - # Team lookup for RBAC check - class TeamTable: - def __init__(self): - # user "admin_user" is team admin - self.members_with_roles = [Member(user_id="admin_user", role="admin")] + async def team_lookup(where): + return TeamTable() if where == {"team_id": "team_admin_team"} else None - async def find_unique(where: dict): - if where == {"team_id": "team_admin_team"}: - return TeamTable() - return None - - self.db.litellm_teamtable = self - self.litellm_teamtable = self - self.find_unique = find_unique - - mock_prisma_client = MockPrismaClient() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - # Override auth dependency to return INTERNAL_USER (who is a team admin via team.members_with_roles) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_team, team_lookup), + ) app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="admin_user" ) try: - start_date = (datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7)).strftime("%Y-%m-%d %H:%M:%S") - end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + start_date, end_date = _default_date_range() response = client.get( "/spend/logs/ui", @@ -765,7 +751,6 @@ async def test_ui_view_spend_logs_team_admin_can_view_team_spend(client, monkeyp @pytest.mark.asyncio async def test_ui_view_spend_logs_pagination(client, monkeypatch): - # Create a larger set of mock data for pagination testing mock_spend_logs = [ { "id": f"log{i}", @@ -780,31 +765,12 @@ async def test_ui_view_spend_logs_pagination(client, monkeypatch): for i in range(1, 26) # 25 records ] - # Create a mock prisma client with pagination support - class MockDB: - async def find_many(self, *args, **kwargs): - # Handle pagination - skip = kwargs.get("skip", 0) - take = kwargs.get("take", 10) - return mock_spend_logs[skip : skip + take] + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, lambda where: mock_spend_logs), + ) - async def count(self, *args, **kwargs): - return len(mock_spend_logs) - - class MockPrismaClient: - def __init__(self): - self.db = MockDB() - self.db.litellm_spendlogs = self.db - - # Apply the monkeypatch - mock_prisma_client = MockPrismaClient() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - # Set up test dates - start_date = ( - datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7) - ).strftime("%Y-%m-%d %H:%M:%S") - end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + start_date, end_date = _default_date_range() # Test first page response = client.get( @@ -867,11 +833,11 @@ async def test_ui_view_session_spend_logs_pagination(client, monkeypatch): assert kwargs.get("where") == {"session_id": "session-123"} return len(mock_spend_logs) - async def find_many(self, *args, **kwargs): - assert kwargs.get("where") == {"session_id": "session-123"} - assert kwargs.get("order") == {"startTime": "asc"} - assert kwargs.get("skip") == 1 # page=2, page_size=1 - assert kwargs.get("take") == 1 + async def query_raw(self, sql_query, session_id, page_size, skip): + # Endpoint uses raw SQL for pagination - verify params + assert session_id == "session-123" + assert page_size == 1 + assert skip == 1 # page=2, page_size=1 return [mock_spend_logs[1]] class MockPrismaClient: @@ -900,9 +866,7 @@ async def test_ui_view_session_spend_logs_pagination(client, monkeypatch): @pytest.mark.asyncio async def test_ui_view_spend_logs_date_range_filter(client, monkeypatch): - # Create mock data with different dates today = datetime.datetime.now(timezone.utc) - mock_spend_logs = [ { "id": "log1", @@ -926,70 +890,15 @@ async def test_ui_view_spend_logs_date_range_filter(client, monkeypatch): }, ] - # Create a mock prisma client with date filtering - class MockDB: - async def find_many(self, *args, **kwargs): - # Check for date range filtering - if "where" in kwargs and "startTime" in kwargs["where"]: - date_filters = kwargs["where"]["startTime"] - filtered_logs = [] + def filter_by_date(where): + return _filter_logs_by_date_range(mock_spend_logs, where) - for log in mock_spend_logs: - log_date = datetime.datetime.fromisoformat( - log["startTime"].replace("Z", "+00:00") - ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_date), + ) - # Apply gte filter if it exists - if "gte" in date_filters: - # Handle ISO format date strings - if "T" in date_filters["gte"]: - filter_date = datetime.datetime.fromisoformat( - date_filters["gte"].replace("Z", "+00:00") - ) - else: - filter_date = datetime.datetime.strptime( - date_filters["gte"], "%Y-%m-%d %H:%M:%S" - ) - - if log_date < filter_date: - continue - - # Apply lte filter if it exists - if "lte" in date_filters: - # Handle ISO format date strings - if "T" in date_filters["lte"]: - filter_date = datetime.datetime.fromisoformat( - date_filters["lte"].replace("Z", "+00:00") - ) - else: - filter_date = datetime.datetime.strptime( - date_filters["lte"], "%Y-%m-%d %H:%M:%S" - ) - - if log_date > filter_date: - continue - - filtered_logs.append(log) - - return filtered_logs - - return mock_spend_logs - - async def count(self, *args, **kwargs): - # For simplicity, we'll just call find_many and count the results - logs = await self.find_many(*args, **kwargs) - return len(logs) - - class MockPrismaClient: - def __init__(self): - self.db = MockDB() - self.db.litellm_spendlogs = self.db - - # Apply the monkeypatch - mock_prisma_client = MockPrismaClient() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - # Test with a date range that should only include the second log + # Date range that should only include the second log (log1 is 10 days ago, log2 is 2 days ago) start_date = (today - datetime.timedelta(days=5)).strftime("%Y-%m-%d %H:%M:%S") end_date = today.strftime("%Y-%m-%d %H:%M:%S") @@ -1025,7 +934,6 @@ async def test_ui_view_spend_logs_unauthorized(client): @pytest.mark.asyncio async def test_ui_view_spend_logs_with_status(client, monkeypatch): - # Mock data for the test mock_spend_logs = [ { "id": "log1", @@ -1051,49 +959,19 @@ async def test_ui_view_spend_logs_with_status(client, monkeypatch): }, ] - # Create a mock prisma client - class MockDB: - async def find_many(self, *args, **kwargs): - # Filter based on status in the where conditions - if "where" in kwargs: - where_conditions = kwargs["where"] - if "OR" in where_conditions: - # Handle success case (which includes None status) - return [mock_spend_logs[0]] - elif ( - "status" in where_conditions - and where_conditions["status"]["equals"] == "failure" - ): - return [mock_spend_logs[1]] - return mock_spend_logs + def filter_by_status(where): + if "OR" in where: + return [mock_spend_logs[0]] # success + if "status" in where and where["status"].get("equals") == "failure": + return [mock_spend_logs[1]] + return mock_spend_logs - async def count(self, *args, **kwargs): - # Return count based on status filter - if "where" in kwargs: - where_conditions = kwargs["where"] - if "OR" in where_conditions: - return 1 - elif ( - "status" in where_conditions - and where_conditions["status"]["equals"] == "failure" - ): - return 1 - return len(mock_spend_logs) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_status), + ) - class MockPrismaClient: - def __init__(self): - self.db = MockDB() - self.db.litellm_spendlogs = self.db - - # Apply the monkeypatch - mock_prisma_client = MockPrismaClient() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - # Set up test dates - start_date = ( - datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7) - ).strftime("%Y-%m-%d %H:%M:%S") - end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + start_date, end_date = _default_date_range() # Test success status response = client.get( @@ -1132,7 +1010,6 @@ async def test_ui_view_spend_logs_with_status(client, monkeypatch): @pytest.mark.asyncio async def test_ui_view_spend_logs_with_model(client, monkeypatch): - # Mock data for the test mock_spend_logs = [ { "id": "log1", @@ -1158,42 +1035,17 @@ async def test_ui_view_spend_logs_with_model(client, monkeypatch): }, ] - # Create a mock prisma client - class MockDB: - async def find_many(self, *args, **kwargs): - # Filter based on model in the where conditions - if ( - "where" in kwargs - and "model" in kwargs["where"] - and kwargs["where"]["model"] == "gpt-3.5-turbo" - ): - return [mock_spend_logs[0]] - return mock_spend_logs + def filter_by_model(where): + if "model" in where and where["model"] == "gpt-3.5-turbo": + return [mock_spend_logs[0]] + return mock_spend_logs - async def count(self, *args, **kwargs): - # Return count based on model filter - if ( - "where" in kwargs - and "model" in kwargs["where"] - and kwargs["where"]["model"] == "gpt-3.5-turbo" - ): - return 1 - return len(mock_spend_logs) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_model), + ) - class MockPrismaClient: - def __init__(self): - self.db = MockDB() - self.db.litellm_spendlogs = self.db - - # Apply the monkeypatch - mock_prisma_client = MockPrismaClient() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - # Set up test dates - start_date = ( - datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7) - ).strftime("%Y-%m-%d %H:%M:%S") - end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + start_date, end_date = _default_date_range() # Make the request with model filter response = client.get( @@ -1246,37 +1098,17 @@ async def test_ui_view_spend_logs_with_model_id(client, monkeypatch): }, ] - class MockDB: - async def find_many(self, *args, **kwargs): - if ( - "where" in kwargs - and "model_id" in kwargs["where"] - and kwargs["where"]["model_id"] == "deployment-id-1" - ): - return [mock_spend_logs[0]] - return mock_spend_logs + def filter_by_model_id(where): + if "model_id" in where and where["model_id"] == "deployment-id-1": + return [mock_spend_logs[0]] + return mock_spend_logs - async def count(self, *args, **kwargs): - if ( - "where" in kwargs - and "model_id" in kwargs["where"] - and kwargs["where"]["model_id"] == "deployment-id-1" - ): - return 1 - return len(mock_spend_logs) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_model_id), + ) - class MockPrismaClient: - def __init__(self): - self.db = MockDB() - self.db.litellm_spendlogs = self.db - - mock_prisma_client = MockPrismaClient() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - start_date = ( - datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7) - ).strftime("%Y-%m-%d %H:%M:%S") - end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + start_date, end_date = _default_date_range() response = client.get( "/spend/logs/ui", @@ -1297,7 +1129,6 @@ async def test_ui_view_spend_logs_with_model_id(client, monkeypatch): @pytest.mark.asyncio async def test_ui_view_spend_logs_with_key_hash(client, monkeypatch): - # Mock data for the test mock_spend_logs = [ { "id": "log1", @@ -1321,42 +1152,17 @@ async def test_ui_view_spend_logs_with_key_hash(client, monkeypatch): }, ] - # Create a mock prisma client - class MockDB: - async def find_many(self, *args, **kwargs): - # Filter based on key_hash in the where conditions - if ( - "where" in kwargs - and "api_key" in kwargs["where"] - and kwargs["where"]["api_key"] == "sk-test-key-1" - ): - return [mock_spend_logs[0]] - return mock_spend_logs + def filter_by_api_key(where): + if "api_key" in where and where["api_key"] == "sk-test-key-1": + return [mock_spend_logs[0]] + return mock_spend_logs - async def count(self, *args, **kwargs): - # Return count based on key_hash filter - if ( - "where" in kwargs - and "api_key" in kwargs["where"] - and kwargs["where"]["api_key"] == "sk-test-key-1" - ): - return 1 - return len(mock_spend_logs) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_api_key), + ) - class MockPrismaClient: - def __init__(self): - self.db = MockDB() - self.db.litellm_spendlogs = self.db - - # Apply the monkeypatch - mock_prisma_client = MockPrismaClient() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - # Set up test dates - start_date = ( - datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7) - ).strftime("%Y-%m-%d %H:%M:%S") - end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + start_date, end_date = _default_date_range() # Make the request with key_hash filter response = client.get( @@ -2163,45 +1969,21 @@ async def test_ui_view_spend_logs_with_error_code(client): }, ] - with patch.object(ps, "prisma_client") as mock_prisma: - # Mock the find_many method to return filtered results - async def mock_find_many(*args, **kwargs): - where_conditions = kwargs.get("where", {}) - if "metadata" in where_conditions: - metadata_filter = where_conditions["metadata"] - if metadata_filter.get("path") == ["error_information", "error_code"]: - error_code = metadata_filter.get("equals") - # Handle both string and integer error codes - # The endpoint wraps error_code in quotes, so strip them for comparison - error_code_value = str(error_code).strip('"') - if error_code_value == "404": - return [mock_spend_logs[0]] - elif error_code_value == "500": - return [mock_spend_logs[1]] - return mock_spend_logs + def filter_by_error_code(where): + if "metadata" in where: + mf = where["metadata"] + if mf.get("path") == ["error_information", "error_code"]: + code = str(mf.get("equals", "")).strip('"') + if code == "404": + return [mock_spend_logs[0]] + if code == "500": + return [mock_spend_logs[1]] + return mock_spend_logs - async def mock_count(*args, **kwargs): - where_conditions = kwargs.get("where", {}) - if "metadata" in where_conditions: - metadata_filter = where_conditions["metadata"] - if metadata_filter.get("path") == ["error_information", "error_code"]: - error_code = metadata_filter.get("equals") - # Handle both string and integer error codes - # The endpoint wraps error_code in quotes, so strip them for comparison - error_code_value = str(error_code).strip('"') - if error_code_value == "404": - return 1 - elif error_code_value == "500": - return 1 - return len(mock_spend_logs) - - mock_prisma.db.litellm_spendlogs.find_many = mock_find_many - mock_prisma.db.litellm_spendlogs.count = mock_count - - start_date = ( - datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7) - ).strftime("%Y-%m-%d %H:%M:%S") - end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + with patch.object( + ps, "prisma_client", make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_error_code) + ): + start_date, end_date = _default_date_range() response = client.get( "/spend/logs/ui", @@ -2251,40 +2033,21 @@ async def test_ui_view_spend_logs_with_error_message(client): }, ] - with patch.object(ps, "prisma_client") as mock_prisma: - # Mock the find_many method to return filtered results - async def mock_find_many(*args, **kwargs): - where_conditions = kwargs.get("where", {}) - if "metadata" in where_conditions: - metadata_filter = where_conditions["metadata"] - if metadata_filter.get("path") == ["error_information", "error_message"]: - error_message_filter = metadata_filter.get("string_contains") - # Check if the error message contains the filter string - if error_message_filter == "Rate limit": - return [mock_spend_logs[0]] - elif error_message_filter == "Invalid API": - return [mock_spend_logs[1]] - return mock_spend_logs + def filter_by_error_message(where): + if "metadata" in where: + mf = where["metadata"] + if mf.get("path") == ["error_information", "error_message"]: + msg = mf.get("string_contains") + if msg == "Rate limit": + return [mock_spend_logs[0]] + if msg == "Invalid API": + return [mock_spend_logs[1]] + return mock_spend_logs - async def mock_count(*args, **kwargs): - where_conditions = kwargs.get("where", {}) - if "metadata" in where_conditions: - metadata_filter = where_conditions["metadata"] - if metadata_filter.get("path") == ["error_information", "error_message"]: - error_message_filter = metadata_filter.get("string_contains") - if error_message_filter == "Rate limit": - return 1 - elif error_message_filter == "Invalid API": - return 1 - return len(mock_spend_logs) - - mock_prisma.db.litellm_spendlogs.find_many = mock_find_many - mock_prisma.db.litellm_spendlogs.count = mock_count - - start_date = ( - datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7) - ).strftime("%Y-%m-%d %H:%M:%S") - end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + with patch.object( + ps, "prisma_client", make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_error_message) + ): + start_date, end_date = _default_date_range() response = client.get( "/spend/logs/ui", @@ -2345,55 +2108,26 @@ async def test_ui_view_spend_logs_with_error_code_and_key_alias(client): }, ] - with patch.object(ps, "prisma_client") as mock_prisma: - # Mock the find_many method to handle AND conditions - async def mock_find_many(*args, **kwargs): - where_conditions = kwargs.get("where", {}) - if "AND" in where_conditions: - key_alias_filter = None - error_code_filter = None - for condition in where_conditions["AND"]: - if "metadata" in condition: - metadata_filter = condition["metadata"] - if metadata_filter.get("path") == ["user_api_key_alias"]: - key_alias_filter = metadata_filter.get("string_contains") - elif metadata_filter.get("path") == ["error_information", "error_code"]: - error_code_filter = metadata_filter.get("equals") + def filter_by_error_code_and_key_alias(where): + if "AND" in where: + key_alias = error_code = None + for cond in where["AND"]: + if "metadata" in cond: + mf = cond["metadata"] + if mf.get("path") == ["user_api_key_alias"]: + key_alias = mf.get("string_contains") + elif mf.get("path") == ["error_information", "error_code"]: + error_code = str(mf.get("equals", "")).strip('"') + if key_alias == "test-key-1" and error_code == "500": + return [mock_spend_logs[2]] + return mock_spend_logs - # Handle both string and integer error codes - # The endpoint wraps error_code in quotes, so strip them for comparison - error_code_value = str(error_code_filter).strip('"') - if key_alias_filter == "test-key-1" and error_code_value == "500": - return [mock_spend_logs[2]] # Only log3 matches both conditions - return mock_spend_logs - - async def mock_count(*args, **kwargs): - where_conditions = kwargs.get("where", {}) - if "AND" in where_conditions: - key_alias_filter = None - error_code_filter = None - for condition in where_conditions["AND"]: - if "metadata" in condition: - metadata_filter = condition["metadata"] - if metadata_filter.get("path") == ["user_api_key_alias"]: - key_alias_filter = metadata_filter.get("string_contains") - elif metadata_filter.get("path") == ["error_information", "error_code"]: - error_code_filter = metadata_filter.get("equals") - - # Handle both string and integer error codes - # The endpoint wraps error_code in quotes, so strip them for comparison - error_code_value = str(error_code_filter).strip('"') - if key_alias_filter == "test-key-1" and error_code_value == "500": - return 1 - return len(mock_spend_logs) - - mock_prisma.db.litellm_spendlogs.find_many = mock_find_many - mock_prisma.db.litellm_spendlogs.count = mock_count - - start_date = ( - datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7) - ).strftime("%Y-%m-%d %H:%M:%S") - end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + with patch.object( + ps, + "prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_error_code_and_key_alias), + ): + start_date, end_date = _default_date_range() response = client.get( "/spend/logs/ui",