101 lines
4.0 KiB
TypeScript
101 lines
4.0 KiB
TypeScript
// @vitest-environment node
|
|
|
|
import { MemoryAdapter, createSharedStore } from "@hivegrid/shared-store";
|
|
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
|
|
|
import {
|
|
CHAT_GATEWAY_TOKEN_SESSION_LIMIT,
|
|
CHAT_GATEWAY_TOKEN_USER_LIMIT,
|
|
__resetChatGatewayTokenRateLimitStoreForTests,
|
|
__setChatGatewayTokenRateLimitStoreForTests,
|
|
checkChatGatewayTokenRateLimit,
|
|
} from "./chatGatewayRateLimit.server";
|
|
|
|
const stores: Array<ReturnType<typeof createSharedStore>> = [];
|
|
|
|
function useMemoryStore() {
|
|
const store = createSharedStore({ namespace: "chat-token-rate-test", adapter: new MemoryAdapter() });
|
|
stores.push(store);
|
|
__setChatGatewayTokenRateLimitStoreForTests(store);
|
|
return store;
|
|
}
|
|
|
|
beforeEach(() => {
|
|
vi.useFakeTimers();
|
|
vi.setSystemTime(new Date("2026-07-15T12:00:00.000Z"));
|
|
useMemoryStore();
|
|
});
|
|
|
|
afterEach(async () => {
|
|
__resetChatGatewayTokenRateLimitStoreForTests();
|
|
vi.useRealTimers();
|
|
await Promise.allSettled(stores.splice(0).map((store) => store.close()));
|
|
});
|
|
|
|
describe("chat gateway token rate limit", () => {
|
|
it("atomically caps token minting per session and resets with the fixed window", async () => {
|
|
const input = { tenantId: "tenant_1", userId: "user_1", sessionId: "session_1" };
|
|
const reservations = await Promise.all(Array.from({ length: CHAT_GATEWAY_TOKEN_SESSION_LIMIT + 1 }, () => checkChatGatewayTokenRateLimit(input)));
|
|
|
|
expect(reservations.filter((result) => result.allowed)).toHaveLength(CHAT_GATEWAY_TOKEN_SESSION_LIMIT);
|
|
expect(reservations.filter((result) => !result.allowed)).toEqual([{ allowed: false, retryAfterSeconds: 60 }]);
|
|
|
|
vi.advanceTimersByTime(60_000);
|
|
await expect(checkChatGatewayTokenRateLimit(input)).resolves.toMatchObject({ allowed: true });
|
|
});
|
|
|
|
it("applies an authoritative tenant-user cap across independent sessions", async () => {
|
|
for (let index = 0; index < CHAT_GATEWAY_TOKEN_USER_LIMIT; index += 1) {
|
|
await expect(
|
|
checkChatGatewayTokenRateLimit({
|
|
tenantId: "tenant_1",
|
|
userId: "user_1",
|
|
sessionId: `session_${Math.floor(index / CHAT_GATEWAY_TOKEN_SESSION_LIMIT)}`,
|
|
})
|
|
).resolves.toMatchObject({ allowed: true });
|
|
}
|
|
|
|
await expect(checkChatGatewayTokenRateLimit({ tenantId: "tenant_1", userId: "user_1", sessionId: "session_fresh" })).resolves.toEqual({
|
|
allowed: false,
|
|
retryAfterSeconds: 60,
|
|
});
|
|
await expect(checkChatGatewayTokenRateLimit({ tenantId: "tenant_2", userId: "user_1", sessionId: "session_fresh" })).resolves.toMatchObject({
|
|
allowed: true,
|
|
});
|
|
});
|
|
|
|
it("does not consume a session budget when the user-wide reservation is denied", async () => {
|
|
for (let index = 0; index < CHAT_GATEWAY_TOKEN_USER_LIMIT; index += 1) {
|
|
await checkChatGatewayTokenRateLimit({
|
|
tenantId: "tenant_1",
|
|
userId: "user_1",
|
|
sessionId: `session_${Math.floor(index / CHAT_GATEWAY_TOKEN_SESSION_LIMIT)}`,
|
|
});
|
|
}
|
|
vi.advanceTimersByTime(30_000);
|
|
await expect(checkChatGatewayTokenRateLimit({ tenantId: "tenant_1", userId: "user_1", sessionId: "session_fresh" })).resolves.toMatchObject({
|
|
allowed: false,
|
|
});
|
|
|
|
vi.advanceTimersByTime(30_001);
|
|
const nextWindow = await Promise.all(
|
|
Array.from({ length: CHAT_GATEWAY_TOKEN_SESSION_LIMIT }, () =>
|
|
checkChatGatewayTokenRateLimit({ tenantId: "tenant_1", userId: "user_1", sessionId: "session_fresh" })
|
|
)
|
|
);
|
|
|
|
expect(nextWindow.every((result) => result.allowed)).toBe(true);
|
|
await expect(checkChatGatewayTokenRateLimit({ tenantId: "tenant_1", userId: "user_1", sessionId: "session_fresh" })).resolves.toMatchObject({
|
|
allowed: false,
|
|
});
|
|
});
|
|
|
|
it("fails closed when the shared rate-limit store is unavailable", async () => {
|
|
const store = stores.at(-1);
|
|
if (!store) throw new Error("Expected a test store");
|
|
await store.close();
|
|
|
|
await expect(checkChatGatewayTokenRateLimit({ tenantId: "tenant_1", userId: "user_1", sessionId: "session_1" })).rejects.toThrow(/closed/i);
|
|
});
|
|
});
|