fix: test

2026-08-06 20:25:29 +00:00 · 2026-01-09 11:23:47 +09:00
parent f129f598a0
commit 48bc5ccb4f
3 changed files with 70 additions and 9 deletions
@@ -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