From cfcd597b9124183e712a8e045f7d8a7cb51584a4 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Sat, 22 Nov 2025 09:44:58 -0800
Subject: [PATCH] Fix tests (#16972)

---
 .../src/components/entity_usage.test.tsx      | 45 ++++++++++---------
 1 file changed, 24 insertions(+), 21 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/entity_usage.test.tsx b/ui/litellm-dashboard/src/components/entity_usage.test.tsx
index 7e2a2d4a23..41fa6ce33e 100644
--- a/ui/litellm-dashboard/src/components/entity_usage.test.tsx
+++ b/ui/litellm-dashboard/src/components/entity_usage.test.tsx
@@ -1,5 +1,5 @@
-import { render, screen, fireEvent, waitFor } from "@testing-library/react";
-import { describe, it, expect, vi, beforeEach, beforeAll } from "vitest";
+import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";
+import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 import EntityUsage from "./entity_usage";
 import * as networking from "./networking";
 
@@ -140,13 +140,13 @@ describe("EntityUsage", () => {
       expect(mockTagDailyActivityCall).toHaveBeenCalled();
     });
 
-    // Check that spend metrics are displayed
     expect(screen.getByText("Tag Spend Overview")).toBeInTheDocument();
     expect(screen.getByText("Total Spend")).toBeInTheDocument();
 
-    // Use getAllByText since $100.50 appears in multiple places
-    const spendElements = screen.getAllByText("$100.50");
-    expect(spendElements.length).toBeGreaterThan(0);
+    await waitFor(() => {
+      const spendElements = screen.getAllByText("$100.50");
+      expect(spendElements.length).toBeGreaterThan(0);
+    });
 
     expect(screen.getByText("1,000")).toBeInTheDocument(); // Total Requests
   });
@@ -161,9 +161,10 @@ describe("EntityUsage", () => {
     // Check that it shows team-specific label
     expect(screen.getByText("Team Spend Overview")).toBeInTheDocument();
 
-    // Use getAllByText since $100.50 appears in multiple places
-    const spendElements = screen.getAllByText("$100.50");
-    expect(spendElements.length).toBeGreaterThan(0);
+    await waitFor(() => {
+      const spendElements = screen.getAllByText("$100.50");
+      expect(spendElements.length).toBeGreaterThan(0);
+    });
   });
 
   it("should render with organization entity type and call organization API", async () => {
@@ -175,8 +176,10 @@ describe("EntityUsage", () => {
 
     expect(getByText("Organization Spend Overview")).toBeInTheDocument();
 
-    const spendElements = getAllByText("$100.50");
-    expect(spendElements.length).toBeGreaterThan(0);
+    await waitFor(() => {
+      const spendElements = getAllByText("$100.50");
+      expect(spendElements.length).toBeGreaterThan(0);
+    });
   });
 
   it("should switch between tabs", async () => {
@@ -186,21 +189,20 @@ describe("EntityUsage", () => {
       expect(mockTagDailyActivityCall).toHaveBeenCalled();
     });
 
-    // Check default tab (Cost) is shown
     expect(screen.getByText("Tag Spend Overview")).toBeInTheDocument();
 
-    // Click Model Activity tab
     const modelActivityTab = screen.getByText("Model Activity");
-    fireEvent.click(modelActivityTab);
+    act(() => {
+      fireEvent.click(modelActivityTab);
+    });
 
-    // Should show activity metrics
     expect(screen.getAllByText("Activity Metrics")[0]).toBeInTheDocument();
 
-    // Click Key Activity tab
     const keyActivityTab = screen.getByText("Key Activity");
-    fireEvent.click(keyActivityTab);
+    act(() => {
+      fireEvent.click(keyActivityTab);
+    });
 
-    // Should show activity metrics again
     expect(screen.getAllByText("Activity Metrics")[1]).toBeInTheDocument();
   });
 
@@ -224,8 +226,9 @@ describe("EntityUsage", () => {
       expect(mockTagDailyActivityCall).toHaveBeenCalled();
     });
 
-    // Check that zero values are displayed (component formats it as $0.00)
-    expect(screen.getByText("$0.00")).toBeInTheDocument();
-    expect(screen.getByText("Total Spend")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("$0.00")).toBeInTheDocument();
+      expect(screen.getByText("Total Spend")).toBeInTheDocument();
+    });
   });
 });