From 0d329826f10da90a6b8205ac58ce0344ef82ee9c Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 27 Nov 2025 17:40:38 -0800
Subject: [PATCH] Fix flaky tests

---
 .../src/components/entity_usage.test.tsx               | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/entity_usage.test.tsx b/ui/litellm-dashboard/src/components/entity_usage.test.tsx
index d5cc503fd10..17016b6479d 100644
--- a/ui/litellm-dashboard/src/components/entity_usage.test.tsx
+++ b/ui/litellm-dashboard/src/components/entity_usage.test.tsx
@@ -168,16 +168,18 @@ describe("EntityUsage", () => {
   });
 
   it("should render with organization entity type and call organization API", async () => {
-    const { getByText, getAllByText } = render(<EntityUsage {...defaultProps} entityType="organization" />);
+    render(<EntityUsage {...defaultProps} entityType="organization" />);
 
     await waitFor(() => {
       expect(mockOrganizationDailyActivityCall).toHaveBeenCalled();
     });
 
-    expect(getByText("Organization Spend Overview")).toBeInTheDocument();
+    expect(screen.getByText("Organization Spend Overview")).toBeInTheDocument();
 
-    const spendElements = getAllByText("$100.50");
-    expect(spendElements.length).toBeGreaterThan(0);
+    await waitFor(() => {
+      const spendElements = screen.getAllByText("$100.50");
+      expect(spendElements.length).toBeGreaterThan(0);
+    });
   });
 
   it("should switch between tabs", async () => {