import { describe, it, expect, vi, beforeEach } from 'vitest'; import { canUseFeature, checkEntitlementOrThrow, QuotaExceededError, } from '@/lib/entitlements'; vi.mock('@/lib/redis', () => ({ redis: { get: vi.fn(), mget: vi.fn(), eval: vi.fn(), }, })); vi.mock('@/lib/byok', () => ({ hasAnyActiveByok: vi.fn().mockResolvedValue(false), })); vi.mock('@/lib/prisma', () => ({ prisma: { subscription: { findUnique: vi.fn(), }, }, })); import { redis } from '@/lib/redis'; import { prisma } from '@/lib/prisma'; function mockActiveSubscription(tier: string) { vi.mocked(prisma.subscription.findUnique).mockResolvedValue({ userId: 'user1', tier: tier as any, status: 'ACTIVE' as any, currentPeriodStart: new Date(), currentPeriodEnd: new Date(), } as any); } describe('entitlements', () => { beforeEach(() => { vi.clearAllMocks(); }); describe('canUseFeature', () => { it('should allow BASIC user when under limit (30)', async () => { mockActiveSubscription('BASIC'); vi.mocked(redis.get).mockResolvedValue('5'); const result = await canUseFeature('user1', 'semantic_search'); expect(result.allowed).toBe(true); expect(result.remaining).toBe(25); expect(result.limit).toBe(30); }); it('should deny BASIC user when at limit (30)', async () => { mockActiveSubscription('BASIC'); vi.mocked(redis.get).mockResolvedValue('30'); const result = await canUseFeature('user1', 'semantic_search'); expect(result.allowed).toBe(false); expect(result.remaining).toBe(0); }); it('should allow PRO user with higher limits', async () => { mockActiveSubscription('PRO'); vi.mocked(redis.get).mockResolvedValue('50'); const result = await canUseFeature('user1', 'semantic_search'); expect(result.allowed).toBe(true); expect(result.limit).toBe(100); expect(result.remaining).toBe(50); }); it('should allow ENTERPRISE user unlimited features', async () => { mockActiveSubscription('ENTERPRISE'); const result = await canUseFeature('user1', 'semantic_search'); expect(result.allowed).toBe(true); expect(result.limit).toBe(Infinity); expect(result.remaining).toBe(Infinity); }); it('should default to BASIC tier when no subscription exists', async () => { vi.mocked(prisma.subscription.findUnique).mockResolvedValue(null); vi.mocked(redis.get).mockResolvedValue('10'); const result = await canUseFeature('user1', 'semantic_search'); expect(result.allowed).toBe(true); expect(result.limit).toBe(30); }); it('should return 0 remaining when usage exceeds limit', async () => { mockActiveSubscription('BASIC'); vi.mocked(redis.get).mockResolvedValue('35'); const result = await canUseFeature('user1', 'semantic_search'); expect(result.allowed).toBe(false); expect(result.remaining).toBe(0); }); it('should deny unknown features', async () => { mockActiveSubscription('BASIC'); const result = await canUseFeature('user1', 'unknownFeature'); expect(result.allowed).toBe(false); expect(result.limit).toBe(0); }); it('should fallback to BASIC for INACTIVE subscription', async () => { vi.mocked(prisma.subscription.findUnique).mockResolvedValue({ userId: 'user1', tier: 'PRO' as any, status: 'INACTIVE' as any, currentPeriodStart: new Date(), currentPeriodEnd: new Date(), } as any); vi.mocked(redis.get).mockResolvedValue('5'); const result = await canUseFeature('user1', 'semantic_search'); expect(result.limit).toBe(30); }); it('should return FEATURE_NOT_AVAILABLE for BASIC user requesting chat', async () => { mockActiveSubscription('BASIC'); const result = await canUseFeature('user1', 'chat'); expect(result.allowed).toBe(false); expect(result.reason).toBe('FEATURE_NOT_AVAILABLE'); }); it('should fail-open when Redis is down', async () => { mockActiveSubscription('BASIC'); vi.mocked(redis.get).mockRejectedValue(new Error('Connection refused')); const result = await canUseFeature('user1', 'semantic_search'); expect(result.allowed).toBe(true); }); }); describe('checkEntitlementOrThrow', () => { it('should throw QuotaExceededError when at limit', async () => { mockActiveSubscription('BASIC'); vi.mocked(redis.get).mockResolvedValue('30'); await expect( checkEntitlementOrThrow('user1', 'semantic_search'), ).rejects.toThrow(QuotaExceededError); }); it('should not throw when under limit', async () => { mockActiveSubscription('BASIC'); vi.mocked(redis.get).mockResolvedValue('5'); await expect( checkEntitlementOrThrow('user1', 'semantic_search'), ).resolves.toBeUndefined(); }); it('should return 402-compatible JSON from QuotaExceededError', () => { const error = new QuotaExceededError('PRO', 'semantic_search', 30, 30, false); const json = error.toJSON(); expect(json).toEqual({ error: 'QUOTA_EXCEEDED', feature: 'semantic_search', upgradeTier: 'PRO', byokConfigured: false, isGuestActor: false, }); }); }); describe('getEffectiveTier — grace period', () => { it('should retain PRO tier for PAST_DUE within billing period', async () => { const futureEnd = new Date(); futureEnd.setDate(futureEnd.getDate() + 10); vi.mocked(prisma.subscription.findUnique).mockResolvedValue({ userId: 'user1', tier: 'PRO' as any, status: 'PAST_DUE' as any, currentPeriodStart: new Date(), currentPeriodEnd: futureEnd, } as any); vi.mocked(redis.get).mockResolvedValue('5'); const result = await canUseFeature('user1', 'semantic_search'); expect(result.limit).toBe(100); }); it('should drop to BASIC for PAST_DUE past billing period', async () => { const pastEnd = new Date(); pastEnd.setDate(pastEnd.getDate() - 5); vi.mocked(prisma.subscription.findUnique).mockResolvedValue({ userId: 'user1', tier: 'PRO' as any, status: 'PAST_DUE' as any, currentPeriodStart: new Date(), currentPeriodEnd: pastEnd, } as any); vi.mocked(redis.get).mockResolvedValue('5'); const result = await canUseFeature('user1', 'semantic_search'); expect(result.limit).toBe(30); }); it('should retain tier for CANCELED with cancelAtPeriodEnd still in period', async () => { const futureEnd = new Date(); futureEnd.setDate(futureEnd.getDate() + 15); vi.mocked(prisma.subscription.findUnique).mockResolvedValue({ userId: 'user1', tier: 'BUSINESS' as any, status: 'CANCELED' as any, currentPeriodStart: new Date(), currentPeriodEnd: futureEnd, } as any); vi.mocked(redis.get).mockResolvedValue('50'); const result = await canUseFeature('user1', 'semantic_search'); expect(result.limit).toBe(1000); }); }); describe('QuotaExceededError', () => { it('should have correct properties', () => { const error = new QuotaExceededError('PRO', 'semantic_search', 30, 30, false); expect(error.code).toBe('QUOTA_EXCEEDED'); expect(error.upgradeTier).toBe('PRO'); expect(error.feature).toBe('semantic_search'); expect(error.currentQuota).toBe(30); expect(error.usedQuota).toBe(30); expect(error.byokConfigured).toBe(false); }); }); });