470 lines
16 KiB
TypeScript
470 lines
16 KiB
TypeScript
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<typeof fetch>[1]) =>
|
|
typeof init?.body === "string" ? init.body : undefined;
|
|
|
|
type EndpointHandler = (init?: Parameters<typeof fetch>[1]) => Promise<Response> | 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<typeof createProviderUsageFetch>) {
|
|
return fetchAntigravityUsage("token-123", 5000, mockFetch as unknown as typeof fetch);
|
|
}
|
|
|
|
function findWindow(snapshot: Awaited<ReturnType<typeof fetchAntigravityUsage>>, label: string) {
|
|
return snapshot.windows.find((window) => window.label === label);
|
|
}
|
|
|
|
function expectTokenExpired(snapshot: Awaited<ReturnType<typeof fetchAntigravityUsage>>) {
|
|
expect(snapshot.error).toBe("Token expired");
|
|
expect(snapshot.windows).toHaveLength(0);
|
|
}
|
|
|
|
function expectSingleWindow(
|
|
snapshot: Awaited<ReturnType<typeof fetchAntigravityUsage>>,
|
|
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();
|
|
});
|
|
});
|