chore: import upstream snapshot with attribution

apps/sim/lib/billing/organizations/member-limits.test.ts

@@ -0,0 +1,161 @@
+/**
+ * @vitest-environment node
+ */
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const {
+  mockDbState,
+  mockInsert,
+  mockInsertValues,
+  mockOnConflictDoUpdate,
+  mockDelete,
+  mockDeleteWhere,
+  mockGetOrganizationSubscription,
+  mockGetOrgWorkspaceUsageCostForUser,
+} = vi.hoisted(() => ({
+  mockDbState: { selectResults: [] as unknown[] },
+  mockInsert: vi.fn(),
+  mockInsertValues: vi.fn(),
+  mockOnConflictDoUpdate: vi.fn(),
+  mockDelete: vi.fn(),
+  mockDeleteWhere: vi.fn(),
+  mockGetOrganizationSubscription: vi.fn(),
+  mockGetOrgWorkspaceUsageCostForUser: vi.fn(),
+}))
+
+vi.mock('@sim/db', () => ({
+  db: {
+    select: vi.fn(() => {
+      const chain: Record<string, unknown> = {}
+      chain.from = vi.fn(() => chain)
+      chain.where = vi.fn(() => chain)
+      chain.limit = vi.fn(() => Promise.resolve(mockDbState.selectResults.shift() ?? []))
+      chain.then = (cb: (rows: unknown) => unknown) =>
+        Promise.resolve(cb(mockDbState.selectResults.shift() ?? []))
+      return chain
+    }),
+    insert: mockInsert,
+    delete: mockDelete,
+  },
+}))
+
+vi.mock('@sim/db/schema', () => ({
+  organizationMemberUsageLimit: {
+    id: 'oml.id',
+    organizationId: 'oml.organizationId',
+    userId: 'oml.userId',
+    usageLimit: 'oml.usageLimit',
+    setBy: 'oml.setBy',
+    createdAt: 'oml.createdAt',
+    updatedAt: 'oml.updatedAt',
+  },
+  workspace: { id: 'workspace.id', organizationId: 'workspace.organizationId' },
+}))
+
+vi.mock('@/lib/billing/core/billing', () => ({
+  getOrganizationSubscription: mockGetOrganizationSubscription,
+}))
+
+vi.mock('@/lib/billing/core/usage-log', () => ({
+  getOrgWorkspaceUsageCostForUser: mockGetOrgWorkspaceUsageCostForUser,
+}))
+
+import { defaultBillingPeriod } from '@/lib/billing/core/billing-period'
+import {
+  getOrgMemberUsageLimit,
+  getOrgMemberWorkspaceUsage,
+  setOrgMemberUsageLimit,
+} from '@/lib/billing/organizations/member-limits'
+
+beforeEach(() => {
+  vi.clearAllMocks()
+  mockDbState.selectResults = []
+  mockInsert.mockReturnValue({ values: mockInsertValues })
+  mockInsertValues.mockReturnValue({ onConflictDoUpdate: mockOnConflictDoUpdate })
+  mockOnConflictDoUpdate.mockResolvedValue(undefined)
+  mockDelete.mockReturnValue({ where: mockDeleteWhere })
+  mockDeleteWhere.mockResolvedValue(undefined)
+})
+
+describe('getOrgMemberUsageLimit', () => {
+  it('returns null when no row exists', async () => {
+    mockDbState.selectResults = [[]]
+    await expect(getOrgMemberUsageLimit('org-1', 'user-2')).resolves.toBeNull()
+  })
+
+  it('returns the stored dollar limit as a number', async () => {
+    mockDbState.selectResults = [[{ usageLimit: '2' }]]
+    await expect(getOrgMemberUsageLimit('org-1', 'user-2')).resolves.toBe(2)
+  })
+})
+
+describe('setOrgMemberUsageLimit', () => {
+  it('upserts when given a dollar amount', async () => {
+    await setOrgMemberUsageLimit('org-1', 'user-2', 2, 'admin-1')
+    expect(mockInsert).toHaveBeenCalledTimes(1)
+    expect(mockDelete).not.toHaveBeenCalled()
+    const values = mockInsertValues.mock.calls[0][0]
+    expect(values).toMatchObject({
+      organizationId: 'org-1',
+      userId: 'user-2',
+      usageLimit: '2',
+      setBy: 'admin-1',
+    })
+    expect(mockOnConflictDoUpdate).toHaveBeenCalledTimes(1)
+  })
+
+  it('deletes the row when limit is null', async () => {
+    await setOrgMemberUsageLimit('org-1', 'user-2', null, 'admin-1')
+    expect(mockDelete).toHaveBeenCalledTimes(1)
+    expect(mockDeleteWhere).toHaveBeenCalledTimes(1)
+    expect(mockInsert).not.toHaveBeenCalled()
+  })
+})
+
+describe('getOrgMemberWorkspaceUsage', () => {
+  it("returns the member's usage within the org subscription window", async () => {
+    const periodStart = new Date('2026-06-01T00:00:00.000Z')
+    const periodEnd = new Date('2026-07-01T00:00:00.000Z')
+    mockGetOrganizationSubscription.mockResolvedValue({ periodStart, periodEnd })
+    mockGetOrgWorkspaceUsageCostForUser.mockResolvedValue(5)
+
+    const result = await getOrgMemberWorkspaceUsage('org-1', 'user-2')
+
+    expect(result).toBe(5)
+    expect(mockGetOrgWorkspaceUsageCostForUser).toHaveBeenCalledWith('org-1', 'user-2', {
+      start: periodStart,
+      end: periodEnd,
+    })
+  })
+
+  it('falls back to the all-time window when the org has no subscription', async () => {
+    mockGetOrganizationSubscription.mockResolvedValue(null)
+    mockGetOrgWorkspaceUsageCostForUser.mockResolvedValue(7)
+
+    const result = await getOrgMemberWorkspaceUsage('org-1', 'user-2')
+
+    expect(result).toBe(7)
+    expect(mockGetOrgWorkspaceUsageCostForUser).toHaveBeenCalledWith(
+      'org-1',
+      'user-2',
+      defaultBillingPeriod()
+    )
+  })
+
+  it('falls back to the all-time window when the subscription is missing periodEnd', async () => {
+    mockGetOrganizationSubscription.mockResolvedValue({
+      periodStart: new Date('2026-06-01T00:00:00.000Z'),
+      periodEnd: null,
+    })
+    mockGetOrgWorkspaceUsageCostForUser.mockResolvedValue(3)
+
+    const result = await getOrgMemberWorkspaceUsage('org-1', 'user-2')
+
+    expect(result).toBe(3)
+    expect(mockGetOrgWorkspaceUsageCostForUser).toHaveBeenCalledWith(
+      'org-1',
+      'user-2',
+      defaultBillingPeriod()
+    )
+  })
+})