import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest' import { RateLimiter } from './rate-limiter' import type { ConsumeResult, RateLimitStorageAdapter, TokenStatus } from './storage' import { MANUAL_EXECUTION_LIMIT, RATE_LIMITS, RateLimitError } from './types' vi.mock('@/lib/core/config/env-flags', () => ({ isBillingEnabled: true })) interface MockAdapter { consumeTokens: Mock getTokenStatus: Mock resetBucket: Mock } const createMockAdapter = (): MockAdapter => ({ consumeTokens: vi.fn(), getTokenStatus: vi.fn(), resetBucket: vi.fn(), }) describe('RateLimiter', () => { const testUserId = 'test-user-123' const freeSubscription = { plan: 'free', referenceId: testUserId } let mockAdapter: MockAdapter let rateLimiter: RateLimiter beforeEach(() => { vi.clearAllMocks() mockAdapter = createMockAdapter() rateLimiter = new RateLimiter(mockAdapter as RateLimitStorageAdapter) }) describe('checkRateLimitWithSubscription', () => { it('should allow unlimited requests for manual trigger type', async () => { const result = await rateLimiter.checkRateLimitWithSubscription( testUserId, freeSubscription, 'manual', false ) expect(result.allowed).toBe(true) expect(result.remaining).toBe(MANUAL_EXECUTION_LIMIT) expect(result.resetAt).toBeInstanceOf(Date) expect(mockAdapter.consumeTokens).not.toHaveBeenCalled() }) it('should consume tokens for API requests', async () => { const mockResult: ConsumeResult = { allowed: true, tokensRemaining: RATE_LIMITS.free.sync.maxTokens - 1, resetAt: new Date(Date.now() + 60000), } mockAdapter.consumeTokens.mockResolvedValue(mockResult) const result = await rateLimiter.checkRateLimitWithSubscription( testUserId, freeSubscription, 'api', false ) expect(result.allowed).toBe(true) expect(result.remaining).toBe(mockResult.tokensRemaining) expect(mockAdapter.consumeTokens).toHaveBeenCalledWith( `${testUserId}:sync`, 1, RATE_LIMITS.free.sync ) }) it('should use async bucket for async requests', async () => { const mockResult: ConsumeResult = { allowed: true, tokensRemaining: RATE_LIMITS.free.async.maxTokens - 1, resetAt: new Date(Date.now() + 60000), } mockAdapter.consumeTokens.mockResolvedValue(mockResult) await rateLimiter.checkRateLimitWithSubscription(testUserId, freeSubscription, 'api', true) expect(mockAdapter.consumeTokens).toHaveBeenCalledWith( `${testUserId}:async`, 1, RATE_LIMITS.free.async ) }) it('should use api-endpoint bucket for api-endpoint trigger', async () => { const mockResult: ConsumeResult = { allowed: true, tokensRemaining: RATE_LIMITS.free.apiEndpoint.maxTokens - 1, resetAt: new Date(Date.now() + 60000), } mockAdapter.consumeTokens.mockResolvedValue(mockResult) await rateLimiter.checkRateLimitWithSubscription( testUserId, freeSubscription, 'api-endpoint', false ) expect(mockAdapter.consumeTokens).toHaveBeenCalledWith( `${testUserId}:api-endpoint`, 1, RATE_LIMITS.free.apiEndpoint ) }) it('should deny requests when rate limit exceeded', async () => { const mockResult: ConsumeResult = { allowed: false, tokensRemaining: 0, resetAt: new Date(Date.now() + 60000), retryAfterMs: 30000, } mockAdapter.consumeTokens.mockResolvedValue(mockResult) const result = await rateLimiter.checkRateLimitWithSubscription( testUserId, freeSubscription, 'api', false ) expect(result.allowed).toBe(false) expect(result.remaining).toBe(0) expect(result.retryAfterMs).toBe(30000) }) it('should use organization key for team subscriptions', async () => { const orgId = 'org-123' const teamSubscription = { plan: 'team', referenceId: orgId } const mockResult: ConsumeResult = { allowed: true, tokensRemaining: RATE_LIMITS.team.sync.maxTokens - 1, resetAt: new Date(Date.now() + 60000), } mockAdapter.consumeTokens.mockResolvedValue(mockResult) await rateLimiter.checkRateLimitWithSubscription(testUserId, teamSubscription, 'api', false) expect(mockAdapter.consumeTokens).toHaveBeenCalledWith( `${orgId}:sync`, 1, RATE_LIMITS.team.sync ) }) it('should use user key when team subscription referenceId matches userId', async () => { const directTeamSubscription = { plan: 'team', referenceId: testUserId } const mockResult: ConsumeResult = { allowed: true, tokensRemaining: RATE_LIMITS.team.sync.maxTokens - 1, resetAt: new Date(Date.now() + 60000), } mockAdapter.consumeTokens.mockResolvedValue(mockResult) await rateLimiter.checkRateLimitWithSubscription( testUserId, directTeamSubscription, 'api', false ) expect(mockAdapter.consumeTokens).toHaveBeenCalledWith( `${testUserId}:sync`, 1, RATE_LIMITS.team.sync ) }) it('should allow on storage error (fail open)', async () => { mockAdapter.consumeTokens.mockRejectedValue(new Error('Storage error')) const result = await rateLimiter.checkRateLimitWithSubscription( testUserId, freeSubscription, 'api', false ) expect(result.allowed).toBe(true) expect(result.remaining).toBe(1) }) it('should work for all non-manual trigger types', async () => { const triggerTypes = ['api', 'webhook', 'schedule', 'chat'] as const const mockResult: ConsumeResult = { allowed: true, tokensRemaining: 10, resetAt: new Date(Date.now() + 60000), } mockAdapter.consumeTokens.mockResolvedValue(mockResult) for (const triggerType of triggerTypes) { await rateLimiter.checkRateLimitWithSubscription( testUserId, freeSubscription, triggerType, false ) expect(mockAdapter.consumeTokens).toHaveBeenCalled() mockAdapter.consumeTokens.mockClear() } }) }) describe('getRateLimitStatusWithSubscription', () => { it('should return unlimited status for manual trigger type', async () => { const status = await rateLimiter.getRateLimitStatusWithSubscription( testUserId, freeSubscription, 'manual', false ) expect(status.requestsPerMinute).toBe(MANUAL_EXECUTION_LIMIT) expect(status.maxBurst).toBe(MANUAL_EXECUTION_LIMIT) expect(status.remaining).toBe(MANUAL_EXECUTION_LIMIT) expect(mockAdapter.getTokenStatus).not.toHaveBeenCalled() }) it('should return status from storage for API requests', async () => { const mockStatus: TokenStatus = { tokensAvailable: 15, maxTokens: RATE_LIMITS.free.sync.maxTokens, lastRefillAt: new Date(), nextRefillAt: new Date(Date.now() + 60000), } mockAdapter.getTokenStatus.mockResolvedValue(mockStatus) const status = await rateLimiter.getRateLimitStatusWithSubscription( testUserId, freeSubscription, 'api', false ) expect(status.remaining).toBe(15) expect(status.requestsPerMinute).toBe(RATE_LIMITS.free.sync.refillRate) expect(status.maxBurst).toBe(RATE_LIMITS.free.sync.maxTokens) expect(mockAdapter.getTokenStatus).toHaveBeenCalledWith( `${testUserId}:sync`, RATE_LIMITS.free.sync ) }) }) describe('resetRateLimit', () => { it('should reset all bucket types for a user', async () => { mockAdapter.resetBucket.mockResolvedValue(undefined) await rateLimiter.resetRateLimit(testUserId) expect(mockAdapter.resetBucket).toHaveBeenCalledTimes(3) expect(mockAdapter.resetBucket).toHaveBeenCalledWith(`${testUserId}:sync`) expect(mockAdapter.resetBucket).toHaveBeenCalledWith(`${testUserId}:async`) expect(mockAdapter.resetBucket).toHaveBeenCalledWith(`${testUserId}:api-endpoint`) }) it('should throw error if reset fails', async () => { mockAdapter.resetBucket.mockRejectedValue(new Error('Reset failed')) await expect(rateLimiter.resetRateLimit(testUserId)).rejects.toThrow('Reset failed') }) }) describe('checkRateLimitDirect', () => { const config = { maxTokens: 3, refillRate: 1, refillIntervalMs: 60_000 } it('should reflect the storage decision when it succeeds', async () => { mockAdapter.consumeTokens.mockResolvedValue({ allowed: true, tokensRemaining: 2, resetAt: new Date(), }) const result = await rateLimiter.checkRateLimitDirect('public:contact:ip', config) expect(result.allowed).toBe(true) expect(result.remaining).toBe(2) }) it('should fail open on storage error by default', async () => { mockAdapter.consumeTokens.mockRejectedValue(new Error('Storage error')) const result = await rateLimiter.checkRateLimitDirect('public:contact:ip', config) expect(result.allowed).toBe(true) expect(result.remaining).toBe(1) }) it('should fail closed on storage error when failClosed is set', async () => { mockAdapter.consumeTokens.mockRejectedValue(new Error('Storage error')) const result = await rateLimiter.checkRateLimitDirect('public:contact:ip', config, { failClosed: true, }) expect(result.allowed).toBe(false) expect(result.remaining).toBe(0) }) }) describe('subscription plan handling', () => { it('should use pro plan limits', async () => { const proSubscription = { plan: 'pro', referenceId: testUserId } const mockResult: ConsumeResult = { allowed: true, tokensRemaining: RATE_LIMITS.pro.sync.maxTokens - 1, resetAt: new Date(Date.now() + 60000), } mockAdapter.consumeTokens.mockResolvedValue(mockResult) await rateLimiter.checkRateLimitWithSubscription(testUserId, proSubscription, 'api', false) expect(mockAdapter.consumeTokens).toHaveBeenCalledWith( `${testUserId}:sync`, 1, RATE_LIMITS.pro.sync ) }) it('should use enterprise plan limits', async () => { const enterpriseSubscription = { plan: 'enterprise', referenceId: 'org-enterprise' } const mockResult: ConsumeResult = { allowed: true, tokensRemaining: RATE_LIMITS.enterprise.sync.maxTokens - 1, resetAt: new Date(Date.now() + 60000), } mockAdapter.consumeTokens.mockResolvedValue(mockResult) await rateLimiter.checkRateLimitWithSubscription( testUserId, enterpriseSubscription, 'api', false ) expect(mockAdapter.consumeTokens).toHaveBeenCalledWith( `org-enterprise:sync`, 1, RATE_LIMITS.enterprise.sync ) }) it('should fall back to free plan when subscription is null', async () => { const mockResult: ConsumeResult = { allowed: true, tokensRemaining: RATE_LIMITS.free.sync.maxTokens - 1, resetAt: new Date(Date.now() + 60000), } mockAdapter.consumeTokens.mockResolvedValue(mockResult) await rateLimiter.checkRateLimitWithSubscription(testUserId, null, 'api', false) expect(mockAdapter.consumeTokens).toHaveBeenCalledWith( `${testUserId}:sync`, 1, RATE_LIMITS.free.sync ) }) }) describe('schedule trigger type', () => { it('should use sync bucket for schedule trigger', async () => { const mockResult: ConsumeResult = { allowed: true, tokensRemaining: 10, resetAt: new Date(Date.now() + 60000), } mockAdapter.consumeTokens.mockResolvedValue(mockResult) await rateLimiter.checkRateLimitWithSubscription( testUserId, freeSubscription, 'schedule', false ) expect(mockAdapter.consumeTokens).toHaveBeenCalledWith( `${testUserId}:sync`, 1, RATE_LIMITS.free.sync ) }) it('should use async bucket for schedule trigger with isAsync true', async () => { const mockResult: ConsumeResult = { allowed: true, tokensRemaining: 10, resetAt: new Date(Date.now() + 60000), } mockAdapter.consumeTokens.mockResolvedValue(mockResult) await rateLimiter.checkRateLimitWithSubscription( testUserId, freeSubscription, 'schedule', true ) expect(mockAdapter.consumeTokens).toHaveBeenCalledWith( `${testUserId}:async`, 1, RATE_LIMITS.free.async ) }) }) describe('getRateLimitStatusWithSubscription error handling', () => { it('should return default config on storage error', async () => { mockAdapter.getTokenStatus.mockRejectedValue(new Error('Storage error')) const status = await rateLimiter.getRateLimitStatusWithSubscription( testUserId, freeSubscription, 'api', false ) expect(status.remaining).toBe(0) expect(status.requestsPerMinute).toBe(RATE_LIMITS.free.sync.refillRate) expect(status.maxBurst).toBe(RATE_LIMITS.free.sync.maxTokens) }) }) }) describe('RateLimitError', () => { it('should create error with default status code 429', () => { const error = new RateLimitError('Rate limit exceeded') expect(error.message).toBe('Rate limit exceeded') expect(error.statusCode).toBe(429) expect(error.name).toBe('RateLimitError') }) it('should create error with custom status code', () => { const error = new RateLimitError('Custom error', 503) expect(error.message).toBe('Custom error') expect(error.statusCode).toBe(503) }) it('should be instanceof Error', () => { const error = new RateLimitError('Test') expect(error instanceof Error).toBe(true) expect(error instanceof RateLimitError).toBe(true) }) it('should have proper stack trace', () => { const error = new RateLimitError('Test error') expect(error.stack).toBeDefined() expect(error.stack).toContain('RateLimitError') }) })