From 1037dc18b9f33221b4652093e661dbe59786e91a Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 12 Dec 2025 12:23:11 -0800
Subject: [PATCH] Adding test

---
 .../EntityUsage/EntityUsage.test.tsx          | 75 +++++++++++++++++++
 1 file changed, 75 insertions(+)

diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx
index 2ace90e0122..afd5c5789b8 100644
--- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx
+++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx
@@ -269,4 +269,79 @@ describe("EntityUsage", () => {
     expect(screen.getByText("Total Spend")).toBeInTheDocument();
     expect(screen.getAllByText("0")[0]).toBeInTheDocument();
   });
+
+  it("should use entityList label when entityList is provided and entity exists", async () => {
+    const customEntityList = [
+      { label: "Custom Tag Label", value: "tag-1" },
+      { label: "Tag 2", value: "tag-2" },
+    ];
+
+    render(<EntityUsage {...defaultProps} entityList={customEntityList} />);
+
+    await waitFor(() => {
+      expect(mockTagDailyActivityCall).toHaveBeenCalled();
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Custom Tag Label")).toBeInTheDocument();
+    });
+  });
+
+  it("should fallback to team_alias when entityList is provided but entity does not exist", async () => {
+    const customEntityList = [{ label: "Tag 2", value: "tag-2" }];
+
+    render(<EntityUsage {...defaultProps} entityList={customEntityList} />);
+
+    await waitFor(() => {
+      expect(mockTagDailyActivityCall).toHaveBeenCalled();
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Tag 1")).toBeInTheDocument();
+    });
+  });
+
+  it("should fallback to team_alias when entityList is null", async () => {
+    render(<EntityUsage {...defaultProps} entityList={null} />);
+
+    await waitFor(() => {
+      expect(mockTagDailyActivityCall).toHaveBeenCalled();
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Tag 1")).toBeInTheDocument();
+    });
+  });
+
+  it("should fallback to entity value when no entityList and no team_alias", async () => {
+    const spendDataWithoutAlias = {
+      ...mockSpendData,
+      results: [
+        {
+          ...mockSpendData.results[0],
+          breakdown: {
+            ...mockSpendData.results[0].breakdown,
+            entities: {
+              "tag-1": {
+                ...mockSpendData.results[0].breakdown.entities["tag-1"],
+                metadata: {},
+              },
+            },
+          },
+        },
+      ],
+    };
+
+    mockTagDailyActivityCall.mockResolvedValue(spendDataWithoutAlias);
+
+    render(<EntityUsage {...defaultProps} entityList={null} />);
+
+    await waitFor(() => {
+      expect(mockTagDailyActivityCall).toHaveBeenCalled();
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("tag-1")).toBeInTheDocument();
+    });
+  });
 });