Momento/memento-note/tests/unit/brainstorm-billing.test.ts

import { describe, it, expect, vi, beforeEach } from 'vitest'
import {
  billingOwnerFromSession,
  getBillingOwner,
} from '@/lib/brainstorm-collab'
import {
  checkSessionEntitlementOrThrow,
  incrementUsageAsync,
  QuotaExceededError,
} from '@/lib/entitlements'

const { prismaMock } = vi.hoisted(() => ({
  prismaMock: {
    subscription: { findUnique: vi.fn() },
    brainstormSession: { findUnique: vi.fn() },
  },
}))

vi.mock('@/lib/prisma', () => ({
  prisma: prismaMock,
  default: prismaMock,
}))

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/config', () => ({
  getSystemConfig: vi.fn().mockResolvedValue({
    AI_PROVIDER_CHAT: 'openai',
  }),
}))

vi.mock('@/lib/ai/provider-for-user', () => ({
  willUseByokForLane: vi.fn().mockResolvedValue({
    providerType: 'openai',
    usedByok: false,
  }),
}))

vi.mock('@/lib/credits', () => ({
  reserveCreditsOrThrow: vi.fn(),
  resolveCreditCost: vi.fn().mockReturnValue(1),
  releaseCredits: vi.fn(),
}))

import { redis } from '@/lib/redis'
import { willUseByokForLane } from '@/lib/ai/provider-for-user'
import { reserveCreditsOrThrow } from '@/lib/credits'

function mockActiveSubscription(tier: string) {
  prismaMock.subscription.findUnique.mockResolvedValue({
    userId: 'host-1',
    tier: tier as any,
    status: 'ACTIVE' as any,
    currentPeriodStart: new Date(),
    currentPeriodEnd: new Date(),
  } as any)
}

describe('brainstorm host-pays billing (Story 3.4)', () => {
  beforeEach(() => {
    vi.clearAllMocks()
    vi.mocked(willUseByokForLane).mockResolvedValue({
      providerType: 'openai',
      usedByok: false,
    })
    vi.mocked(reserveCreditsOrThrow).mockResolvedValue({
      cost: 1,
      unlimited: false,
    })
  })

  describe('billingOwnerFromSession', () => {
    it('returns host id for host actor', () => {
      expect(billingOwnerFromSession('host-1', 'host-1')).toEqual({
        billingOwnerId: 'host-1',
        isGuestActor: false,
      })
    })

    it('returns host id for guest actor', () => {
      expect(billingOwnerFromSession('host-1', 'guest-9')).toEqual({
        billingOwnerId: 'host-1',
        isGuestActor: true,
      })
    })
  })

  describe('getBillingOwner', () => {
    it('loads session and resolves billing owner', async () => {
      prismaMock.brainstormSession.findUnique.mockResolvedValue({ userId: 'host-1' })

      const result = await getBillingOwner('session-abc', 'guest-9')
      expect(result).toEqual({ billingOwnerId: 'host-1', isGuestActor: true })
    })
  })

  describe('checkSessionEntitlementOrThrow', () => {
    it('throws QuotaExceededError with guest metadata when host quota exhausted', async () => {
      mockActiveSubscription('BASIC')
      vi.mocked(reserveCreditsOrThrow).mockRejectedValue(
        new QuotaExceededError('PRO', 'brainstorm_expand', 10, 10, false, {
          currentTier: 'BASIC',
        }),
      )

      await expect(
        checkSessionEntitlementOrThrow('host-1', 'guest-9', true, 'brainstorm_expand'),
      ).rejects.toMatchObject({
        billingOwnerId: 'host-1',
        triggeredByUserId: 'guest-9',
        isGuestActor: true,
      })
    })

    it('increments host credits when guest action bills host', async () => {
      mockActiveSubscription('PRO')

      await checkSessionEntitlementOrThrow('host-1', 'guest-9', true, 'brainstorm_expand')

      expect(reserveCreditsOrThrow).toHaveBeenCalledTimes(1)
      const [billedUserId, , options] = vi.mocked(reserveCreditsOrThrow).mock.calls[0]
      expect(billedUserId).toBe('host-1')
      expect(options).toMatchObject({ feature: 'brainstorm_expand' })
    })
  })

  describe('host BYOK bypass (Story 3.5)', () => {
    it('does not debit credits when host has active BYOK', async () => {
      vi.mocked(willUseByokForLane).mockResolvedValue({
        providerType: 'openai',
        usedByok: true,
      })
      mockActiveSubscription('BASIC')
      vi.mocked(reserveCreditsOrThrow).mockRejectedValue(
        new Error('reserveCreditsOrThrow should not be called when BYOK is active'),
      )

      await checkSessionEntitlementOrThrow('host-1', 'guest-9', true, 'brainstorm_expand')

      expect(reserveCreditsOrThrow).not.toHaveBeenCalled()
    })

    it('AC10: guest action bills host, not guest', async () => {
      mockActiveSubscription('PRO')

      await checkSessionEntitlementOrThrow('host-1', 'guest-9', true, 'brainstorm_expand')

      expect(reserveCreditsOrThrow).toHaveBeenCalledTimes(1)
      const [billedUserId] = vi.mocked(reserveCreditsOrThrow).mock.calls[0]
      expect(billedUserId).toBe('host-1')
      expect(billedUserId).not.toBe('guest-9')
    })
  })

  describe('QuotaExceededError.toJSON', () => {
    it('includes isGuestActor in toJSON for 402 responses', () => {
      const err = new QuotaExceededError('PRO', 'brainstorm_expand', 10, 10, false, {
        billingOwnerId: 'host-1',
        triggeredByUserId: 'guest-9',
        isGuestActor: true,
      })
      expect(err.toJSON()).toEqual({
        error: 'QUOTA_EXCEEDED',
        feature: 'brainstorm_expand',
        upgradeTier: 'PRO',
        byokConfigured: false,
        isGuestActor: true,
      })
      expect(err.billingOwnerId).toBe('host-1')
      expect(err.triggeredByUserId).toBe('guest-9')
    })
  })
})