Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/quota/__tests__/harness.test.ts

Raw
import { mkdtempSync, readdirSync, readFileSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import path from "node:path";
import { afterEach, describe, expect, it } from "vitest";
import { createTestSession, type TestSession } from "../../../test/harness";
import { __test } from "../index.ts";

/** 4h50m ahead, matching the reset window the footer renders. */
const RESET_AT = Date.now() + 4 * 3_600_000 + 50 * 60_000 + 30_000;

describe("quota pi-test-harness", () => {
	let t: TestSession | undefined;

	afterEach(() => {
		t?.dispose();
		t = undefined;
	});

	it("loads in real Pi extension runtime", async () => {
		const extensionPath = path.resolve(import.meta.dirname, "../index.ts");
		t = await createTestSession({ extensions: [extensionPath] });

		expect(t.session.extensionRunner.extensions).toHaveLength(1);
		expect(t.session.extensionRunner.extensions[0].path).toBe(extensionPath);
	}, 15000);

	it("writes its opt-in lifecycle trace inside the harness sandbox", async () => {
		const root = mkdtempSync(path.join(tmpdir(), "pi-ext-quota-debug-"));
		try {
			const extensionPath = path.resolve(import.meta.dirname, "../index.ts");
			t = await createTestSession({
				cwd: root,
				extensions: [extensionPath],
				env: { PI_QUOTA_DEBUG: "1" },
			});
			const runner = t.session.extensionRunner;
			await runner.emit({ type: "session_start", reason: "startup" });
			await runner.emit({ type: "session_shutdown", reason: "quit" });
			const stateHome = t.env.XDG_STATE_HOME;
			if (!stateHome) throw new Error("missing sandbox state home");
			const directory = path.join(stateHome, "pi-ext", "debug", "quota");
			const [file] = readdirSync(directory);
			const events = readFileSync(path.join(directory, file), "utf8")
				.trim()
				.split("\n")
				.map((line) => JSON.parse(line).event);
			expect(events).toEqual(
				expect.arrayContaining(["session.start", "session.shutdown"]),
			);
		} finally {
			t?.dispose();
			t = undefined;
			rmSync(root, { recursive: true, force: true });
		}
	});

	it("parses stolen provider quota shapes", () => {
		expect(
			__test.parseCopilot({
				quota_snapshots: {
					premium_interactions: { percent_remaining: 75 },
					chat: { percent_remaining: 40 },
				},
			}),
		).toMatchObject({
			name: "GitHub Copilot",
			fiveH: { used: 25, total: 100, remaining: 75 },
			weekly: { used: 60, total: 100, remaining: 40 },
		});

		expect(
			__test.parseGemini({
				buckets: [
					{ modelId: "gemini-pro", remainingFraction: 0.25 },
					{ modelId: "gemini-flash", remainingFraction: 0.75 },
				],
			}),
		).toMatchObject({
			name: "Gemini",
			fiveH: { used: 75, total: 100, remaining: 25 },
			weekly: { used: 25, total: 100, remaining: 75 },
		});

		expect(
			__test.parseMinimax({
				base_resp: { status_code: 0, status_msg: "" },
				model_remains: [
					{
						model_name: "MiniMax-M2",
						current_interval_total_count: 100,
						current_interval_usage_count: 30,
						current_weekly_total_count: 1000,
						current_weekly_usage_count: 800,
					},
				],
			}),
		).toMatchObject({
			name: "MiniMax",
			fiveH: { used: 70, total: 100, remaining: 30 },
			weekly: { used: 200, total: 1000, remaining: 800 },
		});

		expect(
			__test.parseMinimax({
				base_resp: { status_code: 0, status_msg: "" },
				model_remains: [
					{
						model_name: "general",
						current_interval_total_count: 0,
						current_interval_usage_count: 0,
						current_interval_remaining_percent: 25,
						current_weekly_total_count: 0,
						current_weekly_usage_count: 0,
						current_weekly_remaining_percent: 80,
					},
				],
			}),
		).toMatchObject({
			name: "MiniMax",
			fiveH: { used: 75, total: 100, remaining: 25 },
			weekly: { used: 20, total: 100, remaining: 80 },
		});
	});

	it("drops the provider name for the active provider only", () => {
		const anthropic = {
			name: "Anthropic",
			fiveH: { used: 4, total: 100, remaining: 96, resetAt: RESET_AT },
			weekly: { used: 0, total: 0, remaining: 0 },
		};

		expect(__test.segmentText(anthropic, true)).toEqual([
			"4% (4h50m)",
			"4%",
			"4%",
		]);
		expect(__test.segmentText(anthropic, false)).toEqual([
			"Anthropic 4% (4h50m)",
			"Anthropic 4%",
			"4%",
		]);
	});

	it("keeps the name in status fallbacks and spells out the non-tui label", () => {
		const providers = [
			{
				name: "Anthropic",
				fiveH: { used: 4, total: 100, remaining: 96, resetAt: RESET_AT },
				weekly: { used: 0, total: 0, remaining: 0 },
			},
		];

		expect(__test.quotaStatusLabel(providers, "Anthropic")).toBe(
			"Anthropic 4% (4h50m)",
		);
		expect(__test.quotaPlainLabel(providers, "Anthropic")).toBe(
			"Anthropic quota 4% used, resets in 4h50m",
		);
		expect(__test.quotaPlainLabel([], "Anthropic")).toBeUndefined();
	});
});