From 11f78d80c7a1c0c283879fd6dedb9296721dab65 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 3 Jun 2026 17:30:48 +0000
Subject: [PATCH] test: align prometheus user-budget + logs quick-select tests
 with merged code

The merge into this branch left two test patterns out of step with the code
they exercise.

test_set_user_budget_metrics_includes_user_email_and_alias_labels_when_opted_in
flipped litellm.prometheus_user_budget_label_include_email_alias after the
fixture had already built the PrometheusLogger. get_labels_for_metric now
snapshots each metric's label set at construction time, so the runtime flip
no longer reached the cached labels. Enable the flag before constructing the
logger, matching how the proxy applies config at startup.

view_logs/index.test.tsx referenced uiSpendLogsCall and moment without
importing them, and the merged index.tsx now fetches through
useLogFilterLogic (the hook the file stubs out) rather than calling
uiSpendLogsCall directly. Add the imports and restore the real hook for the
Quick Select window assertions so the call is actually observed.
---
 .../test_prometheus_user_team_metrics.py      | 39 +++++++++++--------
 .../src/components/view_logs/index.test.tsx   |  9 +++++
 2 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py
index 12f30ab6024..361ab7332f8 100644
--- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py
+++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py
@@ -511,29 +511,34 @@ def test_set_user_budget_metrics_default_no_email_alias_labels(
     )
 
 
-def test_set_user_budget_metrics_includes_user_email_and_alias_labels_when_opted_in(
-    prometheus_logger,
-):
-    """When prometheus_user_budget_label_include_email_alias=True, email+alias labels appear."""
+def test_set_user_budget_metrics_includes_user_email_and_alias_labels_when_opted_in():
+    """When prometheus_user_budget_label_include_email_alias=True, email+alias labels appear.
+
+    The flag is read once per metric at logger construction time and snapshotted,
+    so it must be enabled before the PrometheusLogger is built (mirroring how the
+    proxy applies config at startup before instantiating callbacks).
+    """
     import litellm
     from litellm.proxy._types import LiteLLM_UserTable
 
     litellm.prometheus_user_budget_label_include_email_alias = True
 
-    user = LiteLLM_UserTable(
-        user_id="user-abc-123",
-        user_email="alice@example.com",
-        user_alias="Alice",
-        spend=25.0,
-        max_budget=100.0,
-        budget_reset_at=datetime(2026, 3, 1, tzinfo=timezone.utc),
-    )
-
-    prometheus_logger.litellm_remaining_user_budget_metric = MagicMock()
-    prometheus_logger.litellm_user_max_budget_metric = MagicMock()
-    prometheus_logger.litellm_user_budget_remaining_hours_metric = MagicMock()
-
     try:
+        prometheus_logger = PrometheusLogger()
+
+        user = LiteLLM_UserTable(
+            user_id="user-abc-123",
+            user_email="alice@example.com",
+            user_alias="Alice",
+            spend=25.0,
+            max_budget=100.0,
+            budget_reset_at=datetime(2026, 3, 1, tzinfo=timezone.utc),
+        )
+
+        prometheus_logger.litellm_remaining_user_budget_metric = MagicMock()
+        prometheus_logger.litellm_user_max_budget_metric = MagicMock()
+        prometheus_logger.litellm_user_budget_remaining_hours_metric = MagicMock()
+
         prometheus_logger._set_user_budget_metrics(user)
 
         prometheus_logger.litellm_remaining_user_budget_metric.labels.assert_called_once_with(
diff --git a/ui/litellm-dashboard/src/components/view_logs/index.test.tsx b/ui/litellm-dashboard/src/components/view_logs/index.test.tsx
index 55676c2b05b..6f1c4126282 100644
--- a/ui/litellm-dashboard/src/components/view_logs/index.test.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/index.test.tsx
@@ -1,8 +1,11 @@
 import { screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
+import moment from "moment";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import SpendLogsTable from "./index";
 import { renderWithProviders } from "../../../tests/test-utils";
+import { uiSpendLogsCall } from "../networking";
+import { useLogFilterLogic } from "./log_filter_logic";
 
 const mockHandleFilterResetFromHook = vi.fn();
 vi.mock("./log_filter_logic", async (importOriginal) => {
@@ -117,6 +120,12 @@ describe("SpendLogsTable", () => {
   });
 
   describe("Quick Select time range", () => {
+    // uiSpendLogsCall fires from the real useLogFilterLogic query, so restore it here.
+    beforeEach(async () => {
+      const actual = await vi.importActual<typeof import("./log_filter_logic")>("./log_filter_logic");
+      vi.mocked(useLogFilterLogic).mockImplementation(actual.useLogFilterLogic);
+    });
+
     const waitForWindowSeconds = async (minMinutes: number) => {
       let diff = -1;
       await waitFor(() => {