From 48bc5ccb4f003dee859aec4c6721ded4d89c0426 Mon Sep 17 00:00:00 2001
From: Yuta Saito <uc4w6c@bma.biglobe.ne.jp>
Date: Fri, 9 Jan 2026 11:23:47 +0900
Subject: [PATCH] fix: test

---
 .../integrations/cloudzero/test_cloudzero.py  | 79 ++++++++++++++++---
 ...database.py => test_cloudzero_database.py} |  0
 ...est_database.py => test_focus_database.py} |  0
 3 files changed, 70 insertions(+), 9 deletions(-)
 rename tests/test_litellm/integrations/cloudzero/{test_database.py => test_cloudzero_database.py} (100%)
 rename tests/test_litellm/integrations/focus/{test_database.py => test_focus_database.py} (100%)

diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py
index 31a2f6cbf5..b0aac17e7d 100644
--- a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py
+++ b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py
@@ -66,16 +66,77 @@ class TestCloudZeroHourlyExport:
             fake_client = MagicMock()
             fake_db = MagicMock()
 
-            async def query_raw_mock(query: str):
-                sql_context = pl.SQLContext(
-                    LiteLLM_DailyUserSpend=spend_mock_data,
-                    LiteLLM_VerificationToken=verification_mock_data,
-                    LiteLLM_TeamTable=team_mock_data,
-                    LiteLLM_UserTable=user_mock_data,
-                )
-                result = sql_context.execute(query).collect()
+            async def query_raw_mock(query: str, *params):
+                start_time_utc = params[0] if len(params) > 0 else None
+                end_time_utc = params[1] if len(params) > 1 else None
+                limit = params[2] if len(params) > 2 else None
 
-                return result
+                spend_df = spend_mock_data.collect()
+                verification_df = verification_mock_data.collect().rename(
+                    {"key_alias": "api_key_alias"}
+                )
+                team_df = team_mock_data.collect()
+                user_df = user_mock_data.collect()
+
+                joined = (
+                    spend_df.join(
+                        verification_df, left_on="api_key", right_on="token", how="left"
+                    )
+                    .join(
+                        team_df,
+                        left_on="team_id",
+                        right_on="team_id",
+                        how="left",
+                        suffix="_team",
+                    )
+                    .join(
+                        user_df,
+                        left_on="user_id",
+                        right_on="user_id",
+                        how="left",
+                        suffix="_user",
+                    )
+                )
+
+                for duplicate_column in ("team_id_team", "user_id_user"):
+                    if duplicate_column in joined.columns:
+                        joined = joined.drop(duplicate_column)
+
+                if start_time_utc is not None:
+                    joined = joined.filter(pl.col("updated_at") >= start_time_utc)
+                if end_time_utc is not None:
+                    joined = joined.filter(pl.col("updated_at") <= end_time_utc)
+
+                joined = joined.select(
+                    [
+                        "id",
+                        "date",
+                        "user_id",
+                        "api_key",
+                        "model",
+                        "model_group",
+                        "custom_llm_provider",
+                        "prompt_tokens",
+                        "completion_tokens",
+                        "spend",
+                        "api_requests",
+                        "successful_requests",
+                        "failed_requests",
+                        "cache_creation_input_tokens",
+                        "cache_read_input_tokens",
+                        "created_at",
+                        "updated_at",
+                        "team_id",
+                        "api_key_alias",
+                        "team_alias",
+                        "user_email",
+                    ]
+                ).sort(["date", "created_at"], descending=[True, True])
+
+                if limit is not None:
+                    joined = joined.head(int(limit))
+
+                return joined
 
             fake_db.query_raw = AsyncMock(side_effect=query_raw_mock)
             fake_client.db = fake_db
diff --git a/tests/test_litellm/integrations/cloudzero/test_database.py b/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py
similarity index 100%
rename from tests/test_litellm/integrations/cloudzero/test_database.py
rename to tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py
diff --git a/tests/test_litellm/integrations/focus/test_database.py b/tests/test_litellm/integrations/focus/test_focus_database.py
similarity index 100%
rename from tests/test_litellm/integrations/focus/test_database.py
rename to tests/test_litellm/integrations/focus/test_focus_database.py