Momento/memento-note/tests/unit/entitlements.test.ts

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);
    });
  });
});