|
| 1 | +/** |
| 2 | + * @vitest-environment node |
| 3 | + * |
| 4 | + * Pins the rate-limit headers every v1 endpoint publishes. `X-RateLimit-Limit` |
| 5 | + * and `X-RateLimit-Remaining` must describe the same quantity — the token |
| 6 | + * bucket's capacity and the tokens left in it — or a client computing |
| 7 | + * `used = limit - remaining` gets a negative number. |
| 8 | + */ |
| 9 | + |
| 10 | +import { createMockRequest } from '@sim/testing' |
| 11 | +import { beforeEach, describe, expect, it, vi } from 'vitest' |
| 12 | + |
| 13 | +const { mockAuthenticateV1Request, mockGetSubscription, mockCheckRateLimit, mockGetRateLimit } = |
| 14 | + vi.hoisted(() => ({ |
| 15 | + mockAuthenticateV1Request: vi.fn(), |
| 16 | + mockGetSubscription: vi.fn(), |
| 17 | + mockCheckRateLimit: vi.fn(), |
| 18 | + mockGetRateLimit: vi.fn(), |
| 19 | + })) |
| 20 | + |
| 21 | +vi.mock('@/app/api/v1/auth', () => ({ |
| 22 | + authenticateV1Request: mockAuthenticateV1Request, |
| 23 | +})) |
| 24 | + |
| 25 | +vi.mock('@/lib/billing/core/subscription', () => ({ |
| 26 | + getHighestPrioritySubscription: mockGetSubscription, |
| 27 | +})) |
| 28 | + |
| 29 | +vi.mock('@/lib/core/rate-limiter', () => ({ |
| 30 | + getRateLimit: mockGetRateLimit, |
| 31 | + RateLimiter: class { |
| 32 | + checkRateLimitWithSubscription = mockCheckRateLimit |
| 33 | + }, |
| 34 | +})) |
| 35 | + |
| 36 | +import { checkRateLimit, createRateLimitResponse } from '@/app/api/v1/middleware' |
| 37 | + |
| 38 | +/** Mirrors `createBucketConfig`: capacity is the per-minute rate x burst multiplier. */ |
| 39 | +const TEAM_BUCKET = { maxTokens: 400, refillRate: 200, refillIntervalMs: 60_000 } |
| 40 | + |
| 41 | +function request() { |
| 42 | + return createMockRequest('GET', undefined, {}, 'http://localhost:3000/api/v1/workflows') |
| 43 | +} |
| 44 | + |
| 45 | +describe('checkRateLimit', () => { |
| 46 | + beforeEach(() => { |
| 47 | + vi.clearAllMocks() |
| 48 | + mockAuthenticateV1Request.mockResolvedValue({ |
| 49 | + authenticated: true, |
| 50 | + userId: 'user-1', |
| 51 | + keyType: 'personal', |
| 52 | + }) |
| 53 | + mockGetSubscription.mockResolvedValue({ plan: 'team' }) |
| 54 | + mockGetRateLimit.mockReturnValue(TEAM_BUCKET) |
| 55 | + mockCheckRateLimit.mockResolvedValue({ |
| 56 | + allowed: true, |
| 57 | + remaining: 399, |
| 58 | + resetAt: new Date('2026-07-28T18:28:48.354Z'), |
| 59 | + }) |
| 60 | + }) |
| 61 | + |
| 62 | + it('reports the bucket capacity as the limit, not the refill rate', async () => { |
| 63 | + const result = await checkRateLimit(request(), 'workflows') |
| 64 | + |
| 65 | + expect(result.limit).toBe(TEAM_BUCKET.maxTokens) |
| 66 | + expect(result.limit).not.toBe(TEAM_BUCKET.refillRate) |
| 67 | + }) |
| 68 | + |
| 69 | + it('never reports more remaining than the limit', async () => { |
| 70 | + const result = await checkRateLimit(request(), 'workflows') |
| 71 | + |
| 72 | + expect(result.remaining).toBeLessThanOrEqual(result.limit) |
| 73 | + }) |
| 74 | + |
| 75 | + it('reports a zero limit when authentication fails, since no bucket was consulted', async () => { |
| 76 | + mockAuthenticateV1Request.mockResolvedValue({ |
| 77 | + authenticated: false, |
| 78 | + error: 'API key required', |
| 79 | + }) |
| 80 | + |
| 81 | + const result = await checkRateLimit(request(), 'workflows') |
| 82 | + |
| 83 | + expect(result.allowed).toBe(false) |
| 84 | + expect(result.limit).toBe(0) |
| 85 | + expect(result.error).toBe('API key required') |
| 86 | + expect(mockCheckRateLimit).not.toHaveBeenCalled() |
| 87 | + }) |
| 88 | + |
| 89 | + it('reports a zero limit when the checker itself throws', async () => { |
| 90 | + mockCheckRateLimit.mockRejectedValue(new Error('redis down')) |
| 91 | + |
| 92 | + const result = await checkRateLimit(request(), 'workflows') |
| 93 | + |
| 94 | + expect(result.allowed).toBe(false) |
| 95 | + expect(result.limit).toBe(0) |
| 96 | + }) |
| 97 | +}) |
| 98 | + |
| 99 | +describe('createRateLimitResponse', () => { |
| 100 | + const throttled = { |
| 101 | + allowed: false, |
| 102 | + remaining: 0, |
| 103 | + limit: 400, |
| 104 | + resetAt: new Date('2026-07-28T18:28:48.354Z'), |
| 105 | + retryAfterMs: 30_000, |
| 106 | + } |
| 107 | + |
| 108 | + it('publishes consistent headers on a 429', () => { |
| 109 | + const response = createRateLimitResponse(throttled) |
| 110 | + |
| 111 | + expect(response.status).toBe(429) |
| 112 | + const limit = Number(response.headers.get('X-RateLimit-Limit')) |
| 113 | + const remaining = Number(response.headers.get('X-RateLimit-Remaining')) |
| 114 | + expect(remaining).toBeLessThanOrEqual(limit) |
| 115 | + expect(response.headers.get('X-RateLimit-Reset')).toBe(throttled.resetAt.toISOString()) |
| 116 | + expect(response.headers.get('Retry-After')).toBe('30') |
| 117 | + }) |
| 118 | + |
| 119 | + it('omits rate-limit headers on an auth failure, which never reached the bucket', async () => { |
| 120 | + const response = createRateLimitResponse({ |
| 121 | + allowed: false, |
| 122 | + remaining: 0, |
| 123 | + limit: 0, |
| 124 | + resetAt: new Date(), |
| 125 | + error: 'API key required', |
| 126 | + }) |
| 127 | + |
| 128 | + expect(response.status).toBe(401) |
| 129 | + expect(response.headers.get('X-RateLimit-Limit')).toBeNull() |
| 130 | + expect(response.headers.get('X-RateLimit-Remaining')).toBeNull() |
| 131 | + expect(response.headers.get('X-RateLimit-Reset')).toBeNull() |
| 132 | + await expect(response.json()).resolves.toEqual({ error: 'API key required' }) |
| 133 | + }) |
| 134 | +}) |
0 commit comments