From 2f7508cf70728bfbb85de60b65ba16140cd336ca Mon Sep 17 00:00:00 2001 From: Jack Felke Date: Thu, 5 Mar 2026 01:44:55 -0700 Subject: [PATCH] test: add estimate_cost test suite (9 tests) --- tests/estimate-cost.test.ts | 208 ++++++++++++++++++++++++++++++++++++ 1 file changed, 208 insertions(+) create mode 100644 tests/estimate-cost.test.ts diff --git a/tests/estimate-cost.test.ts b/tests/estimate-cost.test.ts new file mode 100644 index 0000000..a593461 --- /dev/null +++ b/tests/estimate-cost.test.ts @@ -0,0 +1,208 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; + +// Mock dependencies before importing +vi.mock("node:fs", () => ({ + readFileSync: vi.fn(), + statSync: vi.fn(), +})); + +vi.mock("../src/lib/session-parser.js", () => ({ + findSessionDirs: vi.fn(), + findSessionFiles: vi.fn(), +})); + +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import { readFileSync, statSync } from "node:fs"; +import { findSessionDirs, findSessionFiles } from "../src/lib/session-parser.js"; +import { registerEstimateCost } from "../src/tools/estimate-cost.js"; + +describe("estimate_cost", () => { + let handler: any; + + beforeEach(() => { + vi.clearAllMocks(); + const server = { + tool: vi.fn((_name: string, _desc: string, _schema: any, h: any) => { + handler = h; + }), + } as any; + registerEstimateCost(server); + }); + + it("registers the tool", () => { + expect(handler).toBeDefined(); + }); + + it("returns error when no session files found", async () => { + vi.mocked(findSessionDirs).mockReturnValue([]); + const result = await handler({}); + expect(result.content[0].text).toContain("No session files found"); + }); + + it("returns error for missing file path", async () => { + vi.mocked(statSync).mockImplementation(() => { + throw new Error("ENOENT"); + }); + const result = await handler({ session_dir: "/nonexistent/file.jsonl" }); + expect(result.content[0].text).toContain("Session file not found"); + }); + + it("analyzes a simple session with no corrections", async () => { + const sessionData = [ + JSON.stringify({ + type: "user", + timestamp: "2026-01-01T10:00:00Z", + message: { content: "Hello, write me a function" }, + }), + JSON.stringify({ + type: "assistant", + timestamp: "2026-01-01T10:00:05Z", + message: { + content: [{ type: "text", text: "Here is a function that does what you need:\n```\nfunction add(a, b) { return a + b; }\n```" }], + }, + }), + ].join("\n"); + + vi.mocked(statSync).mockReturnValue({} as any); + vi.mocked(readFileSync).mockReturnValue(sessionData); + + const result = await handler({ session_dir: "/tmp/test.jsonl" }); + const text = result.content[0].text; + + expect(text).toContain("Session Cost Estimate"); + expect(text).toContain("1 prompts"); + expect(text).toContain("No corrections detected"); + }); + + it("detects corrections from user signals", async () => { + const sessionData = [ + JSON.stringify({ + type: "user", + timestamp: "2026-01-01T10:00:00Z", + message: { content: "Write a sorting function" }, + }), + JSON.stringify({ + type: "assistant", + timestamp: "2026-01-01T10:00:05Z", + message: { + content: [{ type: "text", text: "Here is bubble sort: " + "x".repeat(400) }], + }, + }), + JSON.stringify({ + type: "user", + timestamp: "2026-01-01T10:00:10Z", + message: { content: "No, I meant quicksort, not that" }, + }), + JSON.stringify({ + type: "assistant", + timestamp: "2026-01-01T10:00:15Z", + message: { + content: [{ type: "text", text: "Here is quicksort: " + "y".repeat(400) }], + }, + }), + ].join("\n"); + + vi.mocked(statSync).mockReturnValue({} as any); + vi.mocked(readFileSync).mockReturnValue(sessionData); + + const result = await handler({ session_dir: "/tmp/test.jsonl" }); + const text = result.content[0].text; + + expect(text).toContain("Corrections detected: 1"); + expect(text).toContain("Waste Analysis"); + }); + + it("detects preflight tool calls", async () => { + const sessionData = [ + JSON.stringify({ + type: "user", + timestamp: "2026-01-01T10:00:00Z", + message: { content: "Check my prompt" }, + }), + JSON.stringify({ + type: "assistant", + timestamp: "2026-01-01T10:00:05Z", + message: { + content: [ + { type: "tool_use", name: "preflight_check", input: { prompt: "test" } }, + { type: "text", text: "Running preflight check..." }, + ], + }, + }), + ].join("\n"); + + vi.mocked(statSync).mockReturnValue({} as any); + vi.mocked(readFileSync).mockReturnValue(sessionData); + + const result = await handler({ session_dir: "/tmp/test.jsonl" }); + const text = result.content[0].text; + + expect(text).toContain("Preflight checks: 1"); + }); + + it("uses specified pricing model", async () => { + const sessionData = [ + JSON.stringify({ + type: "user", + timestamp: "2026-01-01T10:00:00Z", + message: { content: "x".repeat(4000) }, // ~1000 tokens + }), + JSON.stringify({ + type: "assistant", + timestamp: "2026-01-01T10:01:00Z", + message: { + content: [{ type: "text", text: "y".repeat(4000) }], // ~1000 tokens + }, + }), + ].join("\n"); + + vi.mocked(statSync).mockReturnValue({} as any); + vi.mocked(readFileSync).mockReturnValue(sessionData); + + const result = await handler({ session_dir: "/tmp/test.jsonl", model: "claude-opus-4" }); + const text = result.content[0].text; + + expect(text).toContain("Claude Opus 4"); + }); + + it("finds latest session file when no path given", async () => { + vi.mocked(findSessionDirs).mockReturnValue([ + { sessionDir: "/home/user/.claude/projects/test", projectDir: "/test" }, + ]); + vi.mocked(findSessionFiles).mockReturnValue([ + { path: "/home/user/.claude/projects/test/session.jsonl", mtime: new Date("2026-01-01") }, + ]); + vi.mocked(statSync).mockReturnValue({} as any); + vi.mocked(readFileSync).mockReturnValue( + JSON.stringify({ + type: "user", + timestamp: "2026-01-01T10:00:00Z", + message: { content: "hello" }, + }), + ); + + const result = await handler({}); + expect(result.content[0].text).toContain("Session Cost Estimate"); + }); + + it("handles malformed JSONL lines gracefully", async () => { + const sessionData = [ + "not valid json", + JSON.stringify({ + type: "user", + timestamp: "2026-01-01T10:00:00Z", + message: { content: "hello" }, + }), + "{broken", + ].join("\n"); + + vi.mocked(statSync).mockReturnValue({} as any); + vi.mocked(readFileSync).mockReturnValue(sessionData); + + const result = await handler({ session_dir: "/tmp/test.jsonl" }); + const text = result.content[0].text; + + expect(text).toContain("Session Cost Estimate"); + expect(text).toContain("1 prompts"); + }); +});