sentryagent-idp/tests/unit/services/TierService.test.ts

/**
 * Unit tests for src/services/TierService.ts
 */

import { Pool, QueryResult } from 'pg';
import Stripe from 'stripe';
import type { RedisClientType } from 'redis';
import { TierService } from '../../../src/services/TierService';
import { ValidationError, TierLimitError } from '../../../src/utils/errors';

// ── Mock helpers ──────────────────────────────────────────────────────────────

function makePool(queryFn: jest.Mock): Pool {
  return { query: queryFn } as unknown as Pool;
}

function makeRedis(getFn: jest.Mock): RedisClientType {
  return { get: getFn } as unknown as RedisClientType;
}

function makeStripe(overrides: Partial<{
  checkoutUrl: string;
}> = {}): Stripe {
  const url = overrides.checkoutUrl ?? 'https://checkout.stripe.com/tier_upgrade_test';
  return {
    checkout: {
      sessions: {
        create: jest.fn().mockResolvedValue({ url, id: 'cs_tier_1' }),
      },
    },
  } as unknown as Stripe;
}

// ════════════════════════════════════════════════════════════════════════════
// TierService
// ════════════════════════════════════════════════════════════════════════════

describe('TierService', () => {
  beforeEach(() => jest.clearAllMocks());

  // ────────────────────────────────────────────────────────────────
  // getStatus()
  // ────────────────────────────────────────────────────────────────

  describe('getStatus()', () => {
    it('should return correct ITierStatus shape with tier, limits, usage, and resetAt', async () => {
      // Pool: tier query, then agent count query
      const mockQuery = jest.fn()
        .mockResolvedValueOnce({ rows: [{ tier: 'pro' }] } as unknown as QueryResult)
        .mockResolvedValueOnce({ rows: [{ count: '7' }] } as unknown as QueryResult);

      const mockGet = jest.fn()
        .mockResolvedValueOnce('123')   // callsToday
        .mockResolvedValueOnce('456');  // tokensToday

      const service = new TierService(makePool(mockQuery), makeRedis(mockGet), makeStripe());
      const status = await service.getStatus('org-uuid-1');

      expect(status.tier).toBe('pro');
      expect(status.limits).toEqual({
        maxAgents: 100,
        maxCallsPerDay: 50_000,
        maxTokensPerDay: 50_000,
      });
      expect(status.usage).toEqual({
        callsToday: 123,
        tokensToday: 456,
        agentCount: 7,
      });
      expect(typeof status.resetAt).toBe('string');
      // resetAt must be a valid ISO 8601 timestamp
      expect(new Date(status.resetAt).toString()).not.toBe('Invalid Date');
    });

    it('should read usage from Redis keys rate:tier:calls:<orgId> and rate:tier:tokens:<orgId>', async () => {
      const mockQuery = jest.fn()
        .mockResolvedValueOnce({ rows: [{ tier: 'free' }] } as unknown as QueryResult)
        .mockResolvedValueOnce({ rows: [{ count: '3' }] } as unknown as QueryResult);

      const mockGet = jest.fn().mockResolvedValue('0');

      const service = new TierService(makePool(mockQuery), makeRedis(mockGet), makeStripe());
      await service.getStatus('org-redis-test');

      expect(mockGet).toHaveBeenCalledWith('rate:tier:calls:org-redis-test');
      expect(mockGet).toHaveBeenCalledWith('rate:tier:tokens:org-redis-test');
    });

    it('should default to free tier when no organization row is found', async () => {
      const mockQuery = jest.fn()
        .mockResolvedValueOnce({ rows: [] } as unknown as QueryResult)  // no org row
        .mockResolvedValueOnce({ rows: [{ count: '0' }] } as unknown as QueryResult);

      const mockGet = jest.fn().mockResolvedValue(null);

      const service = new TierService(makePool(mockQuery), makeRedis(mockGet), makeStripe());
      const status = await service.getStatus('org-unknown');

      expect(status.tier).toBe('free');
      expect(status.limits.maxAgents).toBe(10);
    });

    it('should return 0 for Redis counters when Redis keys are absent (null)', async () => {
      const mockQuery = jest.fn()
        .mockResolvedValueOnce({ rows: [{ tier: 'free' }] } as unknown as QueryResult)
        .mockResolvedValueOnce({ rows: [{ count: '2' }] } as unknown as QueryResult);

      const mockGet = jest.fn().mockResolvedValue(null);

      const service = new TierService(makePool(mockQuery), makeRedis(mockGet), makeStripe());
      const status = await service.getStatus('org-no-redis');

      expect(status.usage.callsToday).toBe(0);
      expect(status.usage.tokensToday).toBe(0);
    });
  });

  // ────────────────────────────────────────────────────────────────
  // initiateUpgrade()
  // ────────────────────────────────────────────────────────────────

  describe('initiateUpgrade()', () => {
    it('should throw ValidationError if already on target tier', async () => {
      const mockQuery = jest.fn().mockResolvedValue({
        rows: [{ tier: 'pro' }],
      } as unknown as QueryResult);

      const service = new TierService(makePool(mockQuery), makeRedis(jest.fn()), makeStripe());

      await expect(service.initiateUpgrade('org-1', 'pro')).rejects.toThrow(ValidationError);
    });

    it('should throw ValidationError when downgrade is attempted (pro → free)', async () => {
      const mockQuery = jest.fn().mockResolvedValue({
        rows: [{ tier: 'pro' }],
      } as unknown as QueryResult);

      const service = new TierService(makePool(mockQuery), makeRedis(jest.fn()), makeStripe());

      await expect(service.initiateUpgrade('org-1', 'free')).rejects.toThrow(ValidationError);
    });

    it('should create a Stripe checkout session and return checkoutUrl', async () => {
      const mockQuery = jest.fn().mockResolvedValue({
        rows: [{ tier: 'free' }],
      } as unknown as QueryResult);

      const stripe = makeStripe({ checkoutUrl: 'https://checkout.stripe.com/upgrade' });
      const service = new TierService(makePool(mockQuery), makeRedis(jest.fn()), stripe);

      const result = await service.initiateUpgrade('org-free', 'pro');

      expect(result.checkoutUrl).toBe('https://checkout.stripe.com/upgrade');
      expect(stripe.checkout.sessions.create).toHaveBeenCalledWith(
        expect.objectContaining({
          mode: 'subscription',
          client_reference_id: 'org-free',
          metadata: expect.objectContaining({ orgId: 'org-free', targetTier: 'pro' }),
        }),
      );
    });

    it('should throw when Stripe returns no session URL', async () => {
      const mockQuery = jest.fn().mockResolvedValue({
        rows: [{ tier: 'free' }],
      } as unknown as QueryResult);

      const stripe = {
        checkout: {
          sessions: {
            create: jest.fn().mockResolvedValue({ url: null, id: 'cs_no_url' }),
          },
        },
      } as unknown as Stripe;

      const service = new TierService(makePool(mockQuery), makeRedis(jest.fn()), stripe);

      await expect(service.initiateUpgrade('org-free', 'pro')).rejects.toThrow(
        'Stripe did not return a checkout session URL.',
      );
    });
  });

  // ────────────────────────────────────────────────────────────────
  // applyUpgrade()
  // ────────────────────────────────────────────────────────────────

  describe('applyUpgrade()', () => {
    it('should update the organizations table tier column', async () => {
      const mockQuery = jest.fn().mockResolvedValue({ rows: [] } as unknown as QueryResult);

      const service = new TierService(makePool(mockQuery), makeRedis(jest.fn()), makeStripe());
      await service.applyUpgrade('org-upgrade', 'pro');

      expect(mockQuery).toHaveBeenCalledWith(
        expect.stringContaining('UPDATE organizations'),
        ['pro', 'org-upgrade'],
      );
      expect(mockQuery).toHaveBeenCalledWith(
        expect.stringContaining('tier_updated_at'),
        expect.any(Array),
      );
    });

    it('should resolve without throwing on success', async () => {
      const mockQuery = jest.fn().mockResolvedValue({ rows: [] } as unknown as QueryResult);

      const service = new TierService(makePool(mockQuery), makeRedis(jest.fn()), makeStripe());
      await expect(service.applyUpgrade('org-upgrade', 'enterprise')).resolves.toBeUndefined();
    });
  });

  // ────────────────────────────────────────────────────────────────
  // enforceAgentLimit()
  // ────────────────────────────────────────────────────────────────

  describe('enforceAgentLimit()', () => {
    it('should throw TierLimitError when agent count is at or above the limit', async () => {
      // Agent count query returns 10 (= free tier max of 10)
      const mockQuery = jest.fn().mockResolvedValue({
        rows: [{ count: '10' }],
      } as unknown as QueryResult);

      const service = new TierService(makePool(mockQuery), makeRedis(jest.fn()), makeStripe());

      await expect(service.enforceAgentLimit('org-limited', 'free')).rejects.toThrow(TierLimitError);
    });

    it('should not throw when agent count is below the limit', async () => {
      const mockQuery = jest.fn().mockResolvedValue({
        rows: [{ count: '5' }],
      } as unknown as QueryResult);

      const service = new TierService(makePool(mockQuery), makeRedis(jest.fn()), makeStripe());

      await expect(service.enforceAgentLimit('org-ok', 'free')).resolves.toBeUndefined();
    });

    it('should never throw for enterprise tier (Infinity limit)', async () => {
      // pool.query should NOT be called because Infinity bypasses the check
      const mockQuery = jest.fn();

      const service = new TierService(makePool(mockQuery), makeRedis(jest.fn()), makeStripe());

      await expect(service.enforceAgentLimit('org-enterprise', 'enterprise')).resolves.toBeUndefined();
      // No DB query needed for Infinity limit
      expect(mockQuery).not.toHaveBeenCalled();
    });
  });
});