/** * Cost Alert Time Window Tests * * Tests for verifying that cost-alert.ts uses proper time-utils functions * and repository functions with correct filtering (deletedAt, warmup exclusion). * * Key Differences After Fix: * | Window | Before | After | * |---------|--------------------------------|----------------------------------------| * | 5h | now - 5h | getTimeRangeForPeriod("5h") - same | * | Weekly | now - 7 days (rolling) | getTimeRangeForPeriod("weekly") - Monday | * | Monthly | Month start (no timezone) | getTimeRangeForPeriod("monthly") - TZ aware | * * Filters Added by Using sumKeyCostInTimeRange/sumProviderCostInTimeRange: * - deletedAt IS NULL * - blockedBy IS NULL OR blockedBy <> 'warmup' (EXCLUDE_WARMUP_CONDITION) */ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; // Track mock calls const mockGetTimeRangeForPeriod = vi.fn(); const mockSumKeyCostInTimeRange = vi.fn(); const mockSumProviderCostInTimeRange = vi.fn(); const mockDbSelect = vi.fn(); const mockDbFrom = vi.fn(); const mockDbWhere = vi.fn(); // Mock dependencies before importing the module under test vi.mock("@/drizzle/db", () => ({ db: { select: (...args: unknown[]) => { mockDbSelect(...args); return { from: (...fromArgs: unknown[]) => { mockDbFrom(...fromArgs); return { where: (...whereArgs: unknown[]) => mockDbWhere(...whereArgs), }; }, }; }, }, })); vi.mock("@/lib/logger", () => ({ logger: { info: vi.fn(), error: vi.fn(), debug: vi.fn(), warn: vi.fn(), }, })); vi.mock("@/lib/utils/timezone", () => ({ resolveSystemTimezone: vi.fn(async () => "Asia/Shanghai"), })); // Mock the time-utils module vi.mock("@/lib/rate-limit/time-utils", () => ({ getTimeRangeForPeriod: (...args: unknown[]) => mockGetTimeRangeForPeriod(...args), })); // Mock the statistics repository vi.mock("@/repository/statistics", () => ({ sumKeyCostInTimeRange: (...args: unknown[]) => mockSumKeyCostInTimeRange(...args), sumProviderCostInTimeRange: (...args: unknown[]) => mockSumProviderCostInTimeRange(...args), })); describe("Cost Alert Time Windows", () => { const nowMs = 1706000000000; // 2024-01-23 08:53:20 UTC (Tuesday) beforeEach(() => { vi.useFakeTimers(); vi.setSystemTime(new Date(nowMs)); vi.clearAllMocks(); // Reset module cache to ensure fresh imports with our mocks vi.resetModules(); // Default mock implementations for time ranges mockGetTimeRangeForPeriod.mockImplementation(async (period: string) => { const now = new Date(nowMs); switch (period) { case "5h": return { startTime: new Date(nowMs - 5 * 60 * 60 * 1000), endTime: now, }; case "weekly": // Monday 00:00 Shanghai (2024-01-22 00:00 +08:00 = 2024-01-21 16:00 UTC) return { startTime: new Date("2024-01-21T16:00:00.000Z"), endTime: now, }; case "monthly": // Month start (2024-01-01 00:00 +08:00 = 2023-12-31 16:00 UTC) return { startTime: new Date("2023-12-31T16:00:00.000Z"), endTime: now, }; default: throw new Error(`Unknown period: ${period}`); } }); // Default mock for cost queries mockSumKeyCostInTimeRange.mockResolvedValue(0); mockSumProviderCostInTimeRange.mockResolvedValue(0); // Default: return empty arrays for DB queries mockDbWhere.mockResolvedValue([]); }); afterEach(() => { vi.useRealTimers(); }); describe("checkUserQuotas", () => { it("should use getTimeRangeForPeriod('5h') for 5-hour window", async () => { // Setup: Key with 5h limit mockDbWhere.mockResolvedValue([ { id: 1, key: "test-key", userName: "Test User", limit5h: "10.00", limitWeek: null, limitMonth: null, }, ]); mockSumKeyCostInTimeRange.mockResolvedValue(5); const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); await generateCostAlerts(0.5); expect(mockGetTimeRangeForPeriod).toHaveBeenCalledWith("5h"); }); it("should use getTimeRangeForPeriod('weekly') for weekly window (natural week from Monday)", async () => { // Setup: Key with weekly limit mockDbWhere.mockResolvedValue([ { id: 1, key: "test-key", userName: "Test User", limit5h: null, limitWeek: "100.00", limitMonth: null, }, ]); mockSumKeyCostInTimeRange.mockResolvedValue(50); const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); await generateCostAlerts(0.5); expect(mockGetTimeRangeForPeriod).toHaveBeenCalledWith("weekly"); }); it("should use getTimeRangeForPeriod('monthly') for monthly window (natural month)", async () => { // Setup: Key with monthly limit mockDbWhere.mockResolvedValue([ { id: 1, key: "test-key", userName: "Test User", limit5h: null, limitWeek: null, limitMonth: "1000.00", }, ]); mockSumKeyCostInTimeRange.mockResolvedValue(500); const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); await generateCostAlerts(0.5); expect(mockGetTimeRangeForPeriod).toHaveBeenCalledWith("monthly"); }); it("should use sumKeyCostInTimeRange with keyId and correct time range", async () => { const expectedStart = new Date(nowMs - 5 * 60 * 60 * 1000); const expectedEnd = new Date(nowMs); mockDbWhere.mockResolvedValue([ { id: 1, key: "test-key", userName: "Test User", limit5h: "10.00", limitWeek: null, limitMonth: null, }, ]); mockSumKeyCostInTimeRange.mockResolvedValue(5); const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); await generateCostAlerts(0.5); // Should call sumKeyCostInTimeRange with keyId (not key string) and time range expect(mockSumKeyCostInTimeRange).toHaveBeenCalledWith( 1, // keyId expectedStart, expectedEnd ); }); it("should generate alert when cost exceeds threshold", async () => { mockDbWhere.mockResolvedValue([ { id: 1, key: "test-key", userName: "Test User", limit5h: "10.00", limitWeek: null, limitMonth: null, }, ]); mockSumKeyCostInTimeRange.mockResolvedValue(9); // 90% of limit const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); const alerts = await generateCostAlerts(0.8); // 80% threshold expect(alerts).toHaveLength(1); expect(alerts[0]).toMatchObject({ targetType: "user", targetName: "Test User", targetId: 1, currentCost: 9, quotaLimit: 10, threshold: 0.8, period: "5小时", }); }); it("should NOT generate alert when cost is below threshold", async () => { mockDbWhere.mockResolvedValue([ { id: 1, key: "test-key", userName: "Test User", limit5h: "10.00", limitWeek: null, limitMonth: null, }, ]); mockSumKeyCostInTimeRange.mockResolvedValue(7); // 70% of limit const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); const alerts = await generateCostAlerts(0.8); // 80% threshold expect(alerts).toHaveLength(0); }); }); describe("checkProviderQuotas", () => { it("should use getTimeRangeForPeriod('weekly') for provider weekly window", async () => { // First call returns empty keys, second call returns provider mockDbWhere .mockResolvedValueOnce([]) // keys query .mockResolvedValueOnce([ { id: 1, name: "Test Provider", limitWeek: "100.00", limitMonth: null }, ]); mockSumProviderCostInTimeRange.mockResolvedValue(50); const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); await generateCostAlerts(0.5); expect(mockGetTimeRangeForPeriod).toHaveBeenCalledWith("weekly"); }); it("should use getTimeRangeForPeriod('monthly') for provider monthly window", async () => { mockDbWhere .mockResolvedValueOnce([]) // keys query .mockResolvedValueOnce([ { id: 1, name: "Test Provider", limitWeek: null, limitMonth: "1000.00" }, ]); mockSumProviderCostInTimeRange.mockResolvedValue(500); const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); await generateCostAlerts(0.5); expect(mockGetTimeRangeForPeriod).toHaveBeenCalledWith("monthly"); }); it("should use sumProviderCostInTimeRange with correct time range", async () => { const expectedWeeklyStart = new Date("2024-01-21T16:00:00.000Z"); const expectedEnd = new Date(nowMs); mockDbWhere .mockResolvedValueOnce([]) // keys query .mockResolvedValueOnce([ { id: 1, name: "Test Provider", limitWeek: "100.00", limitMonth: null }, ]); mockSumProviderCostInTimeRange.mockResolvedValue(50); const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); await generateCostAlerts(0.5); expect(mockSumProviderCostInTimeRange).toHaveBeenCalledWith( 1, // providerId expectedWeeklyStart, expectedEnd ); }); it("should generate provider alert when cost exceeds threshold", async () => { mockDbWhere .mockResolvedValueOnce([]) // keys query .mockResolvedValueOnce([ { id: 1, name: "Test Provider", limitWeek: "100.00", limitMonth: null }, ]); mockSumProviderCostInTimeRange.mockResolvedValue(90); // 90% of limit const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); const alerts = await generateCostAlerts(0.8); // 80% threshold expect(alerts).toHaveLength(1); expect(alerts[0]).toMatchObject({ targetType: "provider", targetName: "Test Provider", targetId: 1, currentCost: 90, quotaLimit: 100, threshold: 0.8, period: "本周", }); }); }); describe("Time Window Semantics", () => { it("weekly window should use natural week (Monday) not rolling 7 days", async () => { // This test verifies that weekly uses natural week boundaries // If today is Tuesday, weekly should start from Monday 00:00 // NOT from 7 days ago mockDbWhere.mockResolvedValue([ { id: 1, key: "test-key", userName: "Test User", limit5h: null, limitWeek: "100.00", limitMonth: null, }, ]); mockSumKeyCostInTimeRange.mockResolvedValue(50); const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); await generateCostAlerts(0.5); // Verify getTimeRangeForPeriod was called for weekly expect(mockGetTimeRangeForPeriod).toHaveBeenCalledWith("weekly"); // Verify sumKeyCostInTimeRange was called expect(mockSumKeyCostInTimeRange).toHaveBeenCalled(); // Extract the actual startTime passed const callArgs = mockSumKeyCostInTimeRange.mock.calls[0]; const startTime = callArgs[1] as Date; // Should be Monday 00:00 Shanghai = Sunday 16:00 UTC expect(startTime.toISOString()).toBe("2024-01-21T16:00:00.000Z"); }); it("monthly window should use natural month (1st) with timezone awareness", async () => { mockDbWhere.mockResolvedValue([ { id: 1, key: "test-key", userName: "Test User", limit5h: null, limitWeek: null, limitMonth: "1000.00", }, ]); mockSumKeyCostInTimeRange.mockResolvedValue(500); const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); await generateCostAlerts(0.5); expect(mockGetTimeRangeForPeriod).toHaveBeenCalledWith("monthly"); const callArgs = mockSumKeyCostInTimeRange.mock.calls[0]; const startTime = callArgs[1] as Date; // Should be Jan 1st 00:00 Shanghai = Dec 31 16:00 UTC expect(startTime.toISOString()).toBe("2023-12-31T16:00:00.000Z"); }); }); describe("Warmup and Deleted Record Exclusion", () => { it("should use sumKeyCostInTimeRange which excludes warmup records", async () => { // This is a verification test - sumKeyCostInTimeRange already includes EXCLUDE_WARMUP_CONDITION // The old getKeyCostSince did NOT have this filter mockDbWhere.mockResolvedValue([ { id: 1, key: "test-key", userName: "Test User", limit5h: "10.00", limitWeek: null, limitMonth: null, }, ]); mockSumKeyCostInTimeRange.mockResolvedValue(5); const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); await generateCostAlerts(0.5); // Verify sumKeyCostInTimeRange is called (which has EXCLUDE_WARMUP_CONDITION built-in) expect(mockSumKeyCostInTimeRange).toHaveBeenCalled(); }); it("should use sumProviderCostInTimeRange which excludes deleted records", async () => { // sumProviderCostInTimeRange has: isNull(messageRequest.deletedAt) filter // The old getProviderCostSince did NOT have this filter mockDbWhere .mockResolvedValueOnce([]) .mockResolvedValueOnce([ { id: 1, name: "Test Provider", limitWeek: "100.00", limitMonth: null }, ]); mockSumProviderCostInTimeRange.mockResolvedValue(50); const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); await generateCostAlerts(0.5); // Verify sumProviderCostInTimeRange is called (which has deletedAt IS NULL built-in) expect(mockSumProviderCostInTimeRange).toHaveBeenCalled(); }); }); describe("Timezone Consistency", () => { it("should use system timezone for all time calculations", async () => { // getTimeRangeForPeriod internally uses resolveSystemTimezone() // This ensures all calculations are timezone-aware mockDbWhere.mockResolvedValue([ { id: 1, key: "test-key", userName: "Test User", limit5h: null, limitWeek: "100.00", limitMonth: null, }, ]); mockSumKeyCostInTimeRange.mockResolvedValue(50); const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); await generateCostAlerts(0.5); // The time ranges returned by getTimeRangeForPeriod are timezone-aware // This is verified by the mock implementation which uses timezone-aware dates expect(mockGetTimeRangeForPeriod).toHaveBeenCalled(); }); }); describe("Performance Optimization", () => { it("should pre-calculate time ranges once for all keys in checkUserQuotas", async () => { // Multiple keys with various limits mockDbWhere .mockResolvedValueOnce([ { id: 1, key: "key-1", userName: "User 1", limit5h: "10.00", limitWeek: null, limitMonth: null, }, { id: 2, key: "key-2", userName: "User 2", limit5h: null, limitWeek: "100.00", limitMonth: null, }, { id: 3, key: "key-3", userName: "User 3", limit5h: null, limitWeek: null, limitMonth: "1000.00", }, ]) .mockResolvedValueOnce([]); // providers query returns empty mockSumKeyCostInTimeRange.mockResolvedValue(5); const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); await generateCostAlerts(0.5); // getTimeRangeForPeriod should be called for user quotas (3 periods) + provider quotas (2 periods) // Total: 5 calls (5h + weekly + monthly for keys, weekly + monthly for providers) const calls = mockGetTimeRangeForPeriod.mock.calls.map((c) => c[0]); // Keys use all 3 periods, providers use weekly + monthly // So 5h should be called 1 time (keys only) // weekly should be called 2 times (keys + providers) // monthly should be called 2 times (keys + providers) expect(calls.filter((c) => c === "5h")).toHaveLength(1); expect(calls.filter((c) => c === "weekly")).toHaveLength(2); expect(calls.filter((c) => c === "monthly")).toHaveLength(2); }); it("should not call getTimeRangeForPeriod per-key (optimized)", async () => { // This tests that we pre-calculate ranges once, not N times for N keys const manyKeys = Array.from({ length: 10 }, (_, i) => ({ id: i + 1, key: `key-${i + 1}`, userName: `User ${i + 1}`, limit5h: "10.00", limitWeek: "100.00", limitMonth: "1000.00", })); mockDbWhere.mockResolvedValueOnce(manyKeys).mockResolvedValueOnce([]); // empty providers mockSumKeyCostInTimeRange.mockResolvedValue(5); const { generateCostAlerts } = await import("@/lib/notification/tasks/cost-alert"); await generateCostAlerts(0.5); // Even with 10 keys, we should only call getTimeRangeForPeriod once per period // Not 10 times per period const calls = mockGetTimeRangeForPeriod.mock.calls.map((c) => c[0]); expect(calls.filter((c) => c === "5h")).toHaveLength(1); // 1 for keys expect(calls.filter((c) => c === "weekly")).toHaveLength(2); // 1 for keys + 1 for providers expect(calls.filter((c) => c === "monthly")).toHaveLength(2); // 1 for keys + 1 for providers }); }); });