From 063b14582b08b6fc307eb2a51d4922aecf25e573 Mon Sep 17 00:00:00 2001
From: Dexploarer <211557447+Dexploarer@users.noreply.github.com>
Date: Mon, 27 Apr 2026 00:49:18 +0000
Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=AA=20[Test]=20Add=20unit=20tests=20fo?=
 =?UTF-8?q?r=20core=20TUI=20UI=20components?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/tui/components/assistant-message.test.ts | 143 +++++++++++++++
 src/tui/components/tool-execution.test.ts    | 181 +++++++++++++++++++
 src/tui/components/user-message.test.ts      |  73 ++++++++
 3 files changed, 397 insertions(+)
 create mode 100644 src/tui/components/assistant-message.test.ts
 create mode 100644 src/tui/components/tool-execution.test.ts
 create mode 100644 src/tui/components/user-message.test.ts

diff --git a/src/tui/components/assistant-message.test.ts b/src/tui/components/assistant-message.test.ts
new file mode 100644
index 0000000000..339b08d395
--- /dev/null
+++ b/src/tui/components/assistant-message.test.ts
@@ -0,0 +1,143 @@
+import { describe, expect, it, vi, beforeEach } from "vitest";
+import { AssistantMessageComponent } from "./assistant-message.js";
+
+let markdownInstances: any[] = [];
+let imageInstances: any[] = [];
+
+vi.mock("@mariozechner/pi-tui", () => {
+  return {
+    Markdown: class Markdown {
+      constructor(public text: string, top: number, left: number, theme: any, options: any) {
+        markdownInstances.push(this);
+      }
+      setText = vi.fn((text) => { this.text = text; });
+      render = vi.fn().mockImplementation((w) => this.text ? [this.text] : []);
+      invalidate = vi.fn();
+    },
+    Image: class Image {
+      constructor(base64: string, mimeType: string, theme: any, options: any) {
+        imageInstances.push(this);
+      }
+      render = vi.fn().mockReturnValue(["[Image]"]);
+      invalidate = vi.fn();
+    }
+  };
+});
+
+vi.mock("../theme.js", () => {
+  return {
+    miladyMarkdownTheme: { mockTheme: true },
+    tuiTheme: {
+      muted: vi.fn((t) => `muted:${t}`),
+      dim: vi.fn((t) => `dim:${t}`)
+    }
+  };
+});
+
+describe("AssistantMessageComponent", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    markdownInstances = [];
+    imageInstances = [];
+  });
+
+  it("initializes empty and streaming", () => {
+    const component = new AssistantMessageComponent();
+    expect(component).toBeDefined();
+
+    // Empty output initially
+    const rendered = component.render(100);
+    expect(rendered).toEqual([]);
+  });
+
+  it("updates content and renders with streaming cursor", () => {
+    const component = new AssistantMessageComponent();
+
+    component.updateContent("Hello");
+
+    // The main markdown should be instantiated on construction
+    const md = markdownInstances[0];
+    expect(md.setText).toHaveBeenCalledWith("Hello ▊");
+
+    const rendered = component.render(100);
+    expect(rendered).toEqual(["", "Hello ▊"]);
+  });
+
+  it("finalizes and removes streaming cursor", () => {
+    const component = new AssistantMessageComponent();
+
+    component.updateContent("Hello");
+    component.finalize();
+
+    const md = markdownInstances[0];
+    expect(md.setText).toHaveBeenCalledWith("Hello");
+
+    const rendered = component.render(100);
+    expect(rendered).toEqual(["", "Hello"]);
+  });
+
+  it("shows thinking text when enabled", () => {
+    const component = new AssistantMessageComponent(true);
+
+    // updateThinking triggers a rebuild which creates a new Markdown instance
+    component.updateThinking("Hmm...");
+    component.updateContent("Yes");
+    component.finalize();
+
+    // We expect multiple instances due to rebuilds, but we care about the render output
+    const rendered = component.render(100);
+    expect(rendered).toEqual([
+      "",
+      "Hmm...",
+      "",
+      "Yes"
+    ]);
+  });
+
+  it("ignores thinking text when disabled", () => {
+    const component = new AssistantMessageComponent(false);
+
+    component.updateThinking("Hmm...");
+    component.updateContent("Yes");
+    component.finalize();
+
+    const rendered = component.render(100);
+    expect(rendered).toEqual(["", "Yes"]);
+  });
+
+  it("adds and renders images", () => {
+    const component = new AssistantMessageComponent();
+
+    component.updateContent("Look at this");
+    component.finalize();
+
+    component.addImage({
+      base64: "dGVzdA==",
+      mimeType: "image/png",
+      filename: "test.png"
+    });
+
+    expect(imageInstances.length).toBe(1);
+
+    const rendered = component.render(100);
+    expect(rendered).toEqual([
+      "",
+      "Look at this",
+      "", // spacer
+      "[Image]"
+    ]);
+  });
+
+  it("invalidates all sub-components", () => {
+    const component = new AssistantMessageComponent(true);
+    component.updateThinking("Hmm");
+    component.addImage({ base64: "dGVzdA==", mimeType: "image/png" });
+
+    component.invalidate();
+
+    // Check that at least some invalidates were called
+    const someMarkdownInvalidated = markdownInstances.some(md => md.invalidate.mock.calls.length > 0);
+    expect(someMarkdownInvalidated).toBe(true);
+    expect(imageInstances[0].invalidate).toHaveBeenCalled();
+  });
+});
diff --git a/src/tui/components/tool-execution.test.ts b/src/tui/components/tool-execution.test.ts
new file mode 100644
index 0000000000..9c7576c09d
--- /dev/null
+++ b/src/tui/components/tool-execution.test.ts
@@ -0,0 +1,181 @@
+import { describe, expect, it, vi, beforeEach } from "vitest";
+import { ToolExecutionComponent } from "./tool-execution.js";
+
+// Mock tuiTheme first so we can verify color fns
+vi.mock("../theme.js", () => {
+  return {
+    tuiTheme: {
+      warning: vi.fn((s) => `warning:${s}`),
+      muted: vi.fn((s) => `muted:${s}`),
+      toolPendingBg: vi.fn((s) => `toolPendingBg:${s}`),
+      toolSuccessBg: vi.fn((s) => `toolSuccessBg:${s}`),
+      toolErrorBg: vi.fn((s) => `toolErrorBg:${s}`),
+      bold: vi.fn((s) => `bold:${s}`),
+      error: vi.fn((s) => `error:${s}`),
+      dim: vi.fn((s) => `dim:${s}`),
+    }
+  };
+});
+
+// We'll capture added children to verify logic
+let boxChildren: any[] = [];
+let boxBgFn: Function | null = null;
+let loaderInstance: any = null;
+
+vi.mock("@mariozechner/pi-tui", () => {
+  return {
+    TUI: class TUI {},
+    Box: class Box {
+      constructor() {
+        boxChildren = [];
+        boxBgFn = null;
+      }
+      addChild = vi.fn((child) => { boxChildren.push(child); });
+      clear = vi.fn(() => { boxChildren = []; });
+      setBgFn = vi.fn((fn) => { boxBgFn = fn; });
+      render = vi.fn().mockReturnValue(["rendered box"]);
+      invalidate = vi.fn();
+    },
+    Text: class Text {
+      constructor(public text: string, public x: number, public y: number) {}
+    },
+    Loader: class Loader {
+      constructor() {
+        loaderInstance = this;
+      }
+      stop = vi.fn();
+      invalidate = vi.fn();
+    },
+  };
+});
+
+describe("ToolExecutionComponent", () => {
+  let tui: any;
+
+  beforeEach(async () => {
+    vi.clearAllMocks();
+    boxChildren = [];
+    boxBgFn = null;
+    loaderInstance = null;
+    const { TUI } = await import("@mariozechner/pi-tui");
+    tui = new TUI();
+  });
+
+  it("initializes in running state with loader and pending bg", () => {
+    const args = { arg1: "value1" };
+    const component = new ToolExecutionComponent("testTool", args, tui);
+
+    expect(component).toBeDefined();
+
+    // Header text with arg summary
+    expect(boxChildren.length).toBe(2);
+    expect(boxChildren[0].text).toContain("bold:testTool");
+    expect(boxChildren[0].text).toContain("muted:arg1=value1");
+
+    // Loader should be added
+    expect(boxChildren[1]).toBe(loaderInstance);
+
+    // Check render
+    expect(component.render(100)).toEqual(["rendered box"]);
+  });
+
+  it("summarizes long arguments correctly and clips total summary length", () => {
+    const args = {
+      shortStr: "hello",
+      longStr: "this is a very long string that should be truncated because it exceeds thirty characters",
+      num: 42,
+      obj: { key: "val" }
+    };
+
+    new ToolExecutionComponent("testTool", args, tui);
+
+    const headerText = boxChildren[0].text;
+    expect(headerText).toContain("bold:testTool");
+
+    // "longStr" value should be truncated to 27 chars + "..." (30 chars)
+    expect(headerText).toContain("this is a very long string ...");
+    expect(headerText).not.toContain("exceeds thirty characters");
+
+    // other args should be present if they fit
+    expect(headerText).toContain("shortStr=hello");
+    expect(headerText).toContain("num=42");
+    expect(headerText).toContain('obj={"key":"val"}');
+  });
+
+  it("updates result with success", () => {
+    const component = new ToolExecutionComponent("testTool", {}, tui);
+
+    component.updateResult({ text: "Success output", isError: false });
+
+    // Loader should be stopped
+    expect(loaderInstance.stop).toHaveBeenCalled();
+
+    // Box bg should be updated to success
+    expect(boxBgFn).toBeDefined();
+    if (boxBgFn) expect(boxBgFn("test")).toBe("toolSuccessBg:test");
+
+    // Output text should be added
+    expect(boxChildren.length).toBe(2); // Header + Output
+    expect(boxChildren[1].text).toContain("muted:Success output");
+  });
+
+  it("updates result with error", () => {
+    const component = new ToolExecutionComponent("testTool", {}, tui);
+
+    component.updateResult({ text: "Error output", isError: true });
+
+    // Loader should be stopped
+    expect(loaderInstance.stop).toHaveBeenCalled();
+
+    // Box bg should be updated to error
+    expect(boxBgFn).toBeDefined();
+    if (boxBgFn) expect(boxBgFn("test")).toBe("toolErrorBg:test");
+
+    // Output text should be added with error color
+    expect(boxChildren.length).toBe(2); // Header + Output
+    expect(boxChildren[1].text).toContain("error:Error output");
+  });
+
+  it("collapses long output to 5 lines by default", () => {
+    const component = new ToolExecutionComponent("testTool", {}, tui);
+
+    const longOutput = Array.from({length: 10}, (_, i) => `Line ${i+1}`).join("\n");
+    component.updateResult({ text: longOutput, isError: false });
+
+    // Should have header, 5 lines of output, and the "more lines" indicator
+    expect(boxChildren.length).toBe(3);
+
+    const outputText = boxChildren[1].text;
+    expect(outputText).toContain("Line 1");
+    expect(outputText).toContain("Line 5");
+    expect(outputText).not.toContain("Line 6");
+
+    expect(boxChildren[2].text).toContain("dim:... (5 more lines, Ctrl+E to expand)");
+  });
+
+  it("expands long output when requested", () => {
+    const component = new ToolExecutionComponent("testTool", {}, tui);
+
+    const longOutput = Array.from({length: 10}, (_, i) => `Line ${i+1}`).join("\n");
+    component.updateResult({ text: longOutput, isError: false });
+
+    component.setExpanded(true);
+
+    // Should have header, and all lines of output
+    expect(boxChildren.length).toBe(2);
+
+    const outputText = boxChildren[1].text;
+    expect(outputText).toContain("Line 1");
+    expect(outputText).toContain("Line 10");
+  });
+
+  it("invalidates correctly", () => {
+    const component = new ToolExecutionComponent("testTool", {}, tui);
+    // Grab the mock instance from our mock setup
+    const { Box } = require("@mariozechner/pi-tui");
+
+    component.invalidate();
+
+    expect(loaderInstance.invalidate).toHaveBeenCalled();
+  });
+});
diff --git a/src/tui/components/user-message.test.ts b/src/tui/components/user-message.test.ts
new file mode 100644
index 0000000000..a356fa41d2
--- /dev/null
+++ b/src/tui/components/user-message.test.ts
@@ -0,0 +1,73 @@
+import { describe, expect, it, vi, beforeEach } from "vitest";
+import { UserMessageComponent } from "./user-message.js";
+
+let markdownText = "";
+let markdownThemeRef = null;
+let markdownOptions = null;
+
+vi.mock("@mariozechner/pi-tui", () => {
+  return {
+    Markdown: class Markdown {
+      constructor(text, top, left, theme, options) {
+        markdownText = text;
+        markdownThemeRef = theme;
+        markdownOptions = options;
+      }
+      render = vi.fn().mockReturnValue(["md line 1", "md line 2"]);
+      invalidate = vi.fn();
+    }
+  };
+});
+
+vi.mock("../theme.js", () => {
+  return {
+    miladyMarkdownTheme: { mockTheme: true },
+    tuiTheme: {
+      userMsgBg: vi.fn((t) => `bg:${t}`)
+    }
+  };
+});
+
+describe("UserMessageComponent", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    markdownText = "";
+    markdownThemeRef = null;
+    markdownOptions = null;
+  });
+
+  it("initializes with text and default theme", () => {
+    const text = "Hello world";
+    const component = new UserMessageComponent(text);
+
+    expect(component).toBeDefined();
+    expect(markdownText).toBe(text);
+    expect(markdownThemeRef).toEqual({ mockTheme: true });
+
+    // Check bgColor mapping
+    expect(markdownOptions).toBeDefined();
+    expect(markdownOptions.bgColor).toBeDefined();
+    expect(markdownOptions.bgColor("test")).toBe("bg:test");
+  });
+
+  it("renders with blank line prepended", () => {
+    const component = new UserMessageComponent("test");
+    const rendered = component.render(100);
+
+    expect(rendered.length).toBe(3);
+    expect(rendered[0]).toBe("");
+    expect(rendered[1]).toBe("md line 1");
+    expect(rendered[2]).toBe("md line 2");
+  });
+
+  it("invalidates correctly", () => {
+    const component = new UserMessageComponent("test");
+
+    // Get the markdown instance
+    const mdInstance = (component as any).markdown;
+
+    component.invalidate();
+
+    expect(mdInstance.invalidate).toHaveBeenCalled();
+  });
+});