Blame · Line-by-line history

`ai-cost-tracker.test.ts`

Each line is annotated with the commit that last touched it. Click any SHA to jump to that commit and see the surrounding change.

ai-cost-tracker.test.tsBlame313 lines · 1 contributor

Normal view Raw

4cc2287Claude	1	/**
	2	* Tests for src/lib/ai-cost-tracker.ts.
	3	*
	4	* Layered:
	5	* 1. Pure pricing arithmetic — known inputs → exact cents.
	6	* 2. Summary aggregation — synthetic in-memory rows → expected rollups.
	7	* 3. Dashboard helpers (projections, day-bucketing).
	8	* 4. DB-backed recording flow — gated on HAS_DB; uses a real user row.
	9	*/
	10
	11	import { describe, it, expect } from "bun:test";
	12
	13	import {
	14	MODEL_PRICING,
	15	DEFAULT_PRICING,
	16	aggregateEvents,
	17	computeCentsForCall,
	18	dailyAverageCents,
	19	extractUsage,
	20	formatCents,
	21	formatTokens,
	22	projectMonthEndCents,
	23	recordAiCost,
	24	startOfUtcMonth,
	25	toUtcDayKey,
	26	summarizeCostsForUser,
	27	} from "../lib/ai-cost-tracker";
	28
	29	const HAS_DB = Boolean(process.env.DATABASE_URL);
	30
	31	// ─── 1. Pricing arithmetic ───────────────────────────────────────────────
	32
	33	describe("computeCentsForCall — pricing arithmetic", () => {
	34	it("returns 0 for an all-zero call", () => {
	35	expect(computeCentsForCall("claude-sonnet-4-20250514", 0, 0)).toBe(0);
	36	});
	37
	38	it("uses the Sonnet rate card: 10k in + 5k out → 10k·0.3¢ + 5k·1.5¢ = 10.5¢ → 11¢", () => {
	39	const got = computeCentsForCall("claude-sonnet-4-20250514", 10_000, 5_000);
	40	expect(got).toBe(11);
	41	});
	42
	43	it("uses the Haiku rate card: 1M in + 0 out → 100¢ ($1)", () => {
	44	const got = computeCentsForCall("claude-haiku-4-5", 1_000_000, 0);
	45	expect(got).toBe(100);
	46	});
	47
	48	it("uses the Opus rate card: 1k in + 1k out → ceil(1.5+7.5)=9¢", () => {
	49	const got = computeCentsForCall("claude-opus-4", 1000, 1000);
	50	expect(got).toBe(9);
	51	});
	52
	53	it("falls back to DEFAULT_PRICING for an unknown model", () => {
	54	const got = computeCentsForCall("claude-future-9001", 10_000, 5_000);
	55	const expected = Math.ceil(
	56	(10_000 / 1000) * DEFAULT_PRICING.inputCentsPer1k +
	57	(5_000 / 1000) * DEFAULT_PRICING.outputCentsPer1k
	58	);
	59	expect(got).toBe(expected);
	60	});
	61
	62	it("never returns a negative value, even with negative inputs", () => {
	63	expect(computeCentsForCall("claude-sonnet-4-20250514", -5, -10)).toBe(0);
	64	});
	65
	66	it("never under-counts a real but tiny call to 0¢", () => {
	67	// 1 token in, 1 token out → fractional cents → must round UP to 1¢.
	68	const got = computeCentsForCall("claude-sonnet-4-20250514", 1, 1);
	69	expect(got).toBe(1);
	70	});
	71
	72	it("pricing table covers every model id our codebase passes to client.messages.create", () => {
	73	// Source-of-truth list (kept in sync with ai-client.ts + ai-review.ts).
	74	const referenced = [
	75	"claude-sonnet-4-20250514",
	76	"claude-haiku-4-5-20251001",
	77	];
	78	for (const model of referenced) {
	79	expect(MODEL_PRICING[model]).toBeDefined();
	80	}
	81	});
	82	});
	83
	84	// ─── 2. extractUsage / extractor robustness ─────────────────────────────
	85
	86	describe("extractUsage — Anthropic response shape", () => {
	87	it("returns zeros for null / non-object input", () => {
	88	expect(extractUsage(null)).toEqual({ input: 0, output: 0 });
	89	expect(extractUsage(undefined)).toEqual({ input: 0, output: 0 });
	90	expect(extractUsage("oops")).toEqual({ input: 0, output: 0 });
	91	});
	92
	93	it("reads .usage.input_tokens / .usage.output_tokens", () => {
	94	const fake = { usage: { input_tokens: 42, output_tokens: 7 } };
	95	expect(extractUsage(fake)).toEqual({ input: 42, output: 7 });
	96	});
	97
	98	it("tolerates missing fields", () => {
	99	const fake = { usage: { input_tokens: 5 } };
	100	expect(extractUsage(fake)).toEqual({ input: 5, output: 0 });
	101	});
	102
	103	it("ignores non-number fields", () => {
	104	const fake = { usage: { input_tokens: "5", output_tokens: null } };
	105	expect(extractUsage(fake)).toEqual({ input: 0, output: 0 });
	106	});
	107	});
	108
	109	// ─── 3. aggregateEvents — summary rollups ───────────────────────────────
	110
	111	describe("aggregateEvents — summary rollups", () => {
	112	const synth = [
	113	{
	114	occurredAt: new Date("2026-05-01T10:00:00Z"),
	115	model: "claude-sonnet-4-20250514",
	116	category: "ai_review",
	117	repositoryId: "repo-1",
	118	agentSessionId: "agent-A",
	119	centsEstimate: 4,
	120	inputTokens: 1000,
	121	outputTokens: 200,
	122	},
	123	{
	124	occurredAt: new Date("2026-05-01T11:00:00Z"),
	125	model: "claude-sonnet-4-20250514",
	126	category: "ci_healer",
	127	repositoryId: "repo-1",
	128	agentSessionId: null,
	129	centsEstimate: 2,
	130	inputTokens: 500,
	131	outputTokens: 100,
	132	},
	133	{
	134	occurredAt: new Date("2026-05-02T11:00:00Z"),
	135	model: "claude-haiku-4-5",
	136	category: "ai_review",
	137	repositoryId: "repo-2",
	138	agentSessionId: "agent-A",
	139	centsEstimate: 1,
	140	inputTokens: 800,
	141	outputTokens: 50,
	142	},
	143	];
	144
	145	const summary = aggregateEvents(synth);
	146
	147	it("rolls up total cents + token totals correctly", () => {
	148	expect(summary.totalCents).toBe(7);
	149	expect(summary.totalInputTokens).toBe(2300);
	150	expect(summary.totalOutputTokens).toBe(350);
	151	});
	152
	153	it("buckets by category, sorted by cents DESC", () => {
	154	expect(summary.byCategory.map((c) => c.category)).toEqual([
	155	"ai_review",
	156	"ci_healer",
	157	]);
	158	expect(summary.byCategory[0].cents).toBe(5);
	159	expect(summary.byCategory[1].cents).toBe(2);
	160	});
	161
	162	it("buckets by model", () => {
	163	const sonnet = summary.byModel.find(
	164	(m) => m.model === "claude-sonnet-4-20250514"
	165	);
	166	const haiku = summary.byModel.find((m) => m.model === "claude-haiku-4-5");
	167	expect(sonnet?.cents).toBe(6);
	168	expect(haiku?.cents).toBe(1);
	169	});
	170
	171	it("buckets by repository_id", () => {
	172	const repo1 = summary.byRepo.find((r) => r.repositoryId === "repo-1");
	173	const repo2 = summary.byRepo.find((r) => r.repositoryId === "repo-2");
	174	expect(repo1?.cents).toBe(6);
	175	expect(repo2?.cents).toBe(1);
	176	});
	177
	178	it("buckets by agent_session_id (incl. null bucket)", () => {
	179	const agentA = summary.byAgent.find((a) => a.agentSessionId === "agent-A");
	180	const noAgent = summary.byAgent.find((a) => a.agentSessionId === null);
	181	expect(agentA?.cents).toBe(5);
	182	expect(noAgent?.cents).toBe(2);
	183	});
	184
	185	it("buckets by UTC day, sorted ascending", () => {
	186	expect(summary.byDay).toEqual([
	187	{ day: "2026-05-01", cents: 6 },
	188	{ day: "2026-05-02", cents: 1 },
	189	]);
	190	});
	191
	192	it("returns zeros on an empty row set", () => {
	193	const empty = aggregateEvents([]);
	194	expect(empty.totalCents).toBe(0);
	195	expect(empty.byCategory).toHaveLength(0);
	196	expect(empty.byDay).toHaveLength(0);
	197	});
	198	});
	199
	200	// ─── 4. Projection + formatting helpers ─────────────────────────────────
	201
	202	describe("projection + formatting helpers", () => {
	203	it("startOfUtcMonth pins to day 1 at 00:00Z", () => {
	204	const d = new Date("2026-05-25T18:34:00Z");
	205	expect(startOfUtcMonth(d).toISOString()).toBe("2026-05-01T00:00:00.000Z");
	206	});
	207
	208	it("projectMonthEndCents linearly extrapolates from elapsed time", () => {
	209	// 31-day month (May). Halfway through (May 16 UTC noon-ish) →
	210	// projected ≈ 2x cents so far. Allow 10% slack for ms-precision drift.
	211	const now = new Date("2026-05-16T00:00:00Z");
	212	const projected = projectMonthEndCents(100, now);
	213	expect(projected).toBeGreaterThan(180);
	214	expect(projected).toBeLessThan(225);
	215	});
	216
	217	it("projectMonthEndCents returns the same number on the last day", () => {
	218	const now = new Date("2026-05-31T23:59:59Z");
	219	const projected = projectMonthEndCents(500, now);
	220	expect(projected).toBeGreaterThanOrEqual(500);
	221	// Should not balloon by more than a fraction of a percent.
	222	expect(projected).toBeLessThan(502);
	223	});
	224
	225	it("dailyAverageCents divides cents-so-far by elapsed days", () => {
	226	// May 10 UTC → day 10 of the month elapsed → avg = 100¢ / 10 = 10¢.
	227	const now = new Date("2026-05-10T12:00:00Z");
	228	expect(dailyAverageCents(100, now)).toBe(10);
	229	});
	230
	231	it("formatCents renders cents → $X.YY", () => {
	232	expect(formatCents(0)).toBe("$0.00");
	233	expect(formatCents(7)).toBe("$0.07");
	234	expect(formatCents(1234)).toBe("$12.34");
	235	expect(formatCents(123456)).toBe("$1,234.56");
	236	});
	237
	238	it("formatTokens groups thousands", () => {
	239	expect(formatTokens(0)).toBe("0");
	240	expect(formatTokens(1234)).toBe("1,234");
	241	expect(formatTokens(1_234_567)).toBe("1,234,567");
	242	});
	243
	244	it("toUtcDayKey is stable across timezones", () => {
	245	expect(toUtcDayKey(new Date("2026-05-25T23:59:59Z"))).toBe("2026-05-25");
	246	expect(toUtcDayKey(new Date("2026-05-25T00:00:01Z"))).toBe("2026-05-25");
	247	});
	248	});
	249
	250	// ─── 5. DB-backed recordAiCost — gated on HAS_DB ────────────────────────
	251
	252	describe.skipIf(!HAS_DB)("recordAiCost — DB-backed", () => {
	253	it.skipIf(!HAS_DB)(
	254	"inserts a row that summarizeCostsForUser can find",
	255	async () => {
	256	const { db } = await import("../db");
	257	const { users } = await import("../db/schema");
	258	// Mint a throwaway user.
	259	const [u] = await db
	260	.insert(users)
	261	.values({
	262	username: `cost-test-${Date.now()}`,
	263	email: `cost-${Date.now()}@example.invalid`,
	264	passwordHash: "$2a$10$" + "x".repeat(53),
	265	})
	266	.returning();
	267
	268	await recordAiCost({
	269	ownerUserId: u.id,
	270	model: "claude-sonnet-4-20250514",
	271	inputTokens: 10_000,
	272	outputTokens: 5_000,
	273	category: "ai_review",
	274	sourceKind: "pull_request",
	275	});
	276	await recordAiCost({
	277	ownerUserId: u.id,
	278	model: "claude-haiku-4-5",
	279	inputTokens: 1_000_000,
	280	outputTokens: 0,
	281	category: "other",
	282	sourceKind: "commit_message",
	283	});
	284
	285	const summary = await summarizeCostsForUser(u.id);
	286	// Sonnet row → 11¢ ; Haiku row → 100¢.
	287	expect(summary.totalCents).toBe(111);
	288	const cats = summary.byCategory.map((c) => c.category).sort();
	289	expect(cats).toEqual(["ai_review", "other"]);
	290	// byModel should split correctly too.
	291	const sonnet = summary.byModel.find(
	292	(m) => m.model === "claude-sonnet-4-20250514"
	293	);
	294	expect(sonnet?.inputTokens).toBe(10_000);
	295	expect(sonnet?.outputTokens).toBe(5_000);
	296	}
	297	);
	298
	299	it.skipIf(!HAS_DB)(
	300	"never throws on bad input — swallows DB errors",
	301	async () => {
	302	await expect(
	303	recordAiCost({
	304	ownerUserId: "not-a-uuid",
	305	model: "x",
	306	inputTokens: 1,
	307	outputTokens: 1,
	308	category: "ai_review",
	309	})
	310	).resolves.toBeUndefined();
	311	}
	312	);
	313	});