From 3b13a5fda968620d699d5f8cc15d7485f118bc5e Mon Sep 17 00:00:00 2001
From: ryan <ryan@berri.ai>
Date: Thu, 3 Sep 2026 18:31:43 +0000
Subject: [PATCH] test(ui): query the tokens cell by role instead of walking
 the DOM

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../RequestLogsTableColumns.test.tsx          | 56 +++++++++----------
 1 file changed, 25 insertions(+), 31 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx
index d2c84173fd1..3c9e6543c1c 100644
--- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx
@@ -1,4 +1,4 @@
-import { render, screen, within } from "@testing-library/react";
+import { render, screen } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { describe, expect, it, vi } from "vitest";
 
@@ -76,43 +76,37 @@ describe("Cost column", () => {
 });
 
 describe("Tokens column", () => {
-  it("shows the summed session token usage, not the representative call's tokens, for a multi-round session", () => {
-    renderRows([
-      logEntry({
-        request_id: "req-session-tokens",
-        total_tokens: 10,
-        prompt_tokens: 7,
-        completion_tokens: 3,
-        session_id: "sess-1",
-        session_total_count: 3,
-        session_total_tokens: 60,
-        session_total_prompt_tokens: 42,
-        session_total_completion_tokens: 18,
-      }),
-    ]);
+  const sessionRow: Partial<LogEntry> = {
+    request_id: "req-session-tokens",
+    total_tokens: 10,
+    prompt_tokens: 7,
+    completion_tokens: 3,
+    session_id: "sess-1",
+    session_total_count: 3,
+  };
 
-    const tokensCell = screen.getByText("60").closest("td")!;
-    expect(within(tokensCell).getByText("(42+18)")).toBeInTheDocument();
-    expect(within(tokensCell).getByText("session total")).toBeInTheDocument();
+  it("shows the summed session token usage, not the representative call's tokens, for a multi-round session", () => {
+    const aggregatedRow: Partial<LogEntry> = {
+      ...sessionRow,
+      session_total_tokens: 60,
+      session_total_prompt_tokens: 42,
+      session_total_completion_tokens: 18,
+    };
+    renderRows([logEntry(aggregatedRow)]);
+
+    const tokensCell = screen.getByRole("cell", { name: /\(42\+18\)/ });
+    expect(tokensCell).toHaveTextContent("60");
+    expect(tokensCell).toHaveTextContent("session total");
     expect(screen.queryByText("10")).not.toBeInTheDocument();
     expect(screen.queryByText("(7+3)")).not.toBeInTheDocument();
   });
 
   it("falls back to the call's own tokens with no session label when the backend sent no session token sums", () => {
-    renderRows([
-      logEntry({
-        request_id: "req-no-token-aggregate",
-        total_tokens: 10,
-        prompt_tokens: 7,
-        completion_tokens: 3,
-        session_id: "sess-2",
-        session_total_count: 3,
-      }),
-    ]);
+    renderRows([logEntry(sessionRow)]);
 
-    const tokensCell = screen.getByText("10").closest("td")!;
-    expect(within(tokensCell).getByText("(7+3)")).toBeInTheDocument();
-    expect(within(tokensCell).queryByText("session total")).not.toBeInTheDocument();
+    const tokensCell = screen.getByRole("cell", { name: /\(7\+3\)/ });
+    expect(tokensCell).toHaveTextContent("10");
+    expect(tokensCell).not.toHaveTextContent("session total");
   });
 });