465 lines
15 KiB
JavaScript
465 lines
15 KiB
JavaScript
import { describe, it, expect, vi, beforeEach } from "vitest";
|
|
|
|
// Mock proxyAwareFetch before any imports that use it
|
|
vi.mock("../../open-sse/utils/proxyFetch.js", () => ({
|
|
proxyAwareFetch: vi.fn(),
|
|
}));
|
|
|
|
import { proxyAwareFetch } from "../../open-sse/utils/proxyFetch.js";
|
|
import {
|
|
parseWeeklyQuotaSummary,
|
|
fetchAntigravityWeeklyQuota,
|
|
_clearWeeklyCache,
|
|
} from "../../open-sse/services/usage/antigravity-weekly.js";
|
|
|
|
// — Fixtures ——————————————————————————————————————————————
|
|
const GEMINI_GROUP = {
|
|
displayName: "Gemini Models",
|
|
buckets: [
|
|
{
|
|
bucketId: "gemini-weekly-bucket",
|
|
displayName: "Weekly Limit",
|
|
remainingFraction: 0.75,
|
|
resetTime: "2026-09-15T00:00:00Z",
|
|
},
|
|
{
|
|
bucketId: "gemini-daily-bucket",
|
|
displayName: "Daily Limit",
|
|
remainingFraction: 0.9,
|
|
resetTime: "2026-09-09T00:00:00Z",
|
|
},
|
|
],
|
|
};
|
|
|
|
const CLAUDE_GPT_GROUP = {
|
|
displayName: "Claude and GPT models",
|
|
buckets: [
|
|
{
|
|
bucketId: "claude-gpt-weekly",
|
|
displayName: "Weekly Quota",
|
|
remainingFraction: 0.5,
|
|
resetTime: "2026-09-14T00:00:00Z",
|
|
},
|
|
],
|
|
};
|
|
|
|
const FULL_RESPONSE = { groups: [GEMINI_GROUP, CLAUDE_GPT_GROUP] };
|
|
|
|
const NESTED_RESPONSE = {
|
|
quotaSummary: {
|
|
groups: [GEMINI_GROUP, CLAUDE_GPT_GROUP],
|
|
},
|
|
};
|
|
|
|
// — parseWeeklyQuotaSummary ———————————————————————————————
|
|
describe("parseWeeklyQuotaSummary", () => {
|
|
it("extracts Gemini weekly quota from top-level groups", () => {
|
|
const result = parseWeeklyQuotaSummary(FULL_RESPONSE);
|
|
expect(result.gemini_weekly).toMatchObject({
|
|
used: 250,
|
|
total: 1000,
|
|
remainingPercentage: 75,
|
|
displayName: "Gemini (Weekly)",
|
|
unlimited: false,
|
|
});
|
|
expect(result.gemini_weekly.resetAt).toBe("2026-09-15T00:00:00.000Z");
|
|
});
|
|
|
|
it("extracts Claude & GPT weekly quota", () => {
|
|
const result = parseWeeklyQuotaSummary(FULL_RESPONSE);
|
|
expect(result.claude_gpt_weekly).toMatchObject({
|
|
used: 500,
|
|
total: 1000,
|
|
remainingPercentage: 50,
|
|
displayName: "Claude & GPT (Weekly)",
|
|
unlimited: false,
|
|
});
|
|
expect(result.claude_gpt_weekly.resetAt).toBe("2026-09-14T00:00:00.000Z");
|
|
});
|
|
|
|
it("handles alternate nested quotaSummary.groups shape", () => {
|
|
const result = parseWeeklyQuotaSummary(NESTED_RESPONSE);
|
|
expect(result.gemini_weekly).toBeDefined();
|
|
expect(result.claude_gpt_weekly).toBeDefined();
|
|
expect(result.gemini_weekly.remainingPercentage).toBe(75);
|
|
expect(result.claude_gpt_weekly.remainingPercentage).toBe(50);
|
|
});
|
|
|
|
it("skips non-weekly buckets", () => {
|
|
const data = {
|
|
groups: [{
|
|
displayName: "Gemini Models",
|
|
buckets: [
|
|
{
|
|
bucketId: "gemini-daily-bucket",
|
|
displayName: "Daily Limit",
|
|
remainingFraction: 0.9,
|
|
resetTime: "2026-09-09T00:00:00Z",
|
|
},
|
|
],
|
|
}],
|
|
};
|
|
const result = parseWeeklyQuotaSummary(data);
|
|
expect(result).toEqual({});
|
|
});
|
|
|
|
it("skips disabled weekly buckets", () => {
|
|
const data = {
|
|
groups: [{
|
|
displayName: "Gemini Models",
|
|
buckets: [{
|
|
bucketId: "gemini-weekly-bucket",
|
|
displayName: "Weekly Limit",
|
|
remainingFraction: 0.75,
|
|
resetTime: "2026-09-15T00:00:00Z",
|
|
disabled: true,
|
|
}],
|
|
}],
|
|
};
|
|
const result = parseWeeklyQuotaSummary(data);
|
|
expect(result).toEqual({});
|
|
});
|
|
|
|
it("returns empty object for null/undefined input", () => {
|
|
expect(parseWeeklyQuotaSummary(null)).toEqual({});
|
|
expect(parseWeeklyQuotaSummary(undefined)).toEqual({});
|
|
expect(parseWeeklyQuotaSummary("string")).toEqual({});
|
|
});
|
|
|
|
it("returns empty object for response with no groups", () => {
|
|
expect(parseWeeklyQuotaSummary({})).toEqual({});
|
|
expect(parseWeeklyQuotaSummary({ groups: "not-array" })).toEqual({});
|
|
expect(parseWeeklyQuotaSummary({ quotaSummary: {} })).toEqual({});
|
|
});
|
|
|
|
it("handles groups with no buckets gracefully", () => {
|
|
const data = {
|
|
groups: [{ displayName: "Gemini Models" }],
|
|
};
|
|
expect(parseWeeklyQuotaSummary(data)).toEqual({});
|
|
});
|
|
|
|
it("handles bucket with non-finite remainingFraction", () => {
|
|
const data = {
|
|
groups: [{
|
|
displayName: "Gemini Models",
|
|
buckets: [{
|
|
bucketId: "weekly-bucket",
|
|
displayName: "Weekly",
|
|
remainingFraction: "not-a-number",
|
|
}],
|
|
}],
|
|
};
|
|
expect(parseWeeklyQuotaSummary(data)).toEqual({});
|
|
});
|
|
|
|
it("ignores groups that don't match known families", () => {
|
|
const data = {
|
|
groups: [{
|
|
displayName: "Unknown AI Provider",
|
|
buckets: [{
|
|
bucketId: "weekly-bucket",
|
|
displayName: "Weekly",
|
|
remainingFraction: 0.5,
|
|
}],
|
|
}],
|
|
};
|
|
expect(parseWeeklyQuotaSummary(data)).toEqual({});
|
|
});
|
|
});
|
|
|
|
// — fetchAntigravityWeeklyQuota ———————————————————————————
|
|
describe("fetchAntigravityWeeklyQuota", () => {
|
|
beforeEach(() => {
|
|
proxyAwareFetch.mockReset();
|
|
_clearWeeklyCache();
|
|
});
|
|
|
|
it("fetches and returns parsed weekly quota on success", async () => {
|
|
proxyAwareFetch.mockResolvedValue({
|
|
ok: true,
|
|
json: async () => FULL_RESPONSE,
|
|
});
|
|
|
|
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
|
expect(result.gemini_weekly).toBeDefined();
|
|
expect(result.claude_gpt_weekly).toBeDefined();
|
|
});
|
|
|
|
it("returns {} on HTTP 401", async () => {
|
|
proxyAwareFetch.mockResolvedValue({ ok: false, status: 401 });
|
|
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
|
expect(result).toEqual({});
|
|
});
|
|
|
|
it("returns {} on HTTP 403", async () => {
|
|
proxyAwareFetch.mockResolvedValue({ ok: false, status: 403 });
|
|
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
|
expect(result).toEqual({});
|
|
});
|
|
|
|
it("returns {} on HTTP 404", async () => {
|
|
proxyAwareFetch.mockResolvedValue({ ok: false, status: 404 });
|
|
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
|
expect(result).toEqual({});
|
|
});
|
|
|
|
it("returns {} on HTTP 429", async () => {
|
|
proxyAwareFetch.mockResolvedValue({ ok: false, status: 429 });
|
|
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
|
expect(result).toEqual({});
|
|
});
|
|
|
|
it("returns {} on network error", async () => {
|
|
proxyAwareFetch.mockRejectedValue(new Error("network timeout"));
|
|
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
|
expect(result).toEqual({});
|
|
});
|
|
|
|
it("returns {} on malformed JSON response", async () => {
|
|
proxyAwareFetch.mockResolvedValue({
|
|
ok: true,
|
|
json: async () => { throw new SyntaxError("Unexpected token"); },
|
|
});
|
|
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
|
expect(result).toEqual({});
|
|
});
|
|
|
|
it("deduplicates concurrent requests for the same account", async () => {
|
|
let resolveResponse;
|
|
proxyAwareFetch.mockReturnValue(new Promise(resolve => {
|
|
resolveResponse = resolve;
|
|
}));
|
|
|
|
const p1 = fetchAntigravityWeeklyQuota("token", "project-1");
|
|
const p2 = fetchAntigravityWeeklyQuota("token", "project-1");
|
|
|
|
resolveResponse({ ok: true, json: async () => FULL_RESPONSE });
|
|
|
|
const [r1, r2] = await Promise.all([p1, p2]);
|
|
expect(r1).toEqual(r2);
|
|
expect(proxyAwareFetch).toHaveBeenCalledTimes(1);
|
|
});
|
|
|
|
it("serves cached result within TTL", async () => {
|
|
proxyAwareFetch.mockResolvedValue({
|
|
ok: true,
|
|
json: async () => FULL_RESPONSE,
|
|
});
|
|
|
|
await fetchAntigravityWeeklyQuota("token", "project-1");
|
|
const result = await fetchAntigravityWeeklyQuota("token", "project-1");
|
|
|
|
expect(proxyAwareFetch).toHaveBeenCalledTimes(1);
|
|
expect(result.gemini_weekly).toBeDefined();
|
|
});
|
|
|
|
it("sends correct headers and body", async () => {
|
|
proxyAwareFetch.mockResolvedValue({
|
|
ok: true,
|
|
json: async () => ({ groups: [] }),
|
|
});
|
|
|
|
await fetchAntigravityWeeklyQuota("token", "project-1");
|
|
|
|
expect(proxyAwareFetch).toHaveBeenCalledWith(
|
|
"https://daily-cloudcode-pa.googleapis.com/v1internal:retrieveUserQuotaSummary",
|
|
expect.objectContaining({
|
|
method: "POST",
|
|
headers: expect.objectContaining({
|
|
"Authorization": "Bearer token",
|
|
"User-Agent": "antigravity/ide/2.11.0 darwin/arm64",
|
|
"Content-Type": "application/json",
|
|
"X-Client-Name": "antigravity",
|
|
}),
|
|
body: JSON.stringify({ project: "project-1" }),
|
|
}),
|
|
expect.any(Object),
|
|
);
|
|
});
|
|
});
|
|
|
|
// — Integration: weekly failure does not affect existing quotas —————
|
|
describe("weekly quota isolation from existing quota", () => {
|
|
beforeEach(() => {
|
|
proxyAwareFetch.mockReset();
|
|
_clearWeeklyCache();
|
|
});
|
|
|
|
it("existing getAntigravityUsage succeeds even when weekly RPC fails", async () => {
|
|
proxyAwareFetch.mockImplementation(async (url) => {
|
|
if (url.includes(":loadCodeAssist")) {
|
|
return {
|
|
ok: true,
|
|
status: 200,
|
|
json: async () => ({ cloudaicompanionProject: "p1", currentTier: { name: "Pro" }, paidTier: { id: "g1-pro-tier", name: "Google AI Pro" } }),
|
|
};
|
|
}
|
|
if (url.includes(":fetchAvailableModels")) {
|
|
return {
|
|
ok: true,
|
|
status: 200,
|
|
json: async () => ({
|
|
models: {
|
|
"gemini-3.8-flash-high": {
|
|
displayName: "Gemini 3.8 Flash (High)",
|
|
quotaInfo: { remainingFraction: 0.85, resetTime: "2026-09-15T00:00:00Z" },
|
|
},
|
|
},
|
|
}),
|
|
};
|
|
}
|
|
if (url.includes(":retrieveUserQuotaSummary")) {
|
|
throw new Error("weekly endpoint unavailable");
|
|
}
|
|
return { ok: false, status: 404 };
|
|
});
|
|
|
|
const { getAntigravityUsage } = await import("../../open-sse/services/usage/google.js");
|
|
const result = await getAntigravityUsage("token", {});
|
|
|
|
expect(result.quotas["gemini-3.8-flash-high"]).toMatchObject({
|
|
used: 150,
|
|
total: 1000,
|
|
remainingPercentage: 85,
|
|
});
|
|
expect(result.quotas.gemini_weekly).toBeUndefined();
|
|
expect(result.quotas.claude_gpt_weekly).toBeUndefined();
|
|
expect(result.message).toBeUndefined();
|
|
});
|
|
|
|
it("free-tier accounts only show weekly quotas, not per-model short-window quotas", async () => {
|
|
proxyAwareFetch.mockImplementation(async (url) => {
|
|
if (url.includes(":loadCodeAssist")) {
|
|
return {
|
|
ok: true,
|
|
status: 200,
|
|
json: async () => ({ cloudaicompanionProject: "p1", currentTier: { name: "Starter" }, paidTier: { id: "free-tier", name: "Antigravity Starter Quota" } }),
|
|
};
|
|
}
|
|
if (url.includes(":fetchAvailableModels")) {
|
|
return {
|
|
ok: true,
|
|
status: 200,
|
|
json: async () => ({
|
|
models: {
|
|
"gemini-3.8-flash-high": {
|
|
displayName: "Gemini 3.8 Flash (High)",
|
|
quotaInfo: { remainingFraction: 1, resetTime: "2026-09-15T00:00:00Z" },
|
|
},
|
|
"claude-sonnet-4-6": {
|
|
displayName: "Claude Sonnet 4.6",
|
|
// Missing remainingFraction — free tier exhausted
|
|
quotaInfo: { resetTime: "2026-09-13T12:00:00Z" },
|
|
},
|
|
},
|
|
}),
|
|
};
|
|
}
|
|
if (url.includes(":retrieveUserQuotaSummary")) {
|
|
return {
|
|
ok: true,
|
|
status: 200,
|
|
json: async () => ({
|
|
groups: [{
|
|
displayName: "Gemini Models",
|
|
buckets: [{
|
|
bucketId: "gemini-weekly",
|
|
displayName: "Weekly Limit Remaining",
|
|
remainingFraction: 1,
|
|
resetTime: "2026-09-15T00:00:00Z",
|
|
}],
|
|
}, {
|
|
displayName: "Claude and GPT models",
|
|
buckets: [{
|
|
bucketId: "3p-weekly",
|
|
displayName: "Weekly Limit Remaining",
|
|
remainingFraction: 0,
|
|
resetTime: "2026-09-13T12:00:00Z",
|
|
}],
|
|
}],
|
|
}),
|
|
};
|
|
}
|
|
return { ok: false, status: 404 };
|
|
});
|
|
|
|
const { getAntigravityUsage } = await import("../../open-sse/services/usage/google.js");
|
|
const result = await getAntigravityUsage("token", {});
|
|
|
|
// Per-model quotas should be absent (free-tier accounts skip model parsing)
|
|
expect(result.quotas["gemini-3.8-flash-high"]).toBeUndefined();
|
|
expect(result.quotas["claude-sonnet-4-6"]).toBeUndefined();
|
|
|
|
// Only weekly quotas should appear
|
|
expect(result.quotas.gemini_weekly).toMatchObject({
|
|
used: 0,
|
|
total: 1000,
|
|
remainingPercentage: 100,
|
|
});
|
|
expect(result.quotas.claude_gpt_weekly).toMatchObject({
|
|
used: 1000,
|
|
total: 1000,
|
|
remainingPercentage: 0,
|
|
});
|
|
});
|
|
|
|
it("reconciles weekly quota to 0% when all paid-tier family models are exhausted", async () => {
|
|
proxyAwareFetch.mockImplementation(async (url) => {
|
|
if (url.includes(":loadCodeAssist")) {
|
|
return {
|
|
ok: true,
|
|
status: 200,
|
|
json: async () => ({ cloudaicompanionProject: "p1", currentTier: { name: "Pro" }, paidTier: { id: "g1-pro-tier", name: "Google AI Pro" } }),
|
|
};
|
|
}
|
|
if (url.includes(":fetchAvailableModels")) {
|
|
return {
|
|
ok: true,
|
|
status: 200,
|
|
json: async () => ({
|
|
models: {
|
|
"gemini-3.8-flash-high": {
|
|
displayName: "Gemini 3.8 Flash (High)",
|
|
// Exhausted model: no remainingFraction, future resetTime
|
|
quotaInfo: { resetTime: "2026-09-13T12:00:00Z" },
|
|
},
|
|
},
|
|
}),
|
|
};
|
|
}
|
|
if (url.includes(":retrieveUserQuotaSummary")) {
|
|
return {
|
|
ok: true,
|
|
status: 200,
|
|
json: async () => ({
|
|
groups: [{
|
|
displayName: "Gemini Models",
|
|
buckets: [{
|
|
bucketId: "gemini-weekly",
|
|
displayName: "Weekly Limit Remaining",
|
|
remainingFraction: 1,
|
|
resetTime: "2026-09-15T00:00:00Z",
|
|
}],
|
|
}],
|
|
}),
|
|
};
|
|
}
|
|
return { ok: false, status: 404 };
|
|
});
|
|
|
|
const { getAntigravityUsage } = await import("../../open-sse/services/usage/google.js");
|
|
const result = await getAntigravityUsage("token", {});
|
|
|
|
// Per-model quota should show exhausted
|
|
expect(result.quotas["gemini-3.8-flash-high"].remainingPercentage).toBe(0);
|
|
// Weekly quota should be reconciled to 0% with the family reset time
|
|
expect(result.quotas.gemini_weekly).toMatchObject({
|
|
used: 1000,
|
|
total: 1000,
|
|
remainingPercentage: 0,
|
|
resetAt: "2026-09-13T12:00:00.000Z",
|
|
});
|
|
});
|
|
});
|