import { describe, expect, it } from "vitest"; import { createProviderUsageFetch, makeResponse } from "../test-utils/provider-usage-fetch.js"; import { fetchAntigravityUsage } from "./provider-usage.fetch.antigravity.js"; const getRequestBody = (init?: Parameters[1]) => typeof init?.body === "string" ? init.body : undefined; type EndpointHandler = (init?: Parameters[1]) => Promise | Response; function createEndpointFetch(spec: { loadCodeAssist?: EndpointHandler; fetchAvailableModels?: EndpointHandler; }) { return createProviderUsageFetch(async (url, init) => { if (url.includes("loadCodeAssist")) { return (await spec.loadCodeAssist?.(init)) ?? makeResponse(404, "not found"); } if (url.includes("fetchAvailableModels")) { return (await spec.fetchAvailableModels?.(init)) ?? makeResponse(404, "not found"); } return makeResponse(404, "not found"); }); } async function runUsage(mockFetch: ReturnType) { return fetchAntigravityUsage("token-123", 5000, mockFetch as unknown as typeof fetch); } function findWindow(snapshot: Awaited>, label: string) { return snapshot.windows.find((window) => window.label === label); } function expectTokenExpired(snapshot: Awaited>) { expect(snapshot.error).toBe("Token expired"); expect(snapshot.windows).toHaveLength(0); } function expectSingleWindow( snapshot: Awaited>, label: string, ) { expect(snapshot.windows).toHaveLength(1); expect(snapshot.windows[0]?.label).toBe(label); return snapshot.windows[0]; } describe("fetchAntigravityUsage", () => { it("returns 3 windows when both endpoints succeed", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(200, { availablePromptCredits: 750, planInfo: { monthlyPromptCredits: 1000 }, planType: "Standard", currentTier: { id: "tier1", name: "Standard Tier" }, }), fetchAvailableModels: () => makeResponse(200, { models: { "gemini-pro-1.5": { quotaInfo: { remainingFraction: 0.6, resetTime: "2026-01-08T00:00:00Z", isExhausted: false, }, }, "gemini-flash-2.0": { quotaInfo: { remainingFraction: 0.8, resetTime: "2026-01-08T00:00:00Z", isExhausted: false, }, }, }, }), }); const snapshot = await runUsage(mockFetch); expect(snapshot.provider).toBe("google-antigravity"); expect(snapshot.displayName).toBe("Antigravity"); expect(snapshot.windows).toHaveLength(3); expect(snapshot.plan).toBe("Standard Tier"); expect(snapshot.error).toBeUndefined(); const creditsWindow = findWindow(snapshot, "Credits"); expect(creditsWindow?.usedPercent).toBe(25); // (1000 - 750) / 1000 * 100 const proWindow = findWindow(snapshot, "gemini-pro-1.5"); expect(proWindow?.usedPercent).toBe(40); // (1 - 0.6) * 100 expect(proWindow?.resetAt).toBe(new Date("2026-01-08T00:00:00Z").getTime()); const flashWindow = findWindow(snapshot, "gemini-flash-2.0"); expect(flashWindow?.usedPercent).toBeCloseTo(20, 1); // (1 - 0.8) * 100 expect(flashWindow?.resetAt).toBe(new Date("2026-01-08T00:00:00Z").getTime()); expect(mockFetch).toHaveBeenCalledTimes(2); }); it("returns Credits only when loadCodeAssist succeeds but fetchAvailableModels fails", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(200, { availablePromptCredits: 250, planInfo: { monthlyPromptCredits: 1000 }, currentTier: { name: "Free" }, }), fetchAvailableModels: () => makeResponse(403, { error: { message: "Permission denied" } }), }); const snapshot = await runUsage(mockFetch); expect(snapshot.provider).toBe("google-antigravity"); expect(snapshot.windows).toHaveLength(1); expect(snapshot.plan).toBe("Free"); expect(snapshot.error).toBeUndefined(); const creditsWindow = snapshot.windows[0]; expect(creditsWindow?.label).toBe("Credits"); expect(creditsWindow?.usedPercent).toBe(75); // (1000 - 250) / 1000 * 100 expect(mockFetch).toHaveBeenCalledTimes(2); }); it("returns model IDs when fetchAvailableModels succeeds but loadCodeAssist fails", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(500, "Internal server error"), fetchAvailableModels: () => makeResponse(200, { models: { "gemini-pro-1.5": { quotaInfo: { remainingFraction: 0.5, resetTime: "2026-01-08T00:00:00Z" }, }, "gemini-flash-2.0": { quotaInfo: { remainingFraction: 0.7, resetTime: "2026-01-08T00:00:00Z" }, }, }, }), }); const snapshot = await runUsage(mockFetch); expect(snapshot.provider).toBe("google-antigravity"); expect(snapshot.windows).toHaveLength(2); expect(snapshot.error).toBeUndefined(); const proWindow = findWindow(snapshot, "gemini-pro-1.5"); expect(proWindow?.usedPercent).toBe(50); // (1 - 0.5) * 100 const flashWindow = findWindow(snapshot, "gemini-flash-2.0"); expect(flashWindow?.usedPercent).toBeCloseTo(30, 1); // (1 - 0.7) * 100 expect(mockFetch).toHaveBeenCalledTimes(2); }); it.each([ { name: "uses cloudaicompanionProject string as project id", project: "projects/alpha", expectedBody: JSON.stringify({ project: "projects/alpha" }), }, { name: "uses cloudaicompanionProject object id when present", project: { id: "projects/beta" }, expectedBody: JSON.stringify({ project: "projects/beta" }), }, ])("project payload: $name", async ({ project, expectedBody }) => { let capturedBody: string | undefined; const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(200, { availablePromptCredits: 900, planInfo: { monthlyPromptCredits: 1000 }, cloudaicompanionProject: project, }), fetchAvailableModels: (init) => { capturedBody = getRequestBody(init); return makeResponse(200, { models: {} }); }, }); await runUsage(mockFetch); expect(capturedBody).toBe(expectedBody); }); it("returns error snapshot when both endpoints fail", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(403, { error: { message: "Access denied" } }), fetchAvailableModels: () => makeResponse(403, "Forbidden"), }); const snapshot = await runUsage(mockFetch); expect(snapshot.provider).toBe("google-antigravity"); expect(snapshot.windows).toHaveLength(0); expect(snapshot.error).toBe("Access denied"); expect(mockFetch).toHaveBeenCalledTimes(2); }); it("returns Token expired when fetchAvailableModels returns 401 and no windows", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(500, "Boom"), fetchAvailableModels: () => makeResponse(401, { error: { message: "Unauthorized" } }), }); const snapshot = await runUsage(mockFetch); expectTokenExpired(snapshot); }); it.each([ { name: "extracts plan info from currentTier.name", loadCodeAssist: { availablePromptCredits: 500, planInfo: { monthlyPromptCredits: 1000 }, planType: "Basic", currentTier: { id: "tier2", name: "Premium Tier" }, }, expectedPlan: "Premium Tier", }, { name: "falls back to planType when currentTier.name is missing", loadCodeAssist: { availablePromptCredits: 500, planInfo: { monthlyPromptCredits: 1000 }, planType: "Basic Plan", }, expectedPlan: "Basic Plan", }, ])("plan label: $name", async ({ loadCodeAssist, expectedPlan }) => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(200, loadCodeAssist), fetchAvailableModels: () => makeResponse(500, "Error"), }); const snapshot = await runUsage(mockFetch); expect(snapshot.plan).toBe(expectedPlan); }); it("includes reset times in model windows", async () => { const resetTime = "2026-01-10T12:00:00Z"; const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(500, "Error"), fetchAvailableModels: () => makeResponse(200, { models: { "gemini-pro-experimental": { quotaInfo: { remainingFraction: 0.3, resetTime }, }, }, }), }); const snapshot = await runUsage(mockFetch); const proWindow = snapshot.windows.find((w) => w.label === "gemini-pro-experimental"); expect(proWindow?.resetAt).toBe(new Date(resetTime).getTime()); }); it("parses string numbers correctly", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(200, { availablePromptCredits: "600", planInfo: { monthlyPromptCredits: "1000" }, }), fetchAvailableModels: () => makeResponse(200, { models: { "gemini-flash-lite": { quotaInfo: { remainingFraction: "0.9" }, }, }, }), }); const snapshot = await runUsage(mockFetch); expect(snapshot.windows).toHaveLength(2); const creditsWindow = snapshot.windows.find((w) => w.label === "Credits"); expect(creditsWindow?.usedPercent).toBe(40); // (1000 - 600) / 1000 * 100 const flashWindow = snapshot.windows.find((w) => w.label === "gemini-flash-lite"); expect(flashWindow?.usedPercent).toBeCloseTo(10, 1); // (1 - 0.9) * 100 }); it("skips internal models", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(200, { availablePromptCredits: 500, planInfo: { monthlyPromptCredits: 1000 }, cloudaicompanionProject: "projects/internal", }), fetchAvailableModels: () => makeResponse(200, { models: { chat_hidden: { quotaInfo: { remainingFraction: 0.1 } }, tab_hidden: { quotaInfo: { remainingFraction: 0.2 } }, "gemini-pro-1.5": { quotaInfo: { remainingFraction: 0.7 } }, }, }), }); const snapshot = await runUsage(mockFetch); expect(snapshot.windows.map((w) => w.label)).toEqual(["Credits", "gemini-pro-1.5"]); }); it("sorts models by usage and shows individual model IDs", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(500, "Error"), fetchAvailableModels: () => makeResponse(200, { models: { "gemini-pro-1.0": { quotaInfo: { remainingFraction: 0.8 } }, "gemini-pro-1.5": { quotaInfo: { remainingFraction: 0.3 } }, "gemini-flash-1.5": { quotaInfo: { remainingFraction: 0.6 } }, "gemini-flash-2.0": { quotaInfo: { remainingFraction: 0.9 } }, }, }), }); const snapshot = await runUsage(mockFetch); expect(snapshot.windows).toHaveLength(4); expect(snapshot.windows[0]?.label).toBe("gemini-pro-1.5"); expect(snapshot.windows[0]?.usedPercent).toBe(70); // (1 - 0.3) * 100 expect(snapshot.windows[1]?.label).toBe("gemini-flash-1.5"); expect(snapshot.windows[1]?.usedPercent).toBe(40); // (1 - 0.6) * 100 expect(snapshot.windows[2]?.label).toBe("gemini-pro-1.0"); expect(snapshot.windows[2]?.usedPercent).toBeCloseTo(20, 1); // (1 - 0.8) * 100 expect(snapshot.windows[3]?.label).toBe("gemini-flash-2.0"); expect(snapshot.windows[3]?.usedPercent).toBeCloseTo(10, 1); // (1 - 0.9) * 100 }); it("returns Token expired error on 401 from loadCodeAssist", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(401, { error: { message: "Unauthorized" } }), }); const snapshot = await runUsage(mockFetch); expectTokenExpired(snapshot); expect(mockFetch).toHaveBeenCalledTimes(1); // Should stop early on 401 }); it("handles empty models object gracefully", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(200, { availablePromptCredits: 800, planInfo: { monthlyPromptCredits: 1000 }, }), fetchAvailableModels: () => makeResponse(200, { models: {} }), }); const snapshot = await runUsage(mockFetch); expect(snapshot.windows).toHaveLength(1); const creditsWindow = snapshot.windows[0]; expect(creditsWindow?.label).toBe("Credits"); expect(creditsWindow?.usedPercent).toBe(20); }); it("handles missing or invalid model quota payloads", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(500, "Error"), fetchAvailableModels: () => makeResponse(200, { models: { no_quota: {}, missing_fraction: { quotaInfo: {} }, invalid_fraction: { quotaInfo: { remainingFraction: "oops" } }, valid_model: { quotaInfo: { remainingFraction: 0.25 } }, }, }), }); const snapshot = await runUsage(mockFetch); expect(snapshot.windows).toEqual([{ label: "valid_model", usedPercent: 75 }]); }); it("handles non-object models payload gracefully", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(200, { availablePromptCredits: 900, planInfo: { monthlyPromptCredits: 1000 }, }), fetchAvailableModels: () => makeResponse(200, { models: null }), }); const snapshot = await runUsage(mockFetch); expect(snapshot.windows).toEqual([{ label: "Credits", usedPercent: 10 }]); }); it("handles missing credits fields gracefully", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(200, { planType: "Free" }), fetchAvailableModels: () => makeResponse(200, { models: { "gemini-flash-experimental": { quotaInfo: { remainingFraction: 0.5 }, }, }, }), }); const snapshot = await runUsage(mockFetch); const flashWindow = expectSingleWindow(snapshot, "gemini-flash-experimental"); expect(flashWindow?.usedPercent).toBe(50); expect(snapshot.plan).toBe("Free"); }); it("handles invalid reset time gracefully", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(500, "Error"), fetchAvailableModels: () => makeResponse(200, { models: { "gemini-pro-test": { quotaInfo: { remainingFraction: 0.4, resetTime: "invalid-date" }, }, }, }), }); const snapshot = await runUsage(mockFetch); const proWindow = snapshot.windows.find((w) => w.label === "gemini-pro-test"); expect(proWindow?.usedPercent).toBe(60); expect(proWindow?.resetAt).toBeUndefined(); }); it("handles loadCodeAssist network errors with graceful degradation", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => { throw new Error("Network failure"); }, fetchAvailableModels: () => makeResponse(200, { models: { "gemini-flash-stable": { quotaInfo: { remainingFraction: 0.85 }, }, }, }), }); const snapshot = await runUsage(mockFetch); const flashWindow = expectSingleWindow(snapshot, "gemini-flash-stable"); expect(flashWindow?.usedPercent).toBeCloseTo(15, 1); expect(snapshot.error).toBeUndefined(); }); it("handles fetchAvailableModels network errors with graceful degradation", async () => { const mockFetch = createEndpointFetch({ loadCodeAssist: () => makeResponse(200, { availablePromptCredits: 300, planInfo: { monthlyPromptCredits: 1000 }, }), fetchAvailableModels: () => { throw new Error("Network failure"); }, }); const snapshot = await runUsage(mockFetch); expect(snapshot.windows).toEqual([{ label: "Credits", usedPercent: 70 }]); expect(snapshot.error).toBeUndefined(); }); });