diff --git a/apps/sim/app/api/auth/sso/register/route.test.ts b/apps/sim/app/api/auth/sso/register/route.test.ts index cdacc7c6c88..e3539b969e5 100644 --- a/apps/sim/app/api/auth/sso/register/route.test.ts +++ b/apps/sim/app/api/auth/sso/register/route.test.ts @@ -1,8 +1,16 @@ /** * @vitest-environment node */ -import { createEnvMock, createMockRequest } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + createEnvMock, + createMockRequest, + dbChainMock, + dbChainMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, +} from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { mockGetSession, @@ -10,58 +18,29 @@ const { mockHasSSOAccess, mockValidateUrlWithDNS, mockSecureFetchWithPinnedIP, - dbState, - memberTable, - ssoProviderTable, } = vi.hoisted(() => ({ mockGetSession: vi.fn(), mockRegisterSSOProvider: vi.fn(), mockHasSSOAccess: vi.fn(), mockValidateUrlWithDNS: vi.fn(), mockSecureFetchWithPinnedIP: vi.fn(), - dbState: { members: [] as any[], providers: [] as any[] }, - memberTable: { - userId: 'member.userId', - organizationId: 'member.organizationId', - role: 'member.role', - }, - ssoProviderTable: { - id: 'sso.id', - providerId: 'sso.providerId', - domain: 'sso.domain', - issuer: 'sso.issuer', - userId: 'sso.userId', - organizationId: 'sso.organizationId', - oidcConfig: 'sso.oidcConfig', - samlConfig: 'sso.samlConfig', - }, })) -function makeBuilder(rows: any[]): any { - const thenable: any = Promise.resolve(rows) - thenable.where = (condition: any) => { - const values = condition?.values - if (Array.isArray(values) && values.length > 0) { - const target = String(values[values.length - 1]).toLowerCase() - return makeBuilder(rows.filter((r) => String(r.domain ?? '').toLowerCase() === target)) - } - return makeBuilder(rows) - } - thenable.limit = () => Promise.resolve(rows) - thenable.orderBy = () => Promise.resolve(rows) - return thenable +vi.mock('@sim/db', () => ({ ...dbChainMock, ...schemaMock })) + +/** Queues the caller's org membership row(s) for the admin/owner check. */ +function queueMembers(rows: Array>) { + queueTableRows(schemaMock.member, rows) } -vi.mock('@sim/db', () => ({ - db: { - select: () => ({ - from: (table: unknown) => - makeBuilder(table === memberTable ? dbState.members : dbState.providers), - }), - }, - member: memberTable, - ssoProvider: ssoProviderTable, -})) +/** + * Queues existing SSO provider rows for BOTH domain-conflict lookups (the + * pre-registration check and the post-registration re-check). + */ +function queueProviders(rows: Array>) { + queueTableRows(schemaMock.ssoProvider, rows) + queueTableRows(schemaMock.ssoProvider, rows) +} vi.mock('@/lib/auth', () => ({ getSession: mockGetSession, @@ -109,8 +88,7 @@ function request(body: Record) { describe('POST /api/auth/sso/register', () => { beforeEach(() => { vi.clearAllMocks() - dbState.members = [] - dbState.providers = [] + resetDbChainMock() mockGetSession.mockResolvedValue({ user: { id: 'u1' } }) mockHasSSOAccess.mockResolvedValue(true) mockValidateUrlWithDNS.mockResolvedValue({ isValid: true, resolvedIP: '1.2.3.4' }) @@ -118,6 +96,10 @@ describe('POST /api/auth/sso/register', () => { mockRegisterSSOProvider.mockResolvedValue({ providerId: 'acme-oidc' }) }) + afterAll(() => { + resetDbChainMock() + }) + it('rejects callers without an Enterprise plan', async () => { mockHasSSOAccess.mockResolvedValue(false) const res = await POST(request({ ...OIDC_BODY, orgId: 'org1' })) @@ -126,22 +108,22 @@ describe('POST /api/auth/sso/register', () => { }) it('rejects callers who are not an admin/owner of the target org', async () => { - dbState.members = [{ organizationId: 'org1', role: 'member' }] + queueMembers([{ organizationId: 'org1', role: 'member' }]) const res = await POST(request({ ...OIDC_BODY, orgId: 'org1' })) expect(res.status).toBe(403) expect(mockRegisterSSOProvider).not.toHaveBeenCalled() }) it('rejects an invalid domain', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) const res = await POST(request({ ...OIDC_BODY, domain: 'not-a-domain', orgId: 'org1' })) expect(res.status).toBe(400) expect(mockRegisterSSOProvider).not.toHaveBeenCalled() }) it('rejects a domain already registered by another organization', async () => { - dbState.members = [{ organizationId: 'org-attacker', role: 'owner' }] - dbState.providers = [{ domain: 'acme.com', userId: 'u-victim', organizationId: 'org-victim' }] + queueMembers([{ organizationId: 'org-attacker', role: 'owner' }]) + queueProviders([{ domain: 'acme.com', userId: 'u-victim', organizationId: 'org-victim' }]) const res = await POST(request({ ...OIDC_BODY, orgId: 'org-attacker' })) const json = await res.json() expect(res.status).toBe(409) @@ -150,46 +132,51 @@ describe('POST /api/auth/sso/register', () => { }) it('matches conflicts across casing variants', async () => { - dbState.members = [{ organizationId: 'org-attacker', role: 'owner' }] - dbState.providers = [{ domain: 'ACME.com', userId: 'u-victim', organizationId: 'org-victim' }] + queueMembers([{ organizationId: 'org-attacker', role: 'owner' }]) + queueProviders([{ domain: 'ACME.com', userId: 'u-victim', organizationId: 'org-victim' }]) const res = await POST(request({ ...OIDC_BODY, orgId: 'org-attacker' })) expect(res.status).toBe(409) expect(mockRegisterSSOProvider).not.toHaveBeenCalled() + // The conflict lookup itself must be case-insensitive: lower(domain) = . + const conflictWhere = dbChainMockFns.where.mock.calls.find(([condition]) => + condition?.strings?.join('?').includes('lower(') + ) + expect(conflictWhere?.[0]?.values).toContain('acme.com') }) it('registers when the domain is unclaimed', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) const res = await POST(request({ ...OIDC_BODY, orgId: 'org1' })) expect(res.status).toBe(200) expect(mockRegisterSSOProvider).toHaveBeenCalledTimes(1) }) it('allows the owning tenant to update its own provider for the same domain', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] - dbState.providers = [{ domain: 'acme.com', userId: 'u1', organizationId: 'org1' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) + queueProviders([{ domain: 'acme.com', userId: 'u1', organizationId: 'org1' }]) const res = await POST(request({ ...OIDC_BODY, orgId: 'org1' })) expect(res.status).toBe(200) expect(mockRegisterSSOProvider).toHaveBeenCalledTimes(1) }) it('lets an org admin adopt their own user-scoped provider for the same domain', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] - dbState.providers = [{ domain: 'acme.com', userId: 'u1', organizationId: null }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) + queueProviders([{ domain: 'acme.com', userId: 'u1', organizationId: null }]) const res = await POST(request({ ...OIDC_BODY, orgId: 'org1' })) expect(res.status).toBe(200) expect(mockRegisterSSOProvider).toHaveBeenCalledTimes(1) }) it("still blocks an org admin from claiming another user's user-scoped domain", async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] - dbState.providers = [{ domain: 'acme.com', userId: 'someone-else', organizationId: null }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) + queueProviders([{ domain: 'acme.com', userId: 'someone-else', organizationId: null }]) const res = await POST(request({ ...OIDC_BODY, orgId: 'org1' })) expect(res.status).toBe(409) expect(mockRegisterSSOProvider).not.toHaveBeenCalled() }) it('normalizes the domain before persisting it', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) const res = await POST(request({ ...OIDC_BODY, domain: 'ACME.com', orgId: 'org1' })) expect(res.status).toBe(200) expect(mockRegisterSSOProvider).toHaveBeenCalledTimes(1) @@ -198,7 +185,7 @@ describe('POST /api/auth/sso/register', () => { }) it('passes skipDiscovery since Sim already resolved and validated the OIDC endpoints', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) const res = await POST(request({ ...OIDC_BODY, orgId: 'org1' })) expect(res.status).toBe(200) const config = mockRegisterSSOProvider.mock.calls[0][0].body @@ -206,7 +193,7 @@ describe('POST /api/auth/sso/register', () => { }) it('omits userInfoEndpoint when skipUserInfoEndpoint is requested, forcing ID token claims', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) const res = await POST(request({ ...OIDC_BODY, skipUserInfoEndpoint: true, orgId: 'org1' })) expect(res.status).toBe(200) const config = mockRegisterSSOProvider.mock.calls[0][0].body @@ -214,7 +201,7 @@ describe('POST /api/auth/sso/register', () => { }) it('does not SSRF-validate userInfoEndpoint when skipUserInfoEndpoint is requested', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) mockValidateUrlWithDNS.mockImplementation(async (url: string, label: string) => { if (label === 'OIDC userInfoEndpoint') { return { isValid: false, error: 'resolves to a private IP address' } @@ -228,7 +215,7 @@ describe('POST /api/auth/sso/register', () => { }) it('does not SSRF-validate a discovered userinfo_endpoint when skipUserInfoEndpoint is requested', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) mockValidateUrlWithDNS.mockImplementation(async (url: string, label: string) => { if (label === 'OIDC userinfo_endpoint') { return { isValid: false, error: 'resolves to a private IP address' } @@ -258,7 +245,7 @@ describe('POST /api/auth/sso/register', () => { }) it('keeps userInfoEndpoint when skipUserInfoEndpoint is not requested', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) const res = await POST(request({ ...OIDC_BODY, orgId: 'org1' })) expect(res.status).toBe(200) const config = mockRegisterSSOProvider.mock.calls[0][0].body @@ -266,7 +253,7 @@ describe('POST /api/auth/sso/register', () => { }) it('selects tokenEndpointAuthentication from the discovery document when endpoints are auto-discovered', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) mockSecureFetchWithPinnedIP.mockResolvedValue({ ok: true, json: async () => ({ @@ -290,7 +277,7 @@ describe('POST /api/auth/sso/register', () => { }) it('still selects tokenEndpointAuthentication from discovery when all endpoints are explicit', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) mockSecureFetchWithPinnedIP.mockResolvedValue({ ok: true, json: async () => ({ @@ -305,7 +292,7 @@ describe('POST /api/auth/sso/register', () => { }) it('registers successfully when discovery is unreachable and all endpoints are explicit', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) mockSecureFetchWithPinnedIP.mockRejectedValue(new Error('ECONNREFUSED')) const res = await POST(request({ ...OIDC_BODY, orgId: 'org1' })) expect(res.status).toBe(200) @@ -316,7 +303,7 @@ describe('POST /api/auth/sso/register', () => { }) it('prefers client_secret_post over client_secret_basic when an IdP supports both', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) mockSecureFetchWithPinnedIP.mockResolvedValue({ ok: true, json: async () => ({ @@ -330,7 +317,7 @@ describe('POST /api/auth/sso/register', () => { }) it('defaults to client_secret_post when discovery advertises no auth methods', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) mockSecureFetchWithPinnedIP.mockResolvedValue({ ok: true, json: async () => ({}), @@ -342,7 +329,7 @@ describe('POST /api/auth/sso/register', () => { }) it('surfaces the specific discovery failure reason when endpoints are missing', async () => { - dbState.members = [{ organizationId: 'org1', role: 'owner' }] + queueMembers([{ organizationId: 'org1', role: 'owner' }]) mockValidateUrlWithDNS.mockImplementation(async (url: string, label: string) => { if (label === 'OIDC discovery URL') { return { isValid: false, error: 'resolves to a private IP address' } diff --git a/apps/sim/app/api/chat/[identifier]/otp/route.test.ts b/apps/sim/app/api/chat/[identifier]/otp/route.test.ts index 6d32eb61f6b..b67e0942c4c 100644 --- a/apps/sim/app/api/chat/[identifier]/otp/route.test.ts +++ b/apps/sim/app/api/chat/[identifier]/otp/route.test.ts @@ -4,14 +4,19 @@ * @vitest-environment node */ import { + dbChainMock, + dbChainMockFns, + queueTableRows, redisConfigMock, redisConfigMockFns, requestUtilsMockFns, + resetDbChainMock, + schemaMock, workflowsApiUtilsMock, workflowsApiUtilsMockFns, } from '@sim/testing' import { NextRequest } from 'next/server' -import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest' const { mockRedisSet, @@ -20,10 +25,6 @@ const { mockRedisTtl, mockRedisEval, mockRedisClient, - mockDbSelect, - mockDbInsert, - mockDbDelete, - mockDbUpdate, mockSendEmail, mockRenderOTPEmail, mockSetChatAuthCookie, @@ -43,10 +44,6 @@ const { ttl: mockRedisTtl, eval: mockRedisEval, } - const mockDbSelect = vi.fn() - const mockDbInsert = vi.fn() - const mockDbDelete = vi.fn() - const mockDbUpdate = vi.fn() const mockSendEmail = vi.fn() const mockRenderOTPEmail = vi.fn() const mockSetChatAuthCookie = vi.fn() @@ -61,10 +58,6 @@ const { mockRedisTtl, mockRedisEval, mockRedisClient, - mockDbSelect, - mockDbInsert, - mockDbDelete, - mockDbUpdate, mockSendEmail, mockRenderOTPEmail, mockSetChatAuthCookie, @@ -80,30 +73,7 @@ const mockCreateErrorResponse = workflowsApiUtilsMockFns.mockCreateErrorResponse vi.mock('@/lib/core/config/redis', () => redisConfigMock) -vi.mock('@sim/db', () => ({ - db: { - select: mockDbSelect, - insert: mockDbInsert, - delete: mockDbDelete, - update: mockDbUpdate, - transaction: vi.fn(async (callback: (tx: Record) => unknown) => { - return callback({ - select: mockDbSelect, - insert: mockDbInsert, - delete: mockDbDelete, - update: mockDbUpdate, - }) - }), - }, -})) - -vi.mock('drizzle-orm', () => ({ - eq: vi.fn((field: string, value: string) => ({ field, value, type: 'eq' })), - and: vi.fn((...conditions: unknown[]) => ({ conditions, type: 'and' })), - gt: vi.fn((field: string, value: string) => ({ field, value, type: 'gt' })), - lt: vi.fn((field: string, value: string) => ({ field, value, type: 'lt' })), - isNull: vi.fn((field: unknown) => ({ field, type: 'isNull' })), -})) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/core/storage', () => ({ getStorageMethod: mockGetStorageMethod, @@ -201,8 +171,21 @@ describe('Chat OTP API Route', () => { const mockIdentifier = 'test-chat' const mockOTP = '123456' + /** Queues the chat-deployment row the route reads before touching OTP storage. */ + const queueDeployment = (row: Record) => { + queueTableRows(schemaMock.chat, [row]) + } + + const emailDeployment = { + id: mockChatId, + authType: 'email', + allowedEmails: [mockEmail], + title: 'Test Chat', + } + beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() vi.spyOn(Math, 'random').mockReturnValue(0.123456) vi.spyOn(Date, 'now').mockReturnValue(1640995200000) @@ -218,27 +201,6 @@ describe('Chat OTP API Route', () => { mockRedisDel.mockResolvedValue(1) mockRedisTtl.mockResolvedValue(600) - const createDbChain = (result: unknown) => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue(result), - }), - }), - }) - - mockDbSelect.mockImplementation(() => createDbChain([])) - mockDbInsert.mockImplementation(() => ({ - values: vi.fn().mockResolvedValue(undefined), - })) - mockDbDelete.mockImplementation(() => ({ - where: vi.fn().mockResolvedValue(undefined), - })) - mockDbUpdate.mockImplementation(() => ({ - set: vi.fn().mockReturnValue({ - where: vi.fn().mockResolvedValue(undefined), - }), - })) - mockGetStorageMethod.mockReturnValue('redis') mockSendEmail.mockResolvedValue({ success: true }) @@ -271,26 +233,17 @@ describe('Chat OTP API Route', () => { vi.restoreAllMocks() }) + afterAll(() => { + resetDbChainMock() + }) + describe('POST - Store OTP (Redis path)', () => { beforeEach(() => { mockGetStorageMethod.mockReturnValue('redis') }) it('should store OTP in Redis when storage method is redis', async () => { - mockDbSelect.mockImplementationOnce(() => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue([ - { - id: mockChatId, - authType: 'email', - allowedEmails: [mockEmail], - title: 'Test Chat', - }, - ]), - }), - }), - })) + queueDeployment(emailDeployment) const request = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'POST', @@ -306,27 +259,11 @@ describe('Chat OTP API Route', () => { 900 // 15 minutes ) - expect(mockDbInsert).not.toHaveBeenCalled() + expect(dbChainMockFns.insert).not.toHaveBeenCalled() }) }) describe('POST - Rate limiting', () => { - const buildDeploymentSelect = () => - mockDbSelect.mockImplementationOnce(() => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue([ - { - id: mockChatId, - authType: 'email', - allowedEmails: [mockEmail], - title: 'Test Chat', - }, - ]), - }), - }), - })) - it('returns 429 with Retry-After when IP rate limit is exceeded', async () => { mockCheckRateLimitDirect.mockResolvedValueOnce({ allowed: false, @@ -354,7 +291,7 @@ describe('Chat OTP API Route', () => { expect(response.status).toBe(429) expect(headerSet).toHaveBeenCalledWith('Retry-After', '900') expect(mockSendEmail).not.toHaveBeenCalled() - expect(mockDbSelect).not.toHaveBeenCalled() + expect(dbChainMockFns.select).not.toHaveBeenCalled() }) it('returns 429 with Retry-After when email rate limit is exceeded', async () => { @@ -378,7 +315,7 @@ describe('Chat OTP API Route', () => { headers: { set: headerSet }, })) - buildDeploymentSelect() + queueDeployment(emailDeployment) const request = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'POST', @@ -420,7 +357,7 @@ describe('Chat OTP API Route', () => { it('folds spoofed `unknown` client IPs into a single shared bucket', async () => { requestUtilsMockFns.mockGetClientIp.mockReturnValueOnce('unknown') - buildDeploymentSelect() + queueDeployment(emailDeployment) const request = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'POST', @@ -448,30 +385,7 @@ describe('Chat OTP API Route', () => { }) it('should store OTP in database when storage method is database', async () => { - mockDbSelect.mockImplementationOnce(() => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue([ - { - id: mockChatId, - authType: 'email', - allowedEmails: [mockEmail], - title: 'Test Chat', - }, - ]), - }), - }), - })) - - const mockInsertValues = vi.fn().mockResolvedValue(undefined) - mockDbInsert.mockImplementationOnce(() => ({ - values: mockInsertValues, - })) - - const mockDeleteWhere = vi.fn().mockResolvedValue(undefined) - mockDbDelete.mockImplementation(() => ({ - where: mockDeleteWhere, - })) + queueDeployment(emailDeployment) const request = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'POST', @@ -480,10 +394,10 @@ describe('Chat OTP API Route', () => { await POST(request, { params: Promise.resolve({ identifier: mockIdentifier }) }) - expect(mockDbDelete).toHaveBeenCalled() + expect(dbChainMockFns.delete).toHaveBeenCalled() - expect(mockDbInsert).toHaveBeenCalled() - expect(mockInsertValues).toHaveBeenCalledWith({ + expect(dbChainMockFns.insert).toHaveBeenCalled() + expect(dbChainMockFns.values).toHaveBeenCalledWith({ id: expect.any(String), identifier: `chat-otp:${mockChatId}:${mockEmail}`, value: expect.any(String), @@ -503,18 +417,7 @@ describe('Chat OTP API Route', () => { }) it('should retrieve OTP from Redis and verify successfully', async () => { - mockDbSelect.mockImplementationOnce(() => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue([ - { - id: mockChatId, - authType: 'email', - }, - ]), - }), - }), - })) + queueDeployment({ id: mockChatId, authType: 'email' }) const request = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'PUT', @@ -525,7 +428,7 @@ describe('Chat OTP API Route', () => { expect(mockRedisGet).toHaveBeenCalledWith(`otp:${mockEmail}:${mockChatId}`) expect(mockRedisDel).toHaveBeenCalledWith(`otp:${mockEmail}:${mockChatId}`) - expect(mockDbSelect).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(1) }) }) @@ -536,19 +439,11 @@ describe('Chat OTP API Route', () => { }) it('rejects verification when the chat has switched away from email auth', async () => { - mockDbSelect.mockImplementationOnce(() => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue([ - { - id: mockChatId, - authType: 'password', - password: 'encrypted-password', - }, - ]), - }), - }), - })) + queueDeployment({ + id: mockChatId, + authType: 'password', + password: 'encrypted-password', + }) const request = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'PUT', @@ -573,36 +468,13 @@ describe('Chat OTP API Route', () => { }) it('should retrieve OTP from database and verify successfully', async () => { - let selectCallCount = 0 - - mockDbSelect.mockImplementation(() => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockImplementation(() => { - selectCallCount++ - if (selectCallCount === 1) { - return Promise.resolve([ - { - id: mockChatId, - authType: 'email', - }, - ]) - } - return Promise.resolve([ - { - value: `${mockOTP}:0`, - expiresAt: new Date(Date.now() + 10 * 60 * 1000), - }, - ]) - }), - }), - }), - })) - - const mockDeleteWhere = vi.fn().mockResolvedValue(undefined) - mockDbDelete.mockImplementation(() => ({ - where: mockDeleteWhere, - })) + queueDeployment({ id: mockChatId, authType: 'email' }) + queueTableRows(schemaMock.verification, [ + { + value: `${mockOTP}:0`, + expiresAt: new Date(Date.now() + 10 * 60 * 1000), + }, + ]) const request = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'PUT', @@ -611,34 +483,16 @@ describe('Chat OTP API Route', () => { await PUT(request, { params: Promise.resolve({ identifier: mockIdentifier }) }) - expect(mockDbSelect).toHaveBeenCalledTimes(2) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(2) - expect(mockDbDelete).toHaveBeenCalled() + expect(dbChainMockFns.delete).toHaveBeenCalled() expect(mockRedisGet).not.toHaveBeenCalled() }) it('should reject expired OTP from database', async () => { - let selectCallCount = 0 - - mockDbSelect.mockImplementation(() => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockImplementation(() => { - selectCallCount++ - if (selectCallCount === 1) { - return Promise.resolve([ - { - id: mockChatId, - authType: 'email', - }, - ]) - } - return Promise.resolve([]) - }), - }), - }), - })) + queueDeployment({ id: mockChatId, authType: 'email' }) + queueTableRows(schemaMock.verification, []) const request = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'PUT', @@ -662,18 +516,7 @@ describe('Chat OTP API Route', () => { it('should delete OTP from Redis after verification', async () => { mockRedisGet.mockResolvedValue(`${mockOTP}:0`) - mockDbSelect.mockImplementationOnce(() => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue([ - { - id: mockChatId, - authType: 'email', - }, - ]), - }), - }), - })) + queueDeployment({ id: mockChatId, authType: 'email' }) const request = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'PUT', @@ -683,7 +526,7 @@ describe('Chat OTP API Route', () => { await PUT(request, { params: Promise.resolve({ identifier: mockIdentifier }) }) expect(mockRedisDel).toHaveBeenCalledWith(`otp:${mockEmail}:${mockChatId}`) - expect(mockDbDelete).not.toHaveBeenCalled() + expect(dbChainMockFns.delete).not.toHaveBeenCalled() }) }) @@ -694,27 +537,10 @@ describe('Chat OTP API Route', () => { }) it('should delete OTP from database after verification', async () => { - let selectCallCount = 0 - mockDbSelect.mockImplementation(() => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockImplementation(() => { - selectCallCount++ - if (selectCallCount === 1) { - return Promise.resolve([{ id: mockChatId, authType: 'email' }]) - } - return Promise.resolve([ - { value: `${mockOTP}:0`, expiresAt: new Date(Date.now() + 10 * 60 * 1000) }, - ]) - }), - }), - }), - })) - - const mockDeleteWhere = vi.fn().mockResolvedValue(undefined) - mockDbDelete.mockImplementation(() => ({ - where: mockDeleteWhere, - })) + queueDeployment({ id: mockChatId, authType: 'email' }) + queueTableRows(schemaMock.verification, [ + { value: `${mockOTP}:0`, expiresAt: new Date(Date.now() + 10 * 60 * 1000) }, + ]) const request = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'PUT', @@ -723,7 +549,7 @@ describe('Chat OTP API Route', () => { await PUT(request, { params: Promise.resolve({ identifier: mockIdentifier }) }) - expect(mockDbDelete).toHaveBeenCalled() + expect(dbChainMockFns.delete).toHaveBeenCalled() expect(mockRedisDel).not.toHaveBeenCalled() }) }) @@ -737,13 +563,7 @@ describe('Chat OTP API Route', () => { mockRedisGet.mockResolvedValue('654321:0') mockRedisEval.mockResolvedValue('654321:1') - mockDbSelect.mockImplementationOnce(() => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue([{ id: mockChatId, authType: 'email' }]), - }), - }), - })) + queueDeployment({ id: mockChatId, authType: 'email' }) const request = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'PUT', @@ -765,13 +585,7 @@ describe('Chat OTP API Route', () => { mockRedisGet.mockResolvedValue('654321:4') mockRedisEval.mockResolvedValue('LOCKED') - mockDbSelect.mockImplementationOnce(() => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue([{ id: mockChatId, authType: 'email' }]), - }), - }), - })) + queueDeployment({ id: mockChatId, authType: 'email' }) const request = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'PUT', @@ -788,20 +602,7 @@ describe('Chat OTP API Route', () => { }) it('should store OTP with zero attempts on generation', async () => { - mockDbSelect.mockImplementationOnce(() => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue([ - { - id: mockChatId, - authType: 'email', - allowedEmails: [mockEmail], - title: 'Test Chat', - }, - ]), - }), - }), - })) + queueDeployment(emailDeployment) const request = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'POST', @@ -824,13 +625,7 @@ describe('Chat OTP API Route', () => { mockGetStorageMethod.mockReturnValue('redis') mockRedisGet.mockResolvedValue(null) - mockDbSelect.mockImplementation(() => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue([{ id: mockChatId, authType: 'email' }]), - }), - }), - })) + queueDeployment({ id: mockChatId, authType: 'email' }) const requestRedis = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'PUT', @@ -850,20 +645,7 @@ describe('Chat OTP API Route', () => { mockGetStorageMethod.mockReturnValue('redis') - mockDbSelect.mockImplementation(() => ({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue([ - { - id: mockChatId, - authType: 'email', - allowedEmails: [mockEmail], - title: 'Test Chat', - }, - ]), - }), - }), - })) + queueDeployment(emailDeployment) const requestRedis = new NextRequest('http://localhost:3000/api/chat/test/otp', { method: 'POST', diff --git a/apps/sim/app/api/resume/poll/route.test.ts b/apps/sim/app/api/resume/poll/route.test.ts index aee6f33055b..19808c6f2eb 100644 --- a/apps/sim/app/api/resume/poll/route.test.ts +++ b/apps/sim/app/api/resume/poll/route.test.ts @@ -1,13 +1,13 @@ /** * @vitest-environment node */ +import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' import { NextRequest } from 'next/server' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { acquireLockMock, assertBillingAttributionSnapshotMock, - dbSelectMock, dueRowsLimitMock, enqueueOrStartResumeMock, executionSnapshotFromJsonMock, @@ -24,7 +24,6 @@ const { } = vi.hoisted(() => ({ acquireLockMock: vi.fn(), assertBillingAttributionSnapshotMock: vi.fn((value: unknown) => value), - dbSelectMock: vi.fn(), dueRowsLimitMock: vi.fn(), enqueueOrStartResumeMock: vi.fn(), executionSnapshotFromJsonMock: vi.fn(), @@ -42,11 +41,7 @@ const { ), })) -vi.mock('@sim/db', () => ({ - db: { - select: dbSelectMock, - }, -})) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@sim/db/schema', () => ({ pausedExecutions: { @@ -173,7 +168,8 @@ function makeSerializedSnapshot(index: number) { describe('time-pause resume admission', () => { beforeEach(() => { vi.clearAllMocks() - dbSelectMock.mockImplementation((selection: Record) => { + resetDbChainMock() + dbChainMockFns.select.mockImplementation((selection: Record) => { if ('snapshotBytes' in selection) { return { from: vi.fn(() => ({ @@ -219,6 +215,10 @@ describe('time-pause resume admission', () => { executionSnapshotFromJsonMock.mockImplementation((value: string) => JSON.parse(value)) }) + afterAll(() => { + resetDbChainMock() + }) + it('keeps a timed pause unclaimed, records why, and schedules automatic retry', async () => { dueRowsLimitMock.mockResolvedValueOnce([makeDueRow(1)]) preprocessExecutionMock.mockResolvedValueOnce({ @@ -385,14 +385,14 @@ describe('time-pause resume admission', () => { expect(dueRowsLimitMock).toHaveBeenCalledWith(200) expect(legacySizeRowsLimitMock).toHaveBeenCalledWith(200) expect(fallbackRowsLimitMock).toHaveBeenCalledWith(LEGACY_PAUSED_SNAPSHOT_FALLBACK_CHUNK_SIZE) - expect(dbSelectMock).toHaveBeenCalledTimes(3) - expect(dbSelectMock.mock.calls[0]?.[0]).not.toHaveProperty('executionSnapshot') - expect(dbSelectMock.mock.calls[0]?.[0]).toEqual( + expect(dbChainMockFns.select).toHaveBeenCalledTimes(3) + expect(dbChainMockFns.select.mock.calls[0]?.[0]).not.toHaveProperty('executionSnapshot') + expect(dbChainMockFns.select.mock.calls[0]?.[0]).toEqual( expect.objectContaining({ metadata: 'boundedMetadata' }) ) - expect(dbSelectMock.mock.calls[1]?.[0]).toHaveProperty('snapshotBytes') - expect(dbSelectMock.mock.calls[1]?.[0]).not.toHaveProperty('executionSnapshot') - expect(dbSelectMock.mock.calls[2]?.[0]).toHaveProperty('executionSnapshot') + expect(dbChainMockFns.select.mock.calls[1]?.[0]).toHaveProperty('snapshotBytes') + expect(dbChainMockFns.select.mock.calls[1]?.[0]).not.toHaveProperty('executionSnapshot') + expect(dbChainMockFns.select.mock.calls[2]?.[0]).toHaveProperty('executionSnapshot') expect( sqlMock.mock.calls.some(([strings]) => (strings as TemplateStringsArray).join('').includes('octet_length(') @@ -455,9 +455,9 @@ describe('time-pause resume admission', () => { expect(legacySizeRowsLimitMock).toHaveBeenNthCalledWith(1, 200) expect(legacySizeRowsLimitMock).toHaveBeenNthCalledWith(2, 200) expect(fallbackRowsLimitMock).not.toHaveBeenCalled() - expect(dbSelectMock.mock.calls.some(([selection]) => 'executionSnapshot' in selection)).toBe( - false - ) + expect( + dbChainMockFns.select.mock.calls.some(([selection]) => 'executionSnapshot' in selection) + ).toBe(false) expect(executionSnapshotFromJsonMock).not.toHaveBeenCalled() expect(preprocessExecutionMock).not.toHaveBeenCalled() expect(enqueueOrStartResumeMock).not.toHaveBeenCalled() diff --git a/apps/sim/app/api/speech/token/route.test.ts b/apps/sim/app/api/speech/token/route.test.ts index ca31a3cd609..5f42626ba41 100644 --- a/apps/sim/app/api/speech/token/route.test.ts +++ b/apps/sim/app/api/speech/token/route.test.ts @@ -1,15 +1,20 @@ /** * @vitest-environment node */ -import { createMockRequest } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + createMockRequest, + dbChainMock, + queueTableRows, + resetDbChainMock, + schemaMock, +} from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { mockGetSession, mockRecordUsage, mockCheckActorUsageLimits, mockVerifyWorkspaceMembership, - mockChatRows, mockResolveBillingAttribution, mockResolveSystemBillingAttribution, mockCheckAttributedUsageLimits, @@ -20,7 +25,6 @@ const { mockRecordUsage: vi.fn(), mockCheckActorUsageLimits: vi.fn(), mockVerifyWorkspaceMembership: vi.fn(), - mockChatRows: { value: [] as Array> }, mockResolveBillingAttribution: vi.fn(), mockResolveSystemBillingAttribution: vi.fn(), mockCheckAttributedUsageLimits: vi.fn(), @@ -41,18 +45,7 @@ const SYSTEM_BILLING_ATTRIBUTION = { payerSubscription: null, } -vi.mock('@sim/db', () => ({ - db: { - select: () => { - const chain: Record = {} - chain.from = () => chain - chain.leftJoin = () => chain - chain.where = () => chain - chain.limit = () => Promise.resolve(mockChatRows.value) - return chain - }, - }, -})) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/auth', () => ({ getSession: mockGetSession })) @@ -105,7 +98,7 @@ const publicChatRow = { beforeEach(() => { vi.clearAllMocks() - mockChatRows.value = [] + resetDbChainMock() mockGetSession.mockResolvedValue({ user: { id: 'member-1' } }) mockRecordUsage.mockResolvedValue(undefined) mockCheckActorUsageLimits.mockResolvedValue({ isExceeded: false }) @@ -135,6 +128,10 @@ beforeEach(() => { }) as unknown as typeof fetch }) +afterAll(() => { + resetDbChainMock() +}) + describe('POST /api/speech/token — usage attribution', () => { it('editor voice: bills the session user and stamps the verified workspace', async () => { const res = await POST(createMockRequest('POST', { workspaceId: 'ws-1' })) @@ -167,7 +164,7 @@ describe('POST /api/speech/token — usage attribution', () => { }) it('deployed chat: uses one atomic system actor and payer snapshot', async () => { - mockChatRows.value = [publicChatRow] + queueTableRows(schemaMock.chat, [publicChatRow]) const res = await POST(createMockRequest('POST', { chatId: 'chat-1' })) @@ -189,7 +186,7 @@ describe('POST /api/speech/token — usage attribution', () => { }) it('deployed chat: uses the chat owner only when no workspace exists', async () => { - mockChatRows.value = [{ ...publicChatRow, workspaceId: null }] + queueTableRows(schemaMock.chat, [{ ...publicChatRow, workspaceId: null }]) const res = await POST(createMockRequest('POST', { chatId: 'chat-1' })) diff --git a/apps/sim/app/api/tools/custom/route.test.ts b/apps/sim/app/api/tools/custom/route.test.ts index 67e1a186e2e..d84994ab522 100644 --- a/apps/sim/app/api/tools/custom/route.test.ts +++ b/apps/sim/app/api/tools/custom/route.test.ts @@ -6,42 +6,23 @@ import { authMockFns, createMockRequest, + dbChainMock, + dbChainMockFns, hybridAuthMockFns, permissionsMock, permissionsMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, workflowAuthzMockFns, workflowsUtilsMock, } from '@sim/testing' import { NextRequest } from 'next/server' -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const { - mockSelect, - mockFrom, - mockWhere, - mockOrderBy, - mockInsert, - mockValues, - mockUpdate, - mockSet, - mockDelete, - mockLimit, - mockUpsertCustomTools, -} = vi.hoisted(() => { - return { - mockSelect: vi.fn(), - mockFrom: vi.fn(), - mockWhere: vi.fn(), - mockOrderBy: vi.fn(), - mockInsert: vi.fn(), - mockValues: vi.fn(), - mockUpdate: vi.fn(), - mockSet: vi.fn(), - mockDelete: vi.fn(), - mockLimit: vi.fn(), - mockUpsertCustomTools: vi.fn(), - } -}) +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockUpsertCustomTools } = vi.hoisted(() => ({ + mockUpsertCustomTools: vi.fn(), +})) const mockGetUserEntityPermissions = permissionsMockFns.mockGetUserEntityPermissions @@ -102,83 +83,10 @@ const sampleTools = [ }, ] -vi.mock('@sim/db', () => ({ - db: { - select: (...args: unknown[]) => mockSelect(...args), - insert: (...args: unknown[]) => mockInsert(...args), - update: (...args: unknown[]) => mockUpdate(...args), - delete: (...args: unknown[]) => mockDelete(...args), - transaction: vi - .fn() - .mockImplementation(async (callback: (tx: Record) => unknown) => { - const txMockSelect = vi.fn().mockReturnValue({ from: mockFrom }) - const txMockInsert = vi.fn().mockReturnValue({ values: mockValues }) - const txMockUpdate = vi.fn().mockReturnValue({ set: mockSet }) - const txMockDelete = vi.fn().mockReturnValue({ where: mockWhere }) - - const txMockOrderBy = vi.fn().mockImplementation(() => { - const queryBuilder = { - limit: mockLimit, - then: (resolve: (value: typeof sampleTools) => void) => { - resolve(sampleTools) - return queryBuilder - }, - catch: (_reject: (error: Error) => void) => queryBuilder, - } - return queryBuilder - }) - - const txMockWhere = vi.fn().mockImplementation(() => { - const queryBuilder = { - orderBy: txMockOrderBy, - limit: mockLimit, - then: (resolve: (value: typeof sampleTools) => void) => { - resolve(sampleTools) - return queryBuilder - }, - catch: (_reject: (error: Error) => void) => queryBuilder, - } - return queryBuilder - }) - - const txMockFrom = vi.fn().mockReturnValue({ where: txMockWhere }) - txMockSelect.mockReturnValue({ from: txMockFrom }) - - return await callback({ - select: txMockSelect, - insert: txMockInsert, - update: txMockUpdate, - delete: txMockDelete, - }) - }), - }, -})) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/workspaces/permissions/utils', () => permissionsMock) -vi.mock('drizzle-orm', () => ({ - eq: vi.fn().mockImplementation((field: unknown, value: unknown) => ({ - field, - value, - operator: 'eq', - })), - and: vi.fn().mockImplementation((...conditions: unknown[]) => ({ - operator: 'and', - conditions, - })), - or: vi.fn().mockImplementation((...conditions: unknown[]) => ({ - operator: 'or', - conditions, - })), - isNull: vi.fn().mockImplementation((field: unknown) => ({ field, operator: 'isNull' })), - ne: vi.fn().mockImplementation((field: unknown, value: unknown) => ({ - field, - value, - operator: 'ne', - })), - desc: vi.fn().mockImplementation((field: unknown) => ({ field, operator: 'desc' })), -})) - vi.mock('@/lib/workflows/custom-tools/operations', () => ({ upsertCustomTools: (...args: unknown[]) => mockUpsertCustomTools(...args), })) @@ -192,38 +100,7 @@ describe('Custom Tools API Routes', () => { beforeEach(() => { vi.clearAllMocks() - - mockSelect.mockReturnValue({ from: mockFrom }) - mockFrom.mockReturnValue({ where: mockWhere }) - mockWhere.mockImplementation(() => { - const queryBuilder = { - orderBy: mockOrderBy, - limit: mockLimit, - then: (resolve: (value: typeof sampleTools) => void) => { - resolve(sampleTools) - return queryBuilder - }, - catch: (_reject: (error: Error) => void) => queryBuilder, - } - return queryBuilder - }) - mockOrderBy.mockImplementation(() => { - const queryBuilder = { - limit: mockLimit, - then: (resolve: (value: typeof sampleTools) => void) => { - resolve(sampleTools) - return queryBuilder - }, - catch: (_reject: (error: Error) => void) => queryBuilder, - } - return queryBuilder - }) - mockLimit.mockResolvedValue(sampleTools) - mockInsert.mockReturnValue({ values: mockValues }) - mockValues.mockResolvedValue({ id: 'new-tool-id' }) - mockUpdate.mockReturnValue({ set: mockSet }) - mockSet.mockReturnValue({ where: mockWhere }) - mockDelete.mockReturnValue({ where: mockWhere }) + resetDbChainMock() authMockFns.mockGetSession.mockResolvedValue(mockSession) hybridAuthMockFns.mockCheckSessionOrInternalAuth.mockResolvedValue({ @@ -240,6 +117,10 @@ describe('Custom Tools API Routes', () => { }) }) + afterAll(() => { + resetDbChainMock() + }) + /** * Test GET endpoint */ @@ -249,9 +130,7 @@ describe('Custom Tools API Routes', () => { 'http://localhost:3000/api/tools/custom?workspaceId=workspace-123' ) - mockWhere.mockReturnValueOnce({ - orderBy: mockOrderBy.mockReturnValueOnce(Promise.resolve(sampleTools)), - }) + queueTableRows(schemaMock.customTools, sampleTools) const response = await GET(req) const data = await response.json() @@ -260,10 +139,10 @@ describe('Custom Tools API Routes', () => { expect(data).toHaveProperty('data') expect(data.data).toEqual(sampleTools) - expect(mockSelect).toHaveBeenCalled() - expect(mockFrom).toHaveBeenCalled() - expect(mockWhere).toHaveBeenCalled() - expect(mockOrderBy).toHaveBeenCalled() + expect(dbChainMockFns.select).toHaveBeenCalled() + expect(dbChainMockFns.from).toHaveBeenCalled() + expect(dbChainMockFns.where).toHaveBeenCalled() + expect(dbChainMockFns.orderBy).toHaveBeenCalled() }) it('should handle unauthorized access', async () => { @@ -286,13 +165,15 @@ describe('Custom Tools API Routes', () => { it('should handle workflowId parameter', async () => { const req = new NextRequest('http://localhost:3000/api/tools/custom?workflowId=workflow-123') + queueTableRows(schemaMock.customTools, sampleTools) + const response = await GET(req) const data = await response.json() expect(response.status).toBe(200) expect(data).toHaveProperty('data') - expect(mockWhere).toHaveBeenCalled() + expect(dbChainMockFns.where).toHaveBeenCalled() }) }) @@ -336,7 +217,7 @@ describe('Custom Tools API Routes', () => { */ describe('DELETE /api/tools/custom', () => { it('should delete a workspace-scoped tool by ID', async () => { - mockLimit.mockResolvedValueOnce([sampleTools[0]]) + queueTableRows(schemaMock.customTools, [sampleTools[0]]) const req = new NextRequest( 'http://localhost:3000/api/tools/custom?id=tool-1&workspaceId=workspace-123' @@ -348,8 +229,8 @@ describe('Custom Tools API Routes', () => { expect(response.status).toBe(200) expect(data).toHaveProperty('success', true) - expect(mockDelete).toHaveBeenCalled() - expect(mockWhere).toHaveBeenCalled() + expect(dbChainMockFns.delete).toHaveBeenCalled() + expect(dbChainMockFns.where).toHaveBeenCalled() }) it('should reject requests missing tool ID', async () => { @@ -363,8 +244,7 @@ describe('Custom Tools API Routes', () => { }) it('should handle tool not found', async () => { - const mockLimitNotFound = vi.fn().mockResolvedValue([]) - mockWhere.mockReturnValueOnce({ limit: mockLimitNotFound }) + queueTableRows(schemaMock.customTools, []) const req = new NextRequest('http://localhost:3000/api/tools/custom?id=non-existent') @@ -383,8 +263,7 @@ describe('Custom Tools API Routes', () => { }) const userScopedTool = { ...sampleTools[0], workspaceId: null, userId: 'user-123' } - const mockLimitUserScoped = vi.fn().mockResolvedValue([userScopedTool]) - mockWhere.mockReturnValueOnce({ limit: mockLimitUserScoped }) + queueTableRows(schemaMock.customTools, [userScopedTool]) const req = new NextRequest('http://localhost:3000/api/tools/custom?id=tool-1') diff --git a/apps/sim/app/api/v1/audit-logs/auth.test.ts b/apps/sim/app/api/v1/audit-logs/auth.test.ts index 119c7e9c643..d7130879a95 100644 --- a/apps/sim/app/api/v1/audit-logs/auth.test.ts +++ b/apps/sim/app/api/v1/audit-logs/auth.test.ts @@ -1,45 +1,20 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const { - mockIsOrganizationBillingBlocked, - mockMemberWhere, - mockMembersWhere, - mockSelect, - mockSubscriptionWhere, -} = vi.hoisted(() => ({ +import { + dbChainMock, + dbChainMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, +} from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockIsOrganizationBillingBlocked } = vi.hoisted(() => ({ mockIsOrganizationBillingBlocked: vi.fn(), - mockMemberWhere: vi.fn(), - mockMembersWhere: vi.fn(), - mockSelect: vi.fn(), - mockSubscriptionWhere: vi.fn(), -})) - -vi.mock('@sim/db', () => ({ - db: { select: mockSelect }, -})) - -vi.mock('@sim/db/schema', () => ({ - member: { - organizationId: 'member.organizationId', - role: 'member.role', - userId: 'member.userId', - }, - subscription: { - id: 'subscription.id', - plan: 'subscription.plan', - referenceId: 'subscription.referenceId', - status: 'subscription.status', - }, })) -vi.mock('drizzle-orm', () => ({ - and: vi.fn((...conditions: unknown[]) => ({ type: 'and', conditions })), - eq: vi.fn((left: unknown, right: unknown) => ({ type: 'eq', left, right })), - inArray: vi.fn((left: unknown, right: unknown) => ({ type: 'inArray', left, right })), -})) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/billing/core/access', () => ({ isOrganizationBillingBlocked: mockIsOrganizationBillingBlocked, @@ -50,18 +25,15 @@ import { validateEnterpriseAuditAccess } from '@/app/api/v1/audit-logs/auth' describe('enterprise audit access', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockIsOrganizationBillingBlocked.mockResolvedValue(false) - mockMemberWhere.mockReturnValue({ - limit: vi.fn().mockResolvedValue([{ organizationId: 'organization-route', role: 'admin' }]), - }) - mockSubscriptionWhere.mockReturnValue({ - limit: vi.fn().mockResolvedValue([{ id: 'subscription-1' }]), - }) - mockMembersWhere.mockResolvedValue([{ userId: 'viewer' }, { userId: 'member-2' }]) - mockSelect - .mockReturnValueOnce({ from: () => ({ where: mockMemberWhere }) }) - .mockReturnValueOnce({ from: () => ({ where: mockSubscriptionWhere }) }) - .mockReturnValueOnce({ from: () => ({ where: mockMembersWhere }) }) + queueTableRows(schemaMock.member, [{ organizationId: 'organization-route', role: 'admin' }]) + queueTableRows(schemaMock.subscription, [{ id: 'subscription-1' }]) + queueTableRows(schemaMock.member, [{ userId: 'viewer' }, { userId: 'member-2' }]) + }) + + afterAll(() => { + resetDbChainMock() }) it('authorizes and bills against the organization named by the route', async () => { @@ -72,11 +44,11 @@ describe('enterprise audit access', () => { orgMemberIds: ['viewer', 'member-2'], }, }) - expect(mockMemberWhere).toHaveBeenCalledWith({ + expect(dbChainMockFns.where).toHaveBeenNthCalledWith(1, { type: 'and', conditions: [ - { type: 'eq', left: 'member.userId', right: 'viewer' }, - { type: 'eq', left: 'member.organizationId', right: 'organization-route' }, + { type: 'eq', left: schemaMock.member.userId, right: 'viewer' }, + { type: 'eq', left: schemaMock.member.organizationId, right: 'organization-route' }, ], }) expect(mockIsOrganizationBillingBlocked).toHaveBeenCalledWith('organization-route') diff --git a/apps/sim/executor/handlers/agent/agent-handler.test.ts b/apps/sim/executor/handlers/agent/agent-handler.test.ts index f733c885944..a5a5b36b410 100644 --- a/apps/sim/executor/handlers/agent/agent-handler.test.ts +++ b/apps/sim/executor/handlers/agent/agent-handler.test.ts @@ -1,4 +1,5 @@ -import { afterEach, beforeEach, describe, expect, it, type Mock, vi } from 'vitest' +import { dbChainMock, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing' +import { afterAll, afterEach, beforeEach, describe, expect, it, type Mock, vi } from 'vitest' import { getAllBlocks } from '@/blocks' import { BlockType, isMcpTool } from '@/executor/constants' import { AgentBlockHandler } from '@/executor/handlers/agent/agent-handler' @@ -87,19 +88,14 @@ vi.mock('@/executor/utils/http', () => ({ }), })) -vi.mock('@sim/db', () => ({ - db: { - select: vi.fn().mockReturnValue({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockResolvedValue([ - { id: 'mcp-search-server', connectionStatus: 'connected' }, - { id: 'same-server', connectionStatus: 'connected' }, - { id: 'mcp-legacy-server', connectionStatus: 'connected' }, - ]), - }), - }), - }, -})) +vi.mock('@sim/db', () => dbChainMock) + +/** Connected MCP servers every workspace-server lookup in this suite resolves. */ +const MCP_SERVER_ROWS = [ + { id: 'mcp-search-server', connectionStatus: 'connected' }, + { id: 'same-server', connectionStatus: 'connected' }, + { id: 'mcp-legacy-server', connectionStatus: 'connected' }, +] const mockGetCustomToolById = vi.fn() @@ -122,6 +118,10 @@ describe('AgentBlockHandler', () => { beforeEach(() => { handler = new AgentBlockHandler() vi.clearAllMocks() + resetDbChainMock() + // The MCP server lookup awaits select().from(mcpServers).where(...) directly; + // queue a set per lookup so the structural where spy keeps its default wiring. + queueTableRows(schemaMock.mcpServers, MCP_SERVER_ROWS) // unstubGlobals removes any module-scope fetch stub before each test, so re-stub here vi.stubGlobal('fetch', mockFetch) @@ -213,6 +213,10 @@ describe('AgentBlockHandler', () => { } catch (e) {} }) + afterAll(() => { + resetDbChainMock() + }) + describe('canHandle', () => { it('should return true for blocks with metadata id "agent"', () => { expect(handler.canHandle(mockBlock)).toBe(true) diff --git a/apps/sim/executor/handlers/agent/skills-resolver.test.ts b/apps/sim/executor/handlers/agent/skills-resolver.test.ts index e5e74ba4222..e562b63b7ec 100644 --- a/apps/sim/executor/handlers/agent/skills-resolver.test.ts +++ b/apps/sim/executor/handlers/agent/skills-resolver.test.ts @@ -1,22 +1,16 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const { limitMock } = vi.hoisted(() => ({ limitMock: vi.fn() })) - -vi.mock('@sim/db', () => ({ - db: { select: () => ({ from: () => ({ where: () => ({ limit: limitMock }) }) }) }, - skill: { workspaceId: 'workspaceId', name: 'name', content: 'content' }, -})) -vi.mock('@sim/logger', () => ({ - createLogger: () => ({ error: vi.fn(), warn: vi.fn(), info: vi.fn(), debug: vi.fn() }), -})) -vi.mock('drizzle-orm', () => ({ - and: vi.fn(() => ({})), - eq: vi.fn(() => ({})), - inArray: vi.fn(() => ({})), -})) +import { + dbChainMock, + dbChainMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, +} from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' + +vi.mock('@sim/db', () => dbChainMock) import { resolveSkillContent } from './skills-resolver' @@ -25,6 +19,11 @@ import { resolveSkillContent } from './skills-resolver' describe('resolveSkillContent', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() + }) + + afterAll(() => { + resetDbChainMock() }) it('returns null without a skill name or workspace', async () => { @@ -35,16 +34,15 @@ describe('resolveSkillContent', () => { it('resolves builtin skills without touching the database', async () => { const content = await resolveSkillContent('research', 'ws-1') expect(content).toBeTruthy() - expect(limitMock).not.toHaveBeenCalled() + expect(dbChainMockFns.limit).not.toHaveBeenCalled() }) it('resolves a workspace user skill by name', async () => { - limitMock.mockResolvedValue([{ content: '# Playbook', name: 'posthog-playbook' }]) + queueTableRows(schemaMock.skill, [{ content: '# Playbook', name: 'posthog-playbook' }]) expect(await resolveSkillContent('posthog-playbook', 'ws-1')).toBe('# Playbook') }) it('returns null when the user skill is not found', async () => { - limitMock.mockResolvedValue([]) expect(await resolveSkillContent('missing', 'ws-1')).toBeNull() }) }) diff --git a/apps/sim/lib/billing/authorization.test.ts b/apps/sim/lib/billing/authorization.test.ts index 002bb2a4537..457564f13a8 100644 --- a/apps/sim/lib/billing/authorization.test.ts +++ b/apps/sim/lib/billing/authorization.test.ts @@ -1,7 +1,8 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { dbChainMock, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { mockHasPaidSubscription, @@ -15,7 +16,7 @@ const { mockGetOrganizationCoverageForMember: vi.fn(), })) -vi.mock('@sim/db', () => ({ db: {} })) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/billing', () => ({ hasPaidSubscription: mockHasPaidSubscription })) vi.mock('@/lib/billing/core/organization', () => ({ isOrganizationOwnerOrAdmin: mockIsOwnerOrAdmin, @@ -42,6 +43,14 @@ import { } from '@/lib/billing/authorization' import { EnterpriseIssuanceInProgressError } from '@/lib/billing/enterprise-outbox' +beforeEach(() => { + resetDbChainMock() +}) + +afterAll(() => { + resetDbChainMock() +}) + describe('isPersonalCheckoutRequest', () => { it('classifies an explicit self reference as personal regardless of customerType', () => { expect(isPersonalCheckoutRequest({ referenceId: 'user-1' }, 'user-1')).toBe(true) diff --git a/apps/sim/lib/billing/calculations/usage-monitor.test.ts b/apps/sim/lib/billing/calculations/usage-monitor.test.ts index 4f57be3b627..b1a498cd455 100644 --- a/apps/sim/lib/billing/calculations/usage-monitor.test.ts +++ b/apps/sim/lib/billing/calculations/usage-monitor.test.ts @@ -1,17 +1,16 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { mockFlags, - mockDbLimit, mockGetOrgMemberUsageForBillingPeriod, mockGetOrgMemberUsageLimit, mockIsOrganizationBillingBlocked, } = vi.hoisted(() => ({ mockFlags: { isHosted: true, isBillingEnabled: true }, - mockDbLimit: vi.fn(), mockGetOrgMemberUsageForBillingPeriod: vi.fn(), mockGetOrgMemberUsageLimit: vi.fn(), mockIsOrganizationBillingBlocked: vi.fn(), @@ -26,17 +25,7 @@ vi.mock('@/lib/core/config/env-flags', () => ({ }, })) -vi.mock('@sim/db', () => ({ - db: { - select: () => ({ - from: () => ({ - where: () => ({ - limit: mockDbLimit, - }), - }), - }), - }, -})) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/billing/organizations/member-limits', () => ({ getOrgMemberUsageForBillingPeriod: mockGetOrgMemberUsageForBillingPeriod, @@ -60,12 +49,17 @@ import { checkOrganizationMemberUsageLimit, } from '@/lib/billing/calculations/usage-monitor' +afterAll(() => { + resetDbChainMock() +}) + describe('checkBillingBlocked', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockFlags.isHosted = true mockFlags.isBillingEnabled = true - mockDbLimit.mockResolvedValue([{ blocked: false, blockedReason: null }]) + dbChainMockFns.limit.mockResolvedValue([{ blocked: false, blockedReason: null }]) }) it("checks only the actor's own user account without inspecting organization memberships", async () => { @@ -73,7 +67,7 @@ describe('checkBillingBlocked', () => { await expect(checkBillingBlocked('actor-1')).resolves.toEqual({ blocked: false }) - expect(mockDbLimit).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.limit).toHaveBeenCalledTimes(1) expect(mockIsOrganizationBillingBlocked).not.toHaveBeenCalled() }) }) @@ -81,10 +75,11 @@ describe('checkBillingBlocked', () => { describe('checkBillingEntityBlocked', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockFlags.isHosted = true mockFlags.isBillingEnabled = true mockIsOrganizationBillingBlocked.mockResolvedValue(false) - mockDbLimit.mockResolvedValue([]) + dbChainMockFns.limit.mockResolvedValue([]) }) it('checks only the exact organization payer', async () => { @@ -95,11 +90,11 @@ describe('checkBillingEntityBlocked', () => { ).resolves.toMatchObject({ blocked: true }) expect(mockIsOrganizationBillingBlocked).toHaveBeenCalledWith('workspace-org') - expect(mockDbLimit).not.toHaveBeenCalled() + expect(dbChainMockFns.limit).not.toHaveBeenCalled() }) it('checks the exact personal payer directly', async () => { - mockDbLimit.mockResolvedValue([{ blocked: true, blockedReason: 'dispute' }]) + dbChainMockFns.limit.mockResolvedValue([{ blocked: true, blockedReason: 'dispute' }]) await expect( checkBillingEntityBlocked({ type: 'user', id: 'personal-payer' }) @@ -120,6 +115,7 @@ describe('checkOrganizationMemberUsageLimit', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockFlags.isHosted = true mockFlags.isBillingEnabled = true mockGetOrgMemberUsageLimit.mockResolvedValue(2) diff --git a/apps/sim/lib/billing/cleanup-dispatcher.test.ts b/apps/sim/lib/billing/cleanup-dispatcher.test.ts index e696d221895..ed126c26535 100644 --- a/apps/sim/lib/billing/cleanup-dispatcher.test.ts +++ b/apps/sim/lib/billing/cleanup-dispatcher.test.ts @@ -1,15 +1,15 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' -const { mockFlags, mockIsTriggerAvailable, mockSelect } = vi.hoisted(() => ({ +const { mockFlags, mockIsTriggerAvailable } = vi.hoisted(() => ({ mockFlags: { isBillingEnabled: false }, mockIsTriggerAvailable: vi.fn(), - mockSelect: vi.fn(), })) -vi.mock('@sim/db', () => ({ db: { select: mockSelect } })) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/billing/core/billing', () => ({ getOrganizationSubscription: vi.fn() })) vi.mock('@/lib/billing/core/subscription', () => ({ getHighestPriorityPersonalSubscription: vi.fn(), @@ -36,14 +36,19 @@ import { dispatchCleanupJobs } from '@/lib/billing/cleanup-dispatcher' describe('dispatchCleanupJobs billing gate', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockFlags.isBillingEnabled = false }) + afterAll(() => { + resetDbChainMock() + }) + it('never dispatches plan-based retention deletion when billing is disabled', async () => { const result = await dispatchCleanupJobs('cleanup-logs') expect(result).toEqual({ jobIds: [], jobCount: 0, chunkCount: 0, workspaceCount: 0 }) expect(mockIsTriggerAvailable).not.toHaveBeenCalled() - expect(mockSelect).not.toHaveBeenCalled() + expect(dbChainMockFns.select).not.toHaveBeenCalled() }) }) diff --git a/apps/sim/lib/billing/core/billing-attribution.test.ts b/apps/sim/lib/billing/core/billing-attribution.test.ts index fb6105a6bc8..05f42edf884 100644 --- a/apps/sim/lib/billing/core/billing-attribution.test.ts +++ b/apps/sim/lib/billing/core/billing-attribution.test.ts @@ -1,7 +1,8 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { mockFlags, @@ -11,7 +12,6 @@ const { mockCheckUsageStatus, mockGetHighestPriorityPersonalSubscription, mockGetOrganizationSubscription, - mockLimit, } = vi.hoisted(() => ({ mockFlags: { isBillingEnabled: true, isHosted: true }, mockCheckBillingBlocked: vi.fn(), @@ -20,7 +20,6 @@ const { mockCheckUsageStatus: vi.fn(), mockGetHighestPriorityPersonalSubscription: vi.fn(), mockGetOrganizationSubscription: vi.fn(), - mockLimit: vi.fn(), })) vi.mock('@/lib/core/config/env-flags', () => ({ @@ -32,17 +31,7 @@ vi.mock('@/lib/core/config/env-flags', () => ({ }, })) -vi.mock('@sim/db', () => ({ - db: { - select: vi.fn(() => ({ - from: vi.fn(() => ({ - where: vi.fn(() => ({ - limit: mockLimit, - })), - })), - })), - }, -})) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/billing/calculations/usage-monitor', () => ({ checkBillingBlocked: mockCheckBillingBlocked, @@ -76,6 +65,10 @@ import { toBillingContext, } from '@/lib/billing/core/billing-attribution' +afterAll(() => { + resetDbChainMock() +}) + const ORG_SUBSCRIPTION = { id: 'sub-org-b', plan: 'team_25000', @@ -89,6 +82,7 @@ const ORG_SUBSCRIPTION = { describe('resolveBillingAttribution', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockCheckBillingBlocked.mockResolvedValue({ blocked: false }) mockCheckBillingEntityBlocked.mockResolvedValue({ blocked: false }) mockCheckUsageStatus.mockResolvedValue({ @@ -108,7 +102,7 @@ describe('resolveBillingAttribution', () => { }) it('bills the workspace organization while retaining an external session actor', async () => { - mockLimit.mockResolvedValue([ + dbChainMockFns.limit.mockResolvedValue([ { billedAccountUserId: 'owner-b', organizationId: 'org-b', @@ -150,7 +144,7 @@ describe('resolveBillingAttribution', () => { }) it('resolves the system actor and payer from one workspace row', async () => { - mockLimit.mockResolvedValue([ + dbChainMockFns.limit.mockResolvedValue([ { billedAccountUserId: 'owner-b', organizationId: 'org-b', @@ -167,11 +161,11 @@ describe('resolveBillingAttribution', () => { organizationId: 'org-b', workspaceId: 'workspace-b', }) - expect(mockLimit).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.limit).toHaveBeenCalledTimes(1) }) it('uses the workspace organization reference even when its billed owner has other memberships', async () => { - mockLimit.mockResolvedValue([ + dbChainMockFns.limit.mockResolvedValue([ { billedAccountUserId: 'multi-org-owner', organizationId: 'org-b', @@ -192,7 +186,7 @@ describe('resolveBillingAttribution', () => { }) it('bills a personal workspace billed account without changing the API-key actor', async () => { - mockLimit.mockResolvedValue([ + dbChainMockFns.limit.mockResolvedValue([ { billedAccountUserId: 'personal-owner', organizationId: null, @@ -222,7 +216,7 @@ describe('resolveBillingAttribution', () => { }) it('retains the exact personal payer when it has no subscription', async () => { - mockLimit.mockResolvedValue([ + dbChainMockFns.limit.mockResolvedValue([ { billedAccountUserId: 'personal-owner', organizationId: null, @@ -245,7 +239,7 @@ describe('resolveBillingAttribution', () => { }) it('serializes only the payer fields needed by later billing gates', async () => { - mockLimit.mockResolvedValue([ + dbChainMockFns.limit.mockResolvedValue([ { billedAccountUserId: 'owner-b', organizationId: 'org-b', @@ -276,7 +270,7 @@ describe('resolveBillingAttribution', () => { }) it('carries only the normalized Enterprise concurrency metadata needed by admission', async () => { - mockLimit.mockResolvedValue([ + dbChainMockFns.limit.mockResolvedValue([ { billedAccountUserId: 'owner-b', organizationId: 'org-b', @@ -301,7 +295,7 @@ describe('resolveBillingAttribution', () => { }) it('rejects a subscription that does not belong to the exact workspace payer', async () => { - mockLimit.mockResolvedValue([ + dbChainMockFns.limit.mockResolvedValue([ { billedAccountUserId: 'owner-b', organizationId: 'org-b', @@ -321,7 +315,7 @@ describe('resolveBillingAttribution', () => { }) it('fails closed when the workspace payer cannot be resolved', async () => { - mockLimit.mockResolvedValue([]) + dbChainMockFns.limit.mockResolvedValue([]) await expect( resolveBillingAttribution({ @@ -332,7 +326,7 @@ describe('resolveBillingAttribution', () => { }) it('resolves markerless legacy-v0 from the current workspace payer', async () => { - mockLimit.mockResolvedValue([ + dbChainMockFns.limit.mockResolvedValue([ { billedAccountUserId: 'owner-b', organizationId: 'org-b', @@ -354,7 +348,7 @@ describe('resolveBillingAttribution', () => { }) it('returns no workspace payer for an opaque markerless legacy-v0 workspace', async () => { - mockLimit.mockResolvedValue([]) + dbChainMockFns.limit.mockResolvedValue([]) await expect( resolveLegacyV0BillingAttribution({ @@ -367,7 +361,7 @@ describe('resolveBillingAttribution', () => { }) it('converts the serialized period back to the exact runtime billing context', async () => { - mockLimit.mockResolvedValue([ + dbChainMockFns.limit.mockResolvedValue([ { billedAccountUserId: 'owner-b', organizationId: 'org-b', @@ -485,6 +479,7 @@ describe('serialized attribution boundaries', () => { describe('checkAttributedUsageLimits', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockFlags.isBillingEnabled = true mockFlags.isHosted = true mockCheckBillingBlocked.mockResolvedValue({ blocked: false }) diff --git a/apps/sim/lib/billing/core/limit-notifications.test.ts b/apps/sim/lib/billing/core/limit-notifications.test.ts index f3486fdd9f0..5c82b418b72 100644 --- a/apps/sim/lib/billing/core/limit-notifications.test.ts +++ b/apps/sim/lib/billing/core/limit-notifications.test.ts @@ -1,13 +1,17 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + dbChainMock, + dbChainMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, +} from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { billingFlag, - mockClaim, - mockSelectRows, - dbUpdateSpy, sendEmailSpy, getEmailPreferencesMock, renderMock, @@ -15,9 +19,6 @@ const { isOrgAdminRoleMock, } = vi.hoisted(() => ({ billingFlag: { enabled: true }, - mockClaim: vi.fn<[], unknown[]>(() => [{ id: 'u1' }]), - mockSelectRows: vi.fn<[], unknown[]>(() => []), - dbUpdateSpy: vi.fn(), sendEmailSpy: vi.fn(() => Promise.resolve({ success: true })), getEmailPreferencesMock: vi.fn(() => Promise.resolve(null as unknown)), renderMock: vi.fn(() => Promise.resolve('')), @@ -25,26 +26,7 @@ const { isOrgAdminRoleMock: vi.fn(() => true), })) -vi.mock('@sim/db', () => { - const updateBuilder: Record = { - set: () => updateBuilder, - where: () => updateBuilder, - returning: () => Promise.resolve(mockClaim()), - then: (f: (v: unknown) => unknown, r?: (e: unknown) => unknown) => - Promise.resolve(undefined).then(f, r), - } - const selectBuilder: Record = { - from: () => selectBuilder, - where: () => selectBuilder, - innerJoin: () => selectBuilder, - leftJoin: () => selectBuilder, - limit: () => Promise.resolve(mockSelectRows()), - then: (f: (v: unknown) => unknown, r?: (e: unknown) => unknown) => - Promise.resolve(mockSelectRows()).then(f, r), - } - dbUpdateSpy.mockImplementation(() => updateBuilder) - return { db: { update: dbUpdateSpy, select: () => selectBuilder } } -}) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/core/config/env-flags', () => ({ get isBillingEnabled() { @@ -78,12 +60,16 @@ const baseUserParams = { describe('maybeSendLimitThresholdEmail', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() billingFlag.enabled = true - mockClaim.mockReturnValue([{ id: 'u1' }]) - mockSelectRows.mockReturnValue([]) + dbChainMockFns.returning.mockResolvedValue([{ id: 'u1' }]) getEmailPreferencesMock.mockResolvedValue(null) }) + afterAll(() => { + resetDbChainMock() + }) + it('sends a warning email when crossing 80% and the claim wins', async () => { await maybeSendLimitThresholdEmail({ ...baseUserParams, currentUsage: 4.5, limit: 5 }) expect(sendEmailSpy).toHaveBeenCalledTimes(1) @@ -104,66 +90,66 @@ describe('maybeSendLimitThresholdEmail', () => { limit: 5, rearmOnly: true, }) - expect(mockClaim).not.toHaveBeenCalled() + expect(dbChainMockFns.returning).not.toHaveBeenCalled() expect(sendEmailSpy).not.toHaveBeenCalled() }) it('does not send when the atomic claim is lost (already notified)', async () => { - mockClaim.mockReturnValue([]) + dbChainMockFns.returning.mockResolvedValue([]) await maybeSendLimitThresholdEmail({ ...baseUserParams, currentUsage: 4.5, limit: 5 }) expect(sendEmailSpy).not.toHaveBeenCalled() }) it('claims without re-arming on a crossing (re-arm and claim are mutually exclusive)', async () => { await maybeSendLimitThresholdEmail({ ...baseUserParams, currentUsage: 4.5, limit: 5 }) - expect(dbUpdateSpy).toHaveBeenCalledTimes(1) - expect(mockClaim).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.update).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.returning).toHaveBeenCalledTimes(1) expect(sendEmailSpy).toHaveBeenCalledTimes(1) }) it('does not send in the dead band (70%–80%)', async () => { await maybeSendLimitThresholdEmail({ ...baseUserParams, currentUsage: 3.75, limit: 5 }) - expect(mockClaim).not.toHaveBeenCalled() + expect(dbChainMockFns.returning).not.toHaveBeenCalled() expect(sendEmailSpy).not.toHaveBeenCalled() }) it('re-arms below the band without claiming or sending', async () => { await maybeSendLimitThresholdEmail({ ...baseUserParams, currentUsage: 1, limit: 5 }) - expect(mockClaim).not.toHaveBeenCalled() + expect(dbChainMockFns.returning).not.toHaveBeenCalled() expect(sendEmailSpy).not.toHaveBeenCalled() }) it('does not send OR burn the claim when the per-user toggle is off', async () => { - mockSelectRows.mockReturnValue([{ enabled: false }]) + queueTableRows(schemaMock.settings, [{ enabled: false }]) await maybeSendLimitThresholdEmail({ ...baseUserParams, currentUsage: 4.5, limit: 5 }) expect(sendEmailSpy).not.toHaveBeenCalled() - expect(mockClaim).not.toHaveBeenCalled() + expect(dbChainMockFns.returning).not.toHaveBeenCalled() }) it('does not send OR burn the claim when the recipient unsubscribed', async () => { getEmailPreferencesMock.mockResolvedValue({ unsubscribeNotifications: true }) await maybeSendLimitThresholdEmail({ ...baseUserParams, currentUsage: 4.5, limit: 5 }) expect(sendEmailSpy).not.toHaveBeenCalled() - expect(mockClaim).not.toHaveBeenCalled() + expect(dbChainMockFns.returning).not.toHaveBeenCalled() }) it('skips entirely when billing is disabled', async () => { billingFlag.enabled = false await maybeSendLimitThresholdEmail({ ...baseUserParams, currentUsage: 5, limit: 5 }) - expect(mockClaim).not.toHaveBeenCalled() + expect(dbChainMockFns.returning).not.toHaveBeenCalled() expect(sendEmailSpy).not.toHaveBeenCalled() }) it('re-arms but does not send when usage is fully cleared (zero usage)', async () => { await maybeSendLimitThresholdEmail({ ...baseUserParams, currentUsage: 0, limit: 5 }) - expect(dbUpdateSpy).toHaveBeenCalledTimes(1) - expect(mockClaim).not.toHaveBeenCalled() + expect(dbChainMockFns.update).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.returning).not.toHaveBeenCalled() expect(sendEmailSpy).not.toHaveBeenCalled() }) it('skips when the limit is non-positive', async () => { await maybeSendLimitThresholdEmail({ ...baseUserParams, currentUsage: 4, limit: 0 }) - expect(dbUpdateSpy).not.toHaveBeenCalled() + expect(dbChainMockFns.update).not.toHaveBeenCalled() expect(sendEmailSpy).not.toHaveBeenCalled() }) }) diff --git a/apps/sim/lib/billing/enterprise-provisioning.test.ts b/apps/sim/lib/billing/enterprise-provisioning.test.ts index d01e15d1964..d50dd0a8088 100644 --- a/apps/sim/lib/billing/enterprise-provisioning.test.ts +++ b/apps/sim/lib/billing/enterprise-provisioning.test.ts @@ -1,11 +1,16 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + dbChainMock, + dbChainMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, +} from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const mocks = vi.hoisted(() => ({ - select: vi.fn(), - update: vi.fn(), subscriptionsCreate: vi.fn(), subscriptionsList: vi.fn(), subscriptionsUpdate: vi.fn(), @@ -27,49 +32,9 @@ vi.mock('@sim/audit', () => ({ recordAudit: vi.fn(), })) -vi.mock('@sim/db', () => ({ - db: { - select: mocks.select, - update: mocks.update, - transaction: vi.fn(), - }, -})) - -vi.mock('@sim/db/schema', () => ({ - member: { userId: 'userId', organizationId: 'organizationId', role: 'role' }, - organization: { id: 'id', name: 'name' }, - outboxEvent: { - id: 'id', - eventType: 'eventType', - payload: 'payload', - status: 'status', - createdAt: 'createdAt', - }, - subscription: { - id: 'id', - referenceId: 'referenceId', - status: 'status', - stripeSubscriptionId: 'stripeSubscriptionId', - metadata: 'metadata', - }, - user: { - id: 'id', - name: 'name', - email: 'email', - stripeCustomerId: 'stripeCustomerId', - }, -})) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@sim/utils/id', () => ({ generateId: vi.fn(() => 'generated-id') })) -vi.mock('drizzle-orm', () => ({ - and: vi.fn(() => 'and'), - count: vi.fn(() => 'count'), - desc: vi.fn(() => 'desc'), - eq: vi.fn(() => 'eq'), - inArray: vi.fn(() => 'inArray'), - isNull: vi.fn(() => 'isNull'), - sql: vi.fn(() => 'sql'), -})) vi.mock('@/lib/billing/organizations/membership', () => ({ acquireOrganizationMutationLock: vi.fn(), })) @@ -107,26 +72,9 @@ import { syncEnterpriseMetadataInStripe, } from '@/lib/billing/enterprise-provisioning' -function selectChain(rows: unknown[]) { - const chain = { - from: vi.fn(), - innerJoin: vi.fn(), - leftJoin: vi.fn(), - where: vi.fn(), - orderBy: vi.fn(), - for: vi.fn(), - limit: vi.fn().mockResolvedValue(rows), - then: (resolve: (value: unknown[]) => unknown, reject: (reason: unknown) => unknown) => - Promise.resolve(rows).then(resolve, reject), - } - chain.from.mockReturnValue(chain) - chain.innerJoin.mockReturnValue(chain) - chain.leftJoin.mockReturnValue(chain) - chain.where.mockReturnValue(chain) - chain.orderBy.mockReturnValue(chain) - chain.for.mockReturnValue(chain) - return chain -} +afterAll(() => { + resetDbChainMock() +}) function operationPayload(overrides: Record = {}) { return { @@ -256,29 +204,26 @@ function arrangeWorkerReads( finalLocalSubscriptions: unknown[] = localSubscriptions, finalMemberCount = 1 ) { - mocks.select - .mockReturnValueOnce( - selectChain([ - { - ownerId: 'owner-1', - ownerName: 'Owner', - ownerEmail: 'owner@example.com', - ownerStripeCustomerId: 'cus_1', - organizationName: 'Acme', - ownerRole: 'owner', - }, - ]) - ) - .mockReturnValueOnce(selectChain([{ value: 1 }])) - .mockReturnValueOnce(selectChain(localSubscriptions)) - .mockReturnValueOnce(selectChain(finalLocalSubscriptions)) - .mockReturnValueOnce(selectChain([{ value: finalMemberCount }])) + queueTableRows(schemaMock.user, [ + { + ownerId: 'owner-1', + ownerName: 'Owner', + ownerEmail: 'owner@example.com', + ownerStripeCustomerId: 'cus_1', + organizationName: 'Acme', + ownerRole: 'owner', + }, + ]) + queueTableRows(schemaMock.member, [{ value: 1 }]) + queueTableRows(schemaMock.subscription, localSubscriptions) + queueTableRows(schemaMock.subscription, finalLocalSubscriptions) + queueTableRows(schemaMock.member, [{ value: finalMemberCount }]) } describe('Enterprise issuance outbox handler', () => { beforeEach(() => { vi.clearAllMocks() - mocks.select.mockReset() + resetDbChainMock() mocks.subscriptionsList.mockResolvedValue({ data: [], has_more: false }) mocks.customersList.mockResolvedValue({ data: [], has_more: false }) mocks.pricesList.mockResolvedValue({ data: [], has_more: false }) @@ -468,7 +413,7 @@ describe('Enterprise issuance outbox handler', () => { context() ) - expect(mocks.select).not.toHaveBeenCalled() + expect(dbChainMockFns.select).not.toHaveBeenCalled() expect(mocks.subscriptionsCreate).not.toHaveBeenCalled() }) @@ -482,7 +427,7 @@ describe('Enterprise issuance outbox handler', () => { describe('Enterprise metadata outbox handler', () => { beforeEach(() => { vi.clearAllMocks() - mocks.select.mockReset() + resetDbChainMock() }) it('pushes only the latest full desired metadata under an operation-stable key', async () => { @@ -498,13 +443,12 @@ describe('Enterprise metadata outbox handler', () => { concurrencyLimit: 1250, }, } - mocks.select - .mockReturnValueOnce( - selectChain([{ stripeSubscriptionId: 'sub_1', referenceId: 'org-1', metadata: {} }]) - ) - .mockReturnValueOnce(selectChain([{ metadata: {} }])) - .mockReturnValueOnce(selectChain([{ id: 'metadata-event-1', payload }])) - .mockReturnValueOnce(selectChain([{ value: 10 }])) + queueTableRows(schemaMock.subscription, [ + { stripeSubscriptionId: 'sub_1', referenceId: 'org-1', metadata: {} }, + ]) + queueTableRows(schemaMock.subscription, [{ metadata: {} }]) + queueTableRows(schemaMock.outboxEvent, [{ id: 'metadata-event-1', payload }]) + queueTableRows(schemaMock.member, [{ value: 10 }]) mocks.subscriptionsUpdate.mockResolvedValue({ id: 'sub_1' }) await expect( @@ -546,13 +490,12 @@ describe('Enterprise metadata outbox handler', () => { concurrencyLimit: null, }, } - mocks.select - .mockReturnValueOnce( - selectChain([{ stripeSubscriptionId: 'sub_1', referenceId: 'org-1', metadata: {} }]) - ) - .mockReturnValueOnce(selectChain([{ metadata: {} }])) - .mockReturnValueOnce(selectChain([{ id: 'metadata-event-2', payload }])) - .mockReturnValueOnce(selectChain([{ value: 10 }])) + queueTableRows(schemaMock.subscription, [ + { stripeSubscriptionId: 'sub_1', referenceId: 'org-1', metadata: {} }, + ]) + queueTableRows(schemaMock.subscription, [{ metadata: {} }]) + queueTableRows(schemaMock.outboxEvent, [{ id: 'metadata-event-2', payload }]) + queueTableRows(schemaMock.member, [{ value: 10 }]) mocks.subscriptionsUpdate.mockResolvedValue({ id: 'sub_1' }) await expect( @@ -583,24 +526,21 @@ describe('Enterprise metadata outbox handler', () => { deliveryRevision: 0, metadata: { seats: 12 }, } - mocks.select - .mockReturnValueOnce( - selectChain([{ stripeSubscriptionId: 'sub_1', referenceId: 'org-1', metadata: {} }]) - ) - .mockReturnValueOnce(selectChain([{ metadata: {} }])) - .mockReturnValueOnce( - selectChain([ - { - id: 'newer-event', - payload: { - subscriptionId: 'local-sub-1', - revision: 4, - deliveryRevision: 0, - metadata: { seats: 15 }, - }, - }, - ]) - ) + queueTableRows(schemaMock.subscription, [ + { stripeSubscriptionId: 'sub_1', referenceId: 'org-1', metadata: {} }, + ]) + queueTableRows(schemaMock.subscription, [{ metadata: {} }]) + queueTableRows(schemaMock.outboxEvent, [ + { + id: 'newer-event', + payload: { + subscriptionId: 'local-sub-1', + revision: 4, + deliveryRevision: 0, + metadata: { seats: 15 }, + }, + }, + ]) await syncEnterpriseMetadataInStripe(payload, { eventId: 'older-event', @@ -619,15 +559,13 @@ describe('Enterprise metadata outbox handler', () => { deliveryRevision: 0, metadata: { seats: 15 }, } - mocks.select.mockReturnValueOnce( - selectChain([ - { - stripeSubscriptionId: 'sub_1', - referenceId: 'org-1', - metadata: { simConfigOperationId: 'metadata-event-1' }, - }, - ]) - ) + queueTableRows(schemaMock.subscription, [ + { + stripeSubscriptionId: 'sub_1', + referenceId: 'org-1', + metadata: { simConfigOperationId: 'metadata-event-1' }, + }, + ]) await syncEnterpriseMetadataInStripe(payload, { eventId: 'metadata-event-1', diff --git a/apps/sim/lib/billing/organizations/member-limits.test.ts b/apps/sim/lib/billing/organizations/member-limits.test.ts index 51fb47aa021..3558a65337f 100644 --- a/apps/sim/lib/billing/organizations/member-limits.test.ts +++ b/apps/sim/lib/billing/organizations/member-limits.test.ts @@ -1,82 +1,56 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { dbChainMock, dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { - mockDbState, - mockInsert, - mockInsertValues, - mockOnConflictDoUpdate, - mockDelete, - mockDeleteWhere, + schemaTables, mockAnd, mockEq, mockGte, mockIsNull, - mockLeftJoin, mockLt, mockOr, mockGetOrganizationSubscription, } = vi.hoisted(() => ({ - mockDbState: { selectResults: [] as unknown[] }, - mockInsert: vi.fn(), - mockInsertValues: vi.fn(), - mockOnConflictDoUpdate: vi.fn(), - mockDelete: vi.fn(), - mockDeleteWhere: vi.fn(), + schemaTables: { + organizationMemberUsageLimit: { + id: 'oml.id', + organizationId: 'oml.organizationId', + userId: 'oml.userId', + usageLimit: 'oml.usageLimit', + setBy: 'oml.setBy', + createdAt: 'oml.createdAt', + updatedAt: 'oml.updatedAt', + }, + usageLog: { + billingEntityType: 'usageLog.billingEntityType', + billingEntityId: 'usageLog.billingEntityId', + billingPeriodStart: 'usageLog.billingPeriodStart', + billingPeriodEnd: 'usageLog.billingPeriodEnd', + createdAt: 'usageLog.createdAt', + cost: 'usageLog.cost', + userId: 'usageLog.userId', + }, + workspace: { + id: 'workspace.id', + organizationAssignedAt: 'workspace.organizationAssignedAt', + organizationId: 'workspace.organizationId', + }, + }, mockAnd: vi.fn((...conditions: unknown[]) => ({ operator: 'and', conditions })), mockEq: vi.fn((field: unknown, value: unknown) => ({ field, value })), mockGte: vi.fn((field: unknown, value: unknown) => ({ operator: 'gte', field, value })), mockIsNull: vi.fn((field: unknown) => ({ operator: 'isNull', field })), - mockLeftJoin: vi.fn(), mockLt: vi.fn((field: unknown, value: unknown) => ({ operator: 'lt', field, value })), mockOr: vi.fn((...conditions: unknown[]) => ({ operator: 'or', conditions })), mockGetOrganizationSubscription: vi.fn(), })) -vi.mock('@sim/db', () => ({ - db: { - select: vi.fn(() => { - const chain: Record = {} - chain.from = vi.fn(() => chain) - chain.leftJoin = mockLeftJoin.mockImplementation(() => 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', () => dbChainMock) -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', - }, - usageLog: { - billingEntityType: 'usageLog.billingEntityType', - billingEntityId: 'usageLog.billingEntityId', - billingPeriodStart: 'usageLog.billingPeriodStart', - billingPeriodEnd: 'usageLog.billingPeriodEnd', - createdAt: 'usageLog.createdAt', - cost: 'usageLog.cost', - userId: 'usageLog.userId', - }, - workspace: { - id: 'workspace.id', - organizationAssignedAt: 'workspace.organizationAssignedAt', - organizationId: 'workspace.organizationId', - }, -})) +vi.mock('@sim/db/schema', () => schemaTables) vi.mock('drizzle-orm', () => ({ and: mockAnd, @@ -102,22 +76,21 @@ import { beforeEach(() => { vi.clearAllMocks() - mockDbState.selectResults = [] - mockInsert.mockReturnValue({ values: mockInsertValues }) - mockInsertValues.mockReturnValue({ onConflictDoUpdate: mockOnConflictDoUpdate }) - mockOnConflictDoUpdate.mockResolvedValue(undefined) - mockDelete.mockReturnValue({ where: mockDeleteWhere }) - mockDeleteWhere.mockResolvedValue(undefined) + resetDbChainMock() +}) + +afterAll(() => { + resetDbChainMock() }) describe('getOrgMemberUsageLimit', () => { it('returns null when no row exists', async () => { - mockDbState.selectResults = [[]] + queueTableRows(schemaTables.organizationMemberUsageLimit, []) await expect(getOrgMemberUsageLimit('org-1', 'user-2')).resolves.toBeNull() }) it('returns the stored dollar limit as a number', async () => { - mockDbState.selectResults = [[{ usageLimit: '2' }]] + queueTableRows(schemaTables.organizationMemberUsageLimit, [{ usageLimit: '2' }]) await expect(getOrgMemberUsageLimit('org-1', 'user-2')).resolves.toBe(2) }) }) @@ -128,7 +101,7 @@ describe('getOrgMemberUsageForBillingPeriod', () => { start: new Date('2026-06-01T00:00:00.000Z'), end: new Date('2026-07-01T00:00:00.000Z'), } - mockDbState.selectResults = [[{ cost: '4.5' }]] + queueTableRows(schemaTables.usageLog, [{ cost: '4.5' }]) mockGetOrganizationSubscription.mockResolvedValue({ periodStart: new Date('2026-07-01T00:00:00.000Z'), periodEnd: new Date('2026-08-01T00:00:00.000Z'), @@ -150,7 +123,7 @@ describe('getOrgMemberUsageForBillingPeriod', () => { expect(mockGte).toHaveBeenCalledWith('usageLog.createdAt', 'workspace.organizationAssignedAt') expect(mockGte).toHaveBeenCalledWith('usageLog.createdAt', billingPeriod.start) expect(mockLt).toHaveBeenCalledWith('usageLog.createdAt', billingPeriod.end) - expect(mockLeftJoin).toHaveBeenCalledWith( + expect(dbChainMockFns.leftJoin).toHaveBeenCalledWith( expect.objectContaining({ id: 'workspace.id' }), expect.anything() ) @@ -174,23 +147,23 @@ describe('getOrgMemberUsageForBillingPeriod', () => { 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(dbChainMockFns.insert).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.delete).not.toHaveBeenCalled() + const values = dbChainMockFns.values.mock.calls[0][0] expect(values).toMatchObject({ organizationId: 'org-1', userId: 'user-2', usageLimit: '2', setBy: 'admin-1', }) - expect(mockOnConflictDoUpdate).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.onConflictDoUpdate).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() + expect(dbChainMockFns.delete).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.where).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.insert).not.toHaveBeenCalled() }) }) @@ -199,7 +172,7 @@ describe('getOrgMemberUsageForCurrentPeriod', () => { const periodStart = new Date('2026-06-01T00:00:00.000Z') const periodEnd = new Date('2026-07-01T00:00:00.000Z') mockGetOrganizationSubscription.mockResolvedValue({ periodStart, periodEnd }) - mockDbState.selectResults = [[{ cost: '5' }]] + queueTableRows(schemaTables.usageLog, [{ cost: '5' }]) const result = await getOrgMemberUsageForCurrentPeriod('org-1', 'user-2') @@ -213,7 +186,7 @@ describe('getOrgMemberUsageForCurrentPeriod', () => { it('uses a prefetched subscription without a second lookup', async () => { const periodStart = new Date('2026-06-01T00:00:00.000Z') const periodEnd = new Date('2026-07-01T00:00:00.000Z') - mockDbState.selectResults = [[{ cost: '5' }]] + queueTableRows(schemaTables.usageLog, [{ cost: '5' }]) const result = await getOrgMemberUsageForCurrentPeriod('org-1', 'user-2', { periodStart, @@ -226,7 +199,7 @@ describe('getOrgMemberUsageForCurrentPeriod', () => { }) it('falls back to the all-time window when the org has no subscription period', async () => { - mockDbState.selectResults = [[{ cost: '7' }]] + queueTableRows(schemaTables.usageLog, [{ cost: '7' }]) const result = await getOrgMemberUsageForCurrentPeriod('org-1', 'user-2', null) diff --git a/apps/sim/lib/billing/organizations/provision-seat.test.ts b/apps/sim/lib/billing/organizations/provision-seat.test.ts index 126f1c645c6..638d84897d4 100644 --- a/apps/sim/lib/billing/organizations/provision-seat.test.ts +++ b/apps/sim/lib/billing/organizations/provision-seat.test.ts @@ -1,7 +1,8 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { mockGetOrganizationSubscription, @@ -12,7 +13,6 @@ const { mockGetPlanByName, enqueueMock, updateCalls, - globalTransactionMock, } = vi.hoisted(() => ({ mockGetOrganizationSubscription: vi.fn(), mockGetHighestPriorityPersonalSubscription: vi.fn(), @@ -22,26 +22,9 @@ const { mockGetPlanByName: vi.fn(), enqueueMock: vi.fn(), updateCalls: { value: [] as Array> }, - globalTransactionMock: vi.fn(), })) -vi.mock('@sim/db', () => { - const update = () => ({ - set: (values: Record) => { - updateCalls.value.push(values) - return { where: () => Promise.resolve([]) } - }, - }) - const txMock = { update } - globalTransactionMock.mockImplementation(async (cb: (tx: typeof txMock) => Promise) => - cb(txMock) - ) - const dbMock = { - update, - transaction: globalTransactionMock, - } - return { db: dbMock } -}) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/billing/core/billing', () => ({ getOrganizationSubscription: mockGetOrganizationSubscription, @@ -95,6 +78,7 @@ function testExecutor(onUpdate: () => void = () => {}) { describe('ensureTeamOrganizationForAcceptance', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() updateCalls.value = [] mockGetPlanByName.mockReturnValue({ priceId: 'price_team_month', @@ -103,6 +87,10 @@ describe('ensureTeamOrganizationForAcceptance', () => { mockAssertNoUnresolvedEnterpriseIssuance.mockResolvedValue(undefined) }) + afterAll(() => { + resetDbChainMock() + }) + it('is a no-op for enterprise organizations (fixed seats)', async () => { mockGetOrganizationSubscription.mockResolvedValue({ id: 'sub-ent', @@ -187,7 +175,7 @@ describe('ensureTeamOrganizationForAcceptance', () => { 'org-1', expect.objectContaining({ executor }) ) - expect(globalTransactionMock).not.toHaveBeenCalled() + expect(dbChainMockFns.transaction).not.toHaveBeenCalled() }) it('blocks an org-scoped Pro conversion while Enterprise issuance is unresolved', async () => { @@ -273,7 +261,7 @@ describe('ensureTeamOrganizationForAcceptance', () => { 'stripe.sync-subscription-seats', expect.objectContaining({ subscriptionId: 'sub-pro' }) ) - expect(globalTransactionMock).not.toHaveBeenCalled() + expect(dbChainMockFns.transaction).not.toHaveBeenCalled() expect(lockOrder).toEqual(['organization', 'subscription']) // ...but with no scheduled cancellation there is no cancel-sync event. expect(enqueueMock).not.toHaveBeenCalledWith( diff --git a/apps/sim/lib/billing/organizations/seat-drift.test.ts b/apps/sim/lib/billing/organizations/seat-drift.test.ts index 43338b0c615..a95ffa0d2f9 100644 --- a/apps/sim/lib/billing/organizations/seat-drift.test.ts +++ b/apps/sim/lib/billing/organizations/seat-drift.test.ts @@ -1,27 +1,15 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { dbChainMock, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' -const { mockReconcileOrganizationSeats, selectRows, mockFeatureFlags } = vi.hoisted(() => ({ +const { mockReconcileOrganizationSeats, mockFeatureFlags } = vi.hoisted(() => ({ mockReconcileOrganizationSeats: vi.fn(), - selectRows: { value: [] as unknown[] }, mockFeatureFlags: { isBillingEnabled: true }, })) -vi.mock('@sim/db', () => { - const makeChain = () => { - const chain: Record = {} - chain.from = () => chain - chain.innerJoin = () => chain - chain.where = () => chain - chain.groupBy = () => chain - chain.having = () => chain - chain.orderBy = () => Promise.resolve(selectRows.value) - return chain - } - return { db: { select: () => makeChain() } } -}) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/billing/organizations/seats', () => ({ reconcileOrganizationSeats: mockReconcileOrganizationSeats, @@ -38,15 +26,22 @@ import { reconcileTeamSeatDrift } from '@/lib/billing/organizations/seat-drift' describe('reconcileTeamSeatDrift', () => { beforeEach(() => { vi.clearAllMocks() - selectRows.value = [] + resetDbChainMock() mockFeatureFlags.isBillingEnabled = true mockReconcileOrganizationSeats.mockResolvedValue({ changed: true, previousSeats: 1, seats: 2 }) }) + afterAll(() => { + resetDbChainMock() + }) + it('reconciles each drifted Team org returned by the query', async () => { // The SQL WHERE (Team-only) + HAVING (seats != member count) already // restrict the result to drifted Team orgs; the function reconciles each. - selectRows.value = [{ organizationId: 'org-1' }, { organizationId: 'org-2' }] + queueTableRows(schemaMock.subscription, [ + { organizationId: 'org-1' }, + { organizationId: 'org-2' }, + ]) const result = await reconcileTeamSeatDrift() @@ -63,7 +58,7 @@ describe('reconcileTeamSeatDrift', () => { }) it('reconciles a past-due Team candidate returned by the entitlement query', async () => { - selectRows.value = [{ organizationId: 'org-past-due' }] + queueTableRows(schemaMock.subscription, [{ organizationId: 'org-past-due' }]) const result = await reconcileTeamSeatDrift() @@ -75,7 +70,10 @@ describe('reconcileTeamSeatDrift', () => { }) it('counts only reconciles that changed the seat count', async () => { - selectRows.value = [{ organizationId: 'org-a' }, { organizationId: 'org-b' }] + queueTableRows(schemaMock.subscription, [ + { organizationId: 'org-a' }, + { organizationId: 'org-b' }, + ]) mockReconcileOrganizationSeats .mockResolvedValueOnce({ changed: true, seats: 2 }) .mockResolvedValueOnce({ changed: false }) @@ -86,7 +84,10 @@ describe('reconcileTeamSeatDrift', () => { }) it('continues past a reconcile failure', async () => { - selectRows.value = [{ organizationId: 'org-a' }, { organizationId: 'org-b' }] + queueTableRows(schemaMock.subscription, [ + { organizationId: 'org-a' }, + { organizationId: 'org-b' }, + ]) mockReconcileOrganizationSeats .mockRejectedValueOnce(new Error('db error')) .mockResolvedValueOnce({ changed: true, seats: 3 }) @@ -107,9 +108,10 @@ describe('reconcileTeamSeatDrift', () => { }) it('caps reconciles per run while still reporting the full drift count', async () => { - selectRows.value = Array.from({ length: 150 }, (_, i) => ({ - organizationId: `org-${i}`, - })) + queueTableRows( + schemaMock.subscription, + Array.from({ length: 150 }, (_, i) => ({ organizationId: `org-${i}` })) + ) const result = await reconcileTeamSeatDrift() diff --git a/apps/sim/lib/billing/organizations/seats.test.ts b/apps/sim/lib/billing/organizations/seats.test.ts index 1a98033586a..26154aa1987 100644 --- a/apps/sim/lib/billing/organizations/seats.test.ts +++ b/apps/sim/lib/billing/organizations/seats.test.ts @@ -1,43 +1,23 @@ /** * @vitest-environment node */ -import { auditMock } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const { mockSyncSubscriptionUsageLimits, enqueueMock, setMock, queryQueue, mockFeatureFlags } = - vi.hoisted(() => ({ - mockSyncSubscriptionUsageLimits: vi.fn(), - enqueueMock: vi.fn(), - setMock: vi.fn(), - queryQueue: { value: [] as unknown[][] }, - mockFeatureFlags: { isBillingEnabled: true }, - })) - -vi.mock('@sim/db', () => { - const makeSelectChain = () => { - const chain: Record = {} - chain.from = () => chain - chain.where = () => chain - chain.for = () => chain - chain.limit = () => Promise.resolve(queryQueue.value.shift() ?? []) - chain.then = (resolve: (rows: unknown[]) => unknown, reject?: (e: unknown) => unknown) => - Promise.resolve(queryQueue.value.shift() ?? []).then(resolve, reject) - return chain - } - const update = () => ({ - set: (values: Record) => { - setMock(values) - return { where: () => Promise.resolve([]) } - }, - }) - const txMock = { select: () => makeSelectChain(), update } - const dbMock = { - select: () => makeSelectChain(), - update, - transaction: async (cb: (tx: typeof txMock) => Promise) => cb(txMock), - } - return { db: dbMock } -}) +import { + auditMock, + dbChainMock, + dbChainMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, +} from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockSyncSubscriptionUsageLimits, enqueueMock, mockFeatureFlags } = vi.hoisted(() => ({ + mockSyncSubscriptionUsageLimits: vi.fn(), + enqueueMock: vi.fn(), + mockFeatureFlags: { isBillingEnabled: true }, +})) + +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/billing/organization', () => ({ syncSubscriptionUsageLimits: mockSyncSubscriptionUsageLimits, @@ -71,16 +51,26 @@ const teamSub = { stripeSubscriptionId: 'stripe_sub', } +/** Queues the two in-transaction reads: locked subscription, then member count. */ +function queueReconcileReads(subscriptionRows: unknown[], memberCountRows: unknown[] = []) { + queueTableRows(schemaMock.subscription, subscriptionRows) + queueTableRows(schemaMock.member, memberCountRows) +} + describe('reconcileOrganizationSeats', () => { beforeEach(() => { vi.clearAllMocks() - queryQueue.value = [] + resetDbChainMock() enqueueMock.mockResolvedValue('evt-1') mockFeatureFlags.isBillingEnabled = true }) + afterAll(() => { + resetDbChainMock() + }) + it('grows seats to the member count and enqueues a Stripe sync', async () => { - queryQueue.value = [[teamSub], [{ value: 2 }]] + queueReconcileReads([teamSub], [{ value: 2 }]) const result = await reconcileOrganizationSeats({ organizationId: 'org-1', @@ -94,7 +84,7 @@ describe('reconcileOrganizationSeats', () => { reason: undefined, outboxEventId: 'evt-1', }) - expect(setMock).toHaveBeenCalledWith({ seats: 2 }) + expect(dbChainMockFns.set).toHaveBeenCalledWith({ seats: 2 }) expect(enqueueMock).toHaveBeenCalledWith(expect.anything(), 'stripe.sync-subscription-seats', { subscriptionId: 'sub-1', reason: 'member-accepted-invite', @@ -105,7 +95,7 @@ describe('reconcileOrganizationSeats', () => { }) it('reconciles a past-due Team subscription because it remains entitled', async () => { - queryQueue.value = [[{ ...teamSub, status: 'past_due' }], [{ value: 2 }]] + queueReconcileReads([{ ...teamSub, status: 'past_due' }], [{ value: 2 }]) const result = await reconcileOrganizationSeats({ organizationId: 'org-1', @@ -113,12 +103,12 @@ describe('reconcileOrganizationSeats', () => { }) expect(result.changed).toBe(true) - expect(setMock).toHaveBeenCalledWith({ seats: 2 }) + expect(dbChainMockFns.set).toHaveBeenCalledWith({ seats: 2 }) expect(enqueueMock).toHaveBeenCalledOnce() }) it('still records the seat audit when the post-commit usage-limit sync fails', async () => { - queryQueue.value = [[teamSub], [{ value: 2 }]] + queueReconcileReads([teamSub], [{ value: 2 }]) mockSyncSubscriptionUsageLimits.mockRejectedValueOnce(new Error('sync unavailable')) const result = await reconcileOrganizationSeats({ @@ -128,7 +118,7 @@ describe('reconcileOrganizationSeats', () => { }) expect(result.changed).toBe(true) - expect(setMock).toHaveBeenCalledWith({ seats: 2 }) + expect(dbChainMockFns.set).toHaveBeenCalledWith({ seats: 2 }) expect(auditMock.recordAudit).toHaveBeenCalledWith( expect.objectContaining({ actorId: 'user-1', @@ -139,7 +129,7 @@ describe('reconcileOrganizationSeats', () => { }) it('reduces seats to the member count on removal', async () => { - queryQueue.value = [[{ ...teamSub, seats: 3 }], [{ value: 2 }]] + queueReconcileReads([{ ...teamSub, seats: 3 }], [{ value: 2 }]) const result = await reconcileOrganizationSeats({ organizationId: 'org-1', @@ -148,12 +138,12 @@ describe('reconcileOrganizationSeats', () => { expect(result.changed).toBe(true) expect(result.seats).toBe(2) - expect(setMock).toHaveBeenCalledWith({ seats: 2 }) + expect(dbChainMockFns.set).toHaveBeenCalledWith({ seats: 2 }) expect(enqueueMock).toHaveBeenCalled() }) it('is a no-op when seats already match the member count', async () => { - queryQueue.value = [[{ ...teamSub, seats: 2 }], [{ value: 2 }]] + queueReconcileReads([{ ...teamSub, seats: 2 }], [{ value: 2 }]) const result = await reconcileOrganizationSeats({ organizationId: 'org-1', @@ -167,13 +157,13 @@ describe('reconcileOrganizationSeats', () => { reason: undefined, outboxEventId: undefined, }) - expect(setMock).not.toHaveBeenCalled() + expect(dbChainMockFns.set).not.toHaveBeenCalled() expect(enqueueMock).not.toHaveBeenCalled() expect(mockSyncSubscriptionUsageLimits).not.toHaveBeenCalled() }) it('never drops below one seat', async () => { - queryQueue.value = [[{ ...teamSub, seats: 3 }], [{ value: 0 }]] + queueReconcileReads([{ ...teamSub, seats: 3 }], [{ value: 0 }]) const result = await reconcileOrganizationSeats({ organizationId: 'org-1', @@ -181,11 +171,11 @@ describe('reconcileOrganizationSeats', () => { }) expect(result.seats).toBe(1) - expect(setMock).toHaveBeenCalledWith({ seats: 1 }) + expect(dbChainMockFns.set).toHaveBeenCalledWith({ seats: 1 }) }) it('skips non-Team subscriptions', async () => { - queryQueue.value = [[{ ...teamSub, plan: 'pro_6000' }]] + queueReconcileReads([{ ...teamSub, plan: 'pro_6000' }]) const result = await reconcileOrganizationSeats({ organizationId: 'org-1', @@ -194,12 +184,12 @@ describe('reconcileOrganizationSeats', () => { expect(result.changed).toBe(false) expect(result.reason).toMatch(/Team/) - expect(setMock).not.toHaveBeenCalled() + expect(dbChainMockFns.set).not.toHaveBeenCalled() expect(enqueueMock).not.toHaveBeenCalled() }) it('skips when the organization has no usable subscription', async () => { - queryQueue.value = [[]] + queueReconcileReads([]) const result = await reconcileOrganizationSeats({ organizationId: 'org-1', diff --git a/apps/sim/lib/billing/storage/limits.test.ts b/apps/sim/lib/billing/storage/limits.test.ts index c0850a29e97..c690221cdfb 100644 --- a/apps/sim/lib/billing/storage/limits.test.ts +++ b/apps/sim/lib/billing/storage/limits.test.ts @@ -1,31 +1,16 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const { - mockEq, - mockFlags, - mockFrom, - mockGetHighestPrioritySubscription, - mockLimit, - mockSelect, - mockWhere, -} = vi.hoisted(() => ({ +import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockEq, mockFlags, mockGetHighestPrioritySubscription } = vi.hoisted(() => ({ mockEq: vi.fn((field: unknown, value: unknown) => ({ field, value })), mockFlags: { isBillingEnabled: true }, - mockFrom: vi.fn(), mockGetHighestPrioritySubscription: vi.fn(), - mockLimit: vi.fn(), - mockSelect: vi.fn(), - mockWhere: vi.fn(), })) -vi.mock('@sim/db', () => ({ - db: { - select: mockSelect, - }, -})) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@sim/db/schema', () => ({ organization: { @@ -91,17 +76,19 @@ const GIB = 1024 ** 3 describe('storage limits and quota', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockFlags.isBillingEnabled = true mockGetEnv.mockReturnValue(undefined) - mockSelect.mockReturnValue({ from: mockFrom }) - mockFrom.mockReturnValue({ where: mockWhere }) - mockWhere.mockReturnValue({ limit: mockLimit }) - mockLimit.mockResolvedValue([{ storageUsedBytes: 1024 }]) + dbChainMockFns.limit.mockResolvedValue([{ storageUsedBytes: 1024 }]) mockGetHighestPrioritySubscription.mockResolvedValue(null) }) + afterAll(() => { + resetDbChainMock() + }) + it('reads user and organization counters through the same entity-aware path', async () => { - mockLimit + dbChainMockFns.limit .mockResolvedValueOnce([{ storageUsedBytes: 11 }]) .mockResolvedValueOnce([{ storageUsedBytes: 22 }]) .mockResolvedValueOnce([{ storageUsedBytes: 33 }]) @@ -138,7 +125,7 @@ describe('storage limits and quota', () => { }) it('returns the exact same quota result for legacy and workspace organization payers', async () => { - mockLimit.mockResolvedValue([{ storageUsedBytes: GIB }]) + dbChainMockFns.limit.mockResolvedValue([{ storageUsedBytes: GIB }]) mockGetHighestPrioritySubscription.mockResolvedValue({ metadata: { customStorageLimitGB: 1 }, plan: 'team_25000', @@ -170,7 +157,7 @@ describe('storage limits and quota', () => { await expect(checkStorageQuota('workspace-owner', GIB)).resolves.toEqual(expected) await expect(checkStorageQuotaForBillingContext(ORG_CONTEXT, GIB)).resolves.toEqual(expected) expect(mockGetHighestPrioritySubscription).not.toHaveBeenCalled() - expect(mockSelect).not.toHaveBeenCalled() + expect(dbChainMockFns.select).not.toHaveBeenCalled() }) it('opts into free-tier enforcement when FREE_STORAGE_LIMIT_GB is explicitly set', async () => { @@ -178,7 +165,7 @@ describe('storage limits and quota', () => { mockGetEnv.mockImplementation((variable: string) => variable === 'FREE_STORAGE_LIMIT_GB' ? '1' : undefined ) - mockLimit.mockResolvedValue([{ storageUsedBytes: GIB }]) + dbChainMockFns.limit.mockResolvedValue([{ storageUsedBytes: GIB }]) await expect(checkStorageQuota('workspace-owner', GIB / 2)).resolves.toEqual({ allowed: false, @@ -199,12 +186,12 @@ describe('storage limits and quota', () => { mockGetHighestPrioritySubscription.mockRejectedValueOnce(new Error('subscription unavailable')) await expect(checkStorageQuota('workspace-owner', GIB)).resolves.toEqual(expected) - mockLimit.mockRejectedValueOnce(new Error('counter unavailable')) + dbChainMockFns.limit.mockRejectedValueOnce(new Error('counter unavailable')) await expect(checkStorageQuotaForBillingContext(ORG_CONTEXT, GIB)).resolves.toEqual(expected) }) it('retains zero fallback for direct usage readers', async () => { - mockLimit.mockRejectedValueOnce(new Error('counter unavailable')) + dbChainMockFns.limit.mockRejectedValueOnce(new Error('counter unavailable')) await expect(getStorageUsageForBillingContext(ORG_CONTEXT)).resolves.toBe(0) }) diff --git a/apps/sim/lib/billing/threshold-billing.test.ts b/apps/sim/lib/billing/threshold-billing.test.ts index 09f88c9ae09..63838978487 100644 --- a/apps/sim/lib/billing/threshold-billing.test.ts +++ b/apps/sim/lib/billing/threshold-billing.test.ts @@ -1,13 +1,18 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + dbChainMock, + dbChainMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, +} from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { mockCalculateSubscriptionOverage, mockComputeOrgOverageAmount, - mockDbSelect, - mockDbTransaction, mockEnqueueOutboxEvent, mockGetEffectiveBillingStatus, mockGetHighestPrioritySubscription, @@ -20,15 +25,9 @@ const { mockIsOrganizationBillingBlocked, mockRecordAudit, mockCaptureServerEvent, - mockTxExecute, - mockTxSelect, - mockTxStatsLimit, - mockTxUpdate, } = vi.hoisted(() => ({ mockCalculateSubscriptionOverage: vi.fn(), mockComputeOrgOverageAmount: vi.fn(), - mockDbSelect: vi.fn(), - mockDbTransaction: vi.fn(), mockEnqueueOutboxEvent: vi.fn(), mockGetEffectiveBillingStatus: vi.fn(), mockGetHighestPrioritySubscription: vi.fn(), @@ -41,10 +40,6 @@ const { mockIsOrganizationBillingBlocked: vi.fn(), mockRecordAudit: vi.fn(), mockCaptureServerEvent: vi.fn(), - mockTxExecute: vi.fn(), - mockTxSelect: vi.fn(), - mockTxStatsLimit: vi.fn(), - mockTxUpdate: vi.fn(), })) vi.mock('@sim/audit', () => ({ @@ -53,38 +48,7 @@ vi.mock('@sim/audit', () => ({ recordAudit: mockRecordAudit, })) -vi.mock('@sim/db', () => ({ - db: { - select: mockDbSelect, - transaction: mockDbTransaction, - }, -})) - -vi.mock('@sim/db/schema', () => ({ - member: { - organizationId: 'member.organizationId', - role: 'member.role', - userId: 'member.userId', - }, - organization: { - creditBalance: 'organization.creditBalance', - departedMemberUsage: 'organization.departedMemberUsage', - id: 'organization.id', - }, - subscription: { - id: 'subscription.id', - stripeCustomerId: 'subscription.stripeCustomerId', - }, - userStats: { - billedOverageThisPeriod: 'userStats.billedOverageThisPeriod', - creditBalance: 'userStats.creditBalance', - currentPeriodCost: 'userStats.currentPeriodCost', - lastPeriodCost: 'userStats.lastPeriodCost', - proPeriodCostSnapshot: 'userStats.proPeriodCostSnapshot', - proPeriodCostSnapshotAt: 'userStats.proPeriodCostSnapshotAt', - userId: 'userStats.userId', - }, -})) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/billing/core/access', () => ({ getEffectiveBillingStatus: mockGetEffectiveBillingStatus, @@ -140,12 +104,6 @@ import { ThresholdSettlementError, } from '@/lib/billing/threshold-billing' -interface MockTx { - execute: typeof mockTxExecute - select: typeof mockTxSelect - update: typeof mockTxUpdate -} - const userSubscription = { id: 'sub-db-1', plan: 'pro', @@ -162,88 +120,85 @@ const expectedBillingPeriod = { end: new Date('2026-06-01T00:00:00.000Z'), } -function buildSelectChain(rows: T[]) { - const chain = { - from: vi.fn(() => chain), - leftJoin: vi.fn(() => chain), - innerJoin: vi.fn(() => chain), - where: vi.fn(() => result), - } - const result = { - limit: vi.fn(async () => rows), - then: (resolve: (value: T[]) => unknown, reject?: (reason: unknown) => unknown) => - Promise.resolve(rows).then(resolve, reject), - } +const defaultUsageSnapshotRow = { + currentPeriodCost: '0', + proPeriodCostSnapshot: '0', + proPeriodCostSnapshotAt: null as Date | null, + lastPeriodCost: '0', +} + +/** + * Queues the two pre-transaction personal reads: the user_stats usage snapshot + * and the subscription's Stripe customer row. + */ +function queuePersonalReads( + snapshot: Record = defaultUsageSnapshotRow, + customerId = 'cus_1' +) { + queueTableRows(schemaMock.userStats, [snapshot]) + queueTableRows(schemaMock.subscription, [{ stripeCustomerId: customerId }]) +} +/** Builds the locked in-transaction user_stats row. */ +function lockedStatsRow(overrides: Record = {}) { return { - from: chain.from, + ...defaultUsageSnapshotRow, + billedOverageThisPeriod: '0', + creditBalance: '0', + ...overrides, } } -function buildPersonalSelectChain(customerId = 'cus_1') { - return buildSelectChain([ - { - currentPeriodCost: '0', - proPeriodCostSnapshot: '0', - proPeriodCostSnapshotAt: null, - lastPeriodCost: '0', - stripeCustomerId: customerId, - }, - ]) +/** Queues the locked user_stats read taken inside the settlement transaction. */ +function queueLockedStats(row: Record) { + queueTableRows(schemaMock.userStats, [row]) } -function buildPersonalSnapshotSelectChain({ - currentPeriodCost = '0', - proPeriodCostSnapshot = '0', - proPeriodCostSnapshotAt = null, - lastPeriodCost = '0', -}: { - currentPeriodCost?: string - proPeriodCostSnapshot?: string - proPeriodCostSnapshotAt?: Date | null - lastPeriodCost?: string -}) { - return buildSelectChain([ - { - currentPeriodCost, - proPeriodCostSnapshot, - proPeriodCostSnapshotAt, - lastPeriodCost, - }, - ]) +const orgMemberUsageRow = { + userId: 'owner-1', + role: 'owner', + currentPeriodCost: '350', + departedMemberUsage: '25', } -function buildStatsSelectChain() { - const result = { - for: vi.fn(() => result), - limit: mockTxStatsLimit, - then: (resolve: (value: unknown[]) => unknown, reject?: (reason: unknown) => unknown) => - Promise.resolve(mockTxStatsLimit()).then(resolve, reject), - } - - return { - from: vi.fn(() => ({ - leftJoin: vi.fn(() => ({ - innerJoin: vi.fn(() => ({ - where: vi.fn(() => result), - })), - })), - where: vi.fn(() => result), - })), - } +/** + * Queues the organization settlement reads in table order: the pre-transaction + * member usage join, then the locked owner row, owner stats, organization row, + * and locked member usage join inside the transaction. + */ +function queueOrgReads({ + memberUsageRows = [orgMemberUsageRow], + lockedOwnerRows = [{ userId: 'owner-1' }], + ownerStatsRows = [{ billedOverageThisPeriod: '0' }], + organizationRows = [{ creditBalance: '0', departedMemberUsage: '25' }], + lockedMemberUsageRows = memberUsageRows, +}: { + memberUsageRows?: unknown[] + lockedOwnerRows?: unknown[] + ownerStatsRows?: unknown[] + organizationRows?: unknown[] + lockedMemberUsageRows?: unknown[] +} = {}) { + queueTableRows(schemaMock.member, memberUsageRows) + queueTableRows(schemaMock.member, lockedOwnerRows) + queueTableRows(schemaMock.userStats, ownerStatsRows) + queueTableRows(schemaMock.organization, organizationRows) + queueTableRows(schemaMock.member, lockedMemberUsageRows) } -function buildUpdateChain() { - return { - set: vi.fn(() => ({ - where: vi.fn(async () => []), - })), - } +const usableOrgSubscription = { + plan: 'team', + seats: 2, + periodStart: new Date('2026-05-01T00:00:00.000Z'), + periodEnd: new Date('2026-06-01T00:00:00.000Z'), + stripeSubscriptionId: 'sub_team_1', + stripeCustomerId: 'cus_team_1', } describe('checkAndBillOverageThreshold', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockGetHighestPrioritySubscription.mockResolvedValue(userSubscription) mockGetEffectiveBillingStatus.mockResolvedValue({ billingBlocked: false }) @@ -253,16 +208,14 @@ describe('checkAndBillOverageThreshold', () => { mockIsEnterprise.mockReturnValue(false) mockIsOrgScopedSubscription.mockReturnValue(false) mockGetBillingPeriodUsageCost.mockResolvedValue(0) - mockDbSelect.mockImplementation(() => buildPersonalSelectChain()) - mockTxSelect.mockImplementation(() => buildStatsSelectChain()) - mockTxUpdate.mockImplementation(() => buildUpdateChain()) - mockTxExecute.mockResolvedValue(undefined) - mockDbTransaction.mockImplementation(async (callback: (tx: MockTx) => Promise) => - callback({ execute: mockTxExecute, select: mockTxSelect, update: mockTxUpdate }) - ) + }) + + afterAll(() => { + resetDbChainMock() }) it('does not lock user_stats when calculated overage is below threshold', async () => { + queuePersonalReads() mockCalculateSubscriptionOverage.mockResolvedValue(99) await checkAndBillOverageThreshold('user-1') @@ -275,18 +228,20 @@ describe('checkAndBillOverageThreshold', () => { periodStart: userSubscription.periodStart, periodEnd: userSubscription.periodEnd, }) - expect(mockDbTransaction).not.toHaveBeenCalled() - expect(mockDbSelect).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.transaction).not.toHaveBeenCalled() + expect(dbChainMockFns.select).toHaveBeenCalledTimes(1) expect(mockEnqueueOutboxEvent).not.toHaveBeenCalled() }) it('preserves best-effort error handling for existing callers', async () => { + queuePersonalReads() mockCalculateSubscriptionOverage.mockRejectedValue(new Error('Overage lookup unavailable')) await expect(checkAndBillOverageThreshold('user-1')).resolves.toBeUndefined() }) it('wraps provider failures when strict settlement has no expected billing period', async () => { + queuePersonalReads() mockCalculateSubscriptionOverage.mockRejectedValue(new Error('Overage lookup unavailable')) await expect( @@ -299,6 +254,7 @@ describe('checkAndBillOverageThreshold', () => { }) it('wraps provider failures as retryable errors for a frozen modern period', async () => { + queuePersonalReads() mockCalculateSubscriptionOverage.mockRejectedValue(new Error('Overage lookup unavailable')) await expect( @@ -329,7 +285,7 @@ describe('checkAndBillOverageThreshold', () => { }) expect(mockCalculateSubscriptionOverage).not.toHaveBeenCalled() - expect(mockDbTransaction).not.toHaveBeenCalled() + expect(dbChainMockFns.transaction).not.toHaveBeenCalled() expect(mockEnqueueOutboxEvent).not.toHaveBeenCalled() }) @@ -349,6 +305,7 @@ describe('checkAndBillOverageThreshold', () => { }) it('fails retryably when an above-threshold modern settlement lacks payment state', async () => { + queuePersonalReads() mockGetHighestPrioritySubscription.mockResolvedValue({ ...userSubscription, stripeSubscriptionId: null, @@ -366,11 +323,12 @@ describe('checkAndBillOverageThreshold', () => { retryable: true, }) - expect(mockDbTransaction).not.toHaveBeenCalled() + expect(dbChainMockFns.transaction).not.toHaveBeenCalled() expect(mockEnqueueOutboxEvent).not.toHaveBeenCalled() }) it('throws retryably for markerless strict settlement when payment state is missing', async () => { + queuePersonalReads() mockGetHighestPrioritySubscription.mockResolvedValue({ ...userSubscription, stripeSubscriptionId: null, @@ -385,7 +343,7 @@ describe('checkAndBillOverageThreshold', () => { retryable: true, }) - expect(mockDbTransaction).not.toHaveBeenCalled() + expect(dbChainMockFns.transaction).not.toHaveBeenCalled() expect(mockEnqueueOutboxEvent).not.toHaveBeenCalled() }) @@ -410,19 +368,11 @@ describe('checkAndBillOverageThreshold', () => { name: 'already settled', prepare: () => { mockCalculateSubscriptionOverage.mockResolvedValue(250) - mockTxStatsLimit.mockResolvedValue([ - { - currentPeriodCost: '0', - proPeriodCostSnapshot: '0', - proPeriodCostSnapshotAt: null, - lastPeriodCost: '0', - billedOverageThisPeriod: '250', - creditBalance: '0', - }, - ]) + queueLockedStats(lockedStatsRow({ billedOverageThisPeriod: '250' })) }, }, ])('keeps the $name terminal no-op successful in markerless strict mode', async ({ prepare }) => { + queuePersonalReads() prepare() await expect( @@ -431,6 +381,7 @@ describe('checkAndBillOverageThreshold', () => { }) it('returns a distinct modern no-op when overage is below threshold', async () => { + queuePersonalReads() mockCalculateSubscriptionOverage.mockResolvedValue(99) await expect( @@ -455,14 +406,7 @@ describe('checkAndBillOverageThreshold', () => { }) it('wraps organization provider failures through the strict payer helper', async () => { - mockGetOrganizationSubscriptionUsable.mockResolvedValue({ - plan: 'team', - seats: 2, - periodStart: expectedBillingPeriod.start, - periodEnd: expectedBillingPeriod.end, - stripeSubscriptionId: 'sub_team_1', - stripeCustomerId: 'cus_team_1', - }) + mockGetOrganizationSubscriptionUsable.mockResolvedValue(usableOrgSubscription) mockIsOrganizationBillingBlocked.mockRejectedValue(new Error('Organization lookup unavailable')) await expect( @@ -479,14 +423,7 @@ describe('checkAndBillOverageThreshold', () => { it('keeps billing-blocked organizations as terminal no-ops in markerless strict mode', async () => { mockIsOrgScopedSubscription.mockReturnValue(true) - mockGetOrganizationSubscriptionUsable.mockResolvedValue({ - plan: 'team', - seats: 2, - periodStart: expectedBillingPeriod.start, - periodEnd: expectedBillingPeriod.end, - stripeSubscriptionId: 'sub_team_1', - stripeCustomerId: 'cus_team_1', - }) + mockGetOrganizationSubscriptionUsable.mockResolvedValue(usableOrgSubscription) mockIsOrganizationBillingBlocked.mockResolvedValue(true) await expect( @@ -516,52 +453,27 @@ describe('checkAndBillOverageThreshold', () => { }) it('calculates overage before opening the short user_stats transaction', async () => { + queuePersonalReads() + queueLockedStats(lockedStatsRow()) mockCalculateSubscriptionOverage.mockResolvedValue(250) - mockTxStatsLimit.mockResolvedValue([ - { - currentPeriodCost: '0', - proPeriodCostSnapshot: '0', - proPeriodCostSnapshotAt: null, - lastPeriodCost: '0', - billedOverageThisPeriod: '0', - creditBalance: '0', - }, - ]) await checkAndBillOverageThreshold('user-1') expect(mockCalculateSubscriptionOverage).toHaveBeenCalled() - expect(mockDbTransaction).toHaveBeenCalled() + expect(dbChainMockFns.transaction).toHaveBeenCalled() expect(mockCalculateSubscriptionOverage.mock.invocationCallOrder[0]).toBeLessThan( - mockDbTransaction.mock.invocationCallOrder[0] + dbChainMockFns.transaction.mock.invocationCallOrder[0] ) - expect(mockTxExecute).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.execute).toHaveBeenCalledTimes(1) expect(mockEnqueueOutboxEvent).toHaveBeenCalledTimes(1) }) it('emits audit and analytics once when a retry finds overage already settled', async () => { mockCalculateSubscriptionOverage.mockResolvedValue(250) - mockTxStatsLimit - .mockResolvedValueOnce([ - { - currentPeriodCost: '0', - proPeriodCostSnapshot: '0', - proPeriodCostSnapshotAt: null, - lastPeriodCost: '0', - billedOverageThisPeriod: '0', - creditBalance: '0', - }, - ]) - .mockResolvedValueOnce([ - { - currentPeriodCost: '0', - proPeriodCostSnapshot: '0', - proPeriodCostSnapshotAt: null, - lastPeriodCost: '0', - billedOverageThisPeriod: '250', - creditBalance: '0', - }, - ]) + queuePersonalReads() + queueLockedStats(lockedStatsRow()) + queuePersonalReads() + queueLockedStats(lockedStatsRow({ billedOverageThisPeriod: '250' })) await checkAndBillOverageThreshold('user-1', undefined, { onError: 'throw' }) await checkAndBillOverageThreshold('user-1', undefined, { onError: 'throw' }) @@ -572,17 +484,9 @@ describe('checkAndBillOverageThreshold', () => { }) it('distinguishes an already-settled modern period without duplicating side effects', async () => { + queuePersonalReads() + queueLockedStats(lockedStatsRow({ billedOverageThisPeriod: '250' })) mockCalculateSubscriptionOverage.mockResolvedValue(250) - mockTxStatsLimit.mockResolvedValue([ - { - currentPeriodCost: '0', - proPeriodCostSnapshot: '0', - proPeriodCostSnapshotAt: null, - lastPeriodCost: '0', - billedOverageThisPeriod: '250', - creditBalance: '0', - }, - ]) await expect( checkAndBillOverageThreshold('user-1', undefined, { @@ -597,64 +501,34 @@ describe('checkAndBillOverageThreshold', () => { }) it('rechecks billed overage while locked before enqueueing an invoice', async () => { + queuePersonalReads() + queueLockedStats(lockedStatsRow({ billedOverageThisPeriod: '200' })) mockCalculateSubscriptionOverage.mockResolvedValue(250) - mockTxStatsLimit.mockResolvedValue([ - { - currentPeriodCost: '0', - proPeriodCostSnapshot: '0', - proPeriodCostSnapshotAt: null, - lastPeriodCost: '0', - billedOverageThisPeriod: '200', - creditBalance: '0', - }, - ]) await checkAndBillOverageThreshold('user-1') - expect(mockDbTransaction).toHaveBeenCalled() - expect(mockTxExecute).toHaveBeenCalledTimes(1) - expect(mockTxUpdate).not.toHaveBeenCalled() + expect(dbChainMockFns.transaction).toHaveBeenCalled() + expect(dbChainMockFns.execute).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.update).not.toHaveBeenCalled() expect(mockEnqueueOutboxEvent).not.toHaveBeenCalled() }) it('skips personal threshold billing when locked usage inputs changed', async () => { + queuePersonalReads({ ...defaultUsageSnapshotRow, currentPeriodCost: '250' }) + queueLockedStats(lockedStatsRow({ lastPeriodCost: '250' })) mockCalculateSubscriptionOverage.mockResolvedValue(250) - mockDbSelect - .mockImplementationOnce(() => buildPersonalSnapshotSelectChain({ currentPeriodCost: '250' })) - .mockImplementationOnce(() => buildPersonalSelectChain()) - mockTxStatsLimit.mockResolvedValue([ - { - currentPeriodCost: '0', - proPeriodCostSnapshot: '0', - proPeriodCostSnapshotAt: null, - lastPeriodCost: '250', - billedOverageThisPeriod: '0', - creditBalance: '0', - }, - ]) await checkAndBillOverageThreshold('user-1') - expect(mockDbTransaction).toHaveBeenCalled() - expect(mockTxUpdate).not.toHaveBeenCalled() + expect(dbChainMockFns.transaction).toHaveBeenCalled() + expect(dbChainMockFns.update).not.toHaveBeenCalled() expect(mockEnqueueOutboxEvent).not.toHaveBeenCalled() }) it('throws retryably in markerless strict mode when locked personal usage changes', async () => { + queuePersonalReads({ ...defaultUsageSnapshotRow, currentPeriodCost: '250' }) + queueLockedStats(lockedStatsRow({ lastPeriodCost: '250' })) mockCalculateSubscriptionOverage.mockResolvedValue(250) - mockDbSelect - .mockImplementationOnce(() => buildPersonalSnapshotSelectChain({ currentPeriodCost: '250' })) - .mockImplementationOnce(() => buildPersonalSelectChain()) - mockTxStatsLimit.mockResolvedValue([ - { - currentPeriodCost: '0', - proPeriodCostSnapshot: '0', - proPeriodCostSnapshotAt: null, - lastPeriodCost: '250', - billedOverageThisPeriod: '0', - creditBalance: '0', - }, - ]) await expect( checkAndBillOverageThreshold('user-1', undefined, { onError: 'throw' }) @@ -663,13 +537,16 @@ describe('checkAndBillOverageThreshold', () => { code: 'concurrent_state_change', retryable: true, }) - expect(mockTxUpdate).not.toHaveBeenCalled() + expect(dbChainMockFns.update).not.toHaveBeenCalled() expect(mockEnqueueOutboxEvent).not.toHaveBeenCalled() }) it('wraps lock timeouts in markerless strict mode', async () => { + queuePersonalReads() mockCalculateSubscriptionOverage.mockResolvedValue(250) - mockDbTransaction.mockRejectedValueOnce(new Error('canceling statement due to lock timeout')) + dbChainMockFns.transaction.mockRejectedValueOnce( + new Error('canceling statement due to lock timeout') + ) await expect( checkAndBillOverageThreshold('user-1', undefined, { onError: 'throw' }) @@ -683,41 +560,13 @@ describe('checkAndBillOverageThreshold', () => { it('computes organization overage before opening the locked transaction', async () => { mockIsOrgScopedSubscription.mockReturnValue(true) mockIsOrganizationBillingBlocked.mockResolvedValue(false) - mockGetOrganizationSubscriptionUsable.mockResolvedValue({ - plan: 'team', - seats: 2, - periodStart: new Date('2026-05-01T00:00:00.000Z'), - periodEnd: new Date('2026-06-01T00:00:00.000Z'), - stripeSubscriptionId: 'sub_team_1', - stripeCustomerId: 'cus_team_1', - }) - mockDbSelect.mockImplementationOnce(() => - buildSelectChain([ - { - userId: 'owner-1', - role: 'owner', - currentPeriodCost: '350', - departedMemberUsage: '25', - }, - ]) - ) + mockGetOrganizationSubscriptionUsable.mockResolvedValue(usableOrgSubscription) + queueOrgReads() mockComputeOrgOverageAmount.mockResolvedValue({ totalOverage: 250, baseSubscriptionAmount: 100, effectiveUsage: 350, }) - mockTxStatsLimit - .mockResolvedValueOnce([{ userId: 'owner-1' }]) - .mockResolvedValueOnce([{ billedOverageThisPeriod: '0' }]) - .mockResolvedValueOnce([{ creditBalance: '0', departedMemberUsage: '25' }]) - .mockResolvedValueOnce([ - { - userId: 'owner-1', - role: 'owner', - currentPeriodCost: '350', - departedMemberUsage: '25', - }, - ]) await checkAndBillOverageThreshold('user-1') @@ -731,143 +580,69 @@ describe('checkAndBillOverageThreshold', () => { departedMemberUsage: 25, memberIds: ['owner-1'], }) - expect(mockDbTransaction).toHaveBeenCalled() + expect(dbChainMockFns.transaction).toHaveBeenCalled() expect(mockComputeOrgOverageAmount.mock.invocationCallOrder[0]).toBeLessThan( - mockDbTransaction.mock.invocationCallOrder[0] + dbChainMockFns.transaction.mock.invocationCallOrder[0] ) - expect(mockTxExecute).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.execute).toHaveBeenCalledTimes(1) expect(mockEnqueueOutboxEvent).toHaveBeenCalledTimes(1) }) it('skips stale organization overage when locked usage inputs changed', async () => { mockIsOrgScopedSubscription.mockReturnValue(true) mockIsOrganizationBillingBlocked.mockResolvedValue(false) - mockGetOrganizationSubscriptionUsable.mockResolvedValue({ - plan: 'team', - seats: 2, - periodStart: new Date('2026-05-01T00:00:00.000Z'), - periodEnd: new Date('2026-06-01T00:00:00.000Z'), - stripeSubscriptionId: 'sub_team_1', - stripeCustomerId: 'cus_team_1', + mockGetOrganizationSubscriptionUsable.mockResolvedValue(usableOrgSubscription) + queueOrgReads({ + organizationRows: [{ creditBalance: '0', departedMemberUsage: '75' }], + lockedMemberUsageRows: [{ ...orgMemberUsageRow, departedMemberUsage: '75' }], }) - mockDbSelect.mockImplementationOnce(() => - buildSelectChain([ - { - userId: 'owner-1', - role: 'owner', - currentPeriodCost: '350', - departedMemberUsage: '25', - }, - ]) - ) mockComputeOrgOverageAmount.mockResolvedValue({ totalOverage: 250, baseSubscriptionAmount: 100, effectiveUsage: 350, }) - mockTxStatsLimit - .mockResolvedValueOnce([{ userId: 'owner-1' }]) - .mockResolvedValueOnce([{ billedOverageThisPeriod: '0' }]) - .mockResolvedValueOnce([{ creditBalance: '0', departedMemberUsage: '75' }]) - .mockResolvedValueOnce([ - { - userId: 'owner-1', - role: 'owner', - currentPeriodCost: '350', - departedMemberUsage: '75', - }, - ]) await checkAndBillOverageThreshold('user-1') - expect(mockDbTransaction).toHaveBeenCalled() + expect(dbChainMockFns.transaction).toHaveBeenCalled() expect(mockEnqueueOutboxEvent).not.toHaveBeenCalled() - expect(mockTxUpdate).not.toHaveBeenCalled() + expect(dbChainMockFns.update).not.toHaveBeenCalled() }) it('rechecks organization billed overage on the locked owner tracker', async () => { mockIsOrgScopedSubscription.mockReturnValue(true) mockIsOrganizationBillingBlocked.mockResolvedValue(false) - mockGetOrganizationSubscriptionUsable.mockResolvedValue({ - plan: 'team', - seats: 2, - periodStart: new Date('2026-05-01T00:00:00.000Z'), - periodEnd: new Date('2026-06-01T00:00:00.000Z'), - stripeSubscriptionId: 'sub_team_1', - stripeCustomerId: 'cus_team_1', - }) - mockDbSelect.mockImplementationOnce(() => - buildSelectChain([ - { - userId: 'owner-1', - role: 'owner', - currentPeriodCost: '350', - departedMemberUsage: '25', - }, - ]) - ) + mockGetOrganizationSubscriptionUsable.mockResolvedValue(usableOrgSubscription) + queueOrgReads({ ownerStatsRows: [{ billedOverageThisPeriod: '200' }] }) mockComputeOrgOverageAmount.mockResolvedValue({ totalOverage: 250, baseSubscriptionAmount: 100, effectiveUsage: 350, }) - mockTxStatsLimit - .mockResolvedValueOnce([{ userId: 'owner-1' }]) - .mockResolvedValueOnce([{ billedOverageThisPeriod: '200' }]) - .mockResolvedValueOnce([{ creditBalance: '0', departedMemberUsage: '25' }]) - .mockResolvedValueOnce([ - { - userId: 'owner-1', - role: 'owner', - currentPeriodCost: '350', - departedMemberUsage: '25', - }, - ]) await checkAndBillOverageThreshold('user-1') - expect(mockDbTransaction).toHaveBeenCalled() + expect(dbChainMockFns.transaction).toHaveBeenCalled() expect(mockEnqueueOutboxEvent).not.toHaveBeenCalled() - expect(mockTxUpdate).not.toHaveBeenCalled() + expect(dbChainMockFns.update).not.toHaveBeenCalled() }) it('skips stale organization overage when owner identity changed', async () => { mockIsOrgScopedSubscription.mockReturnValue(true) mockIsOrganizationBillingBlocked.mockResolvedValue(false) - mockGetOrganizationSubscriptionUsable.mockResolvedValue({ - plan: 'team', - seats: 2, - periodStart: new Date('2026-05-01T00:00:00.000Z'), - periodEnd: new Date('2026-06-01T00:00:00.000Z'), - stripeSubscriptionId: 'sub_team_1', - stripeCustomerId: 'cus_team_1', - }) - mockDbSelect.mockImplementationOnce(() => - buildSelectChain([ - { - userId: 'owner-1', - role: 'owner', - currentPeriodCost: '350', - departedMemberUsage: '25', - }, + mockGetOrganizationSubscriptionUsable.mockResolvedValue(usableOrgSubscription) + queueOrgReads({ + memberUsageRows: [ + orgMemberUsageRow, { userId: 'member-1', role: 'member', currentPeriodCost: '25', departedMemberUsage: '25', }, - ]) - ) - mockComputeOrgOverageAmount.mockResolvedValue({ - totalOverage: 250, - baseSubscriptionAmount: 100, - effectiveUsage: 350, - }) - mockTxStatsLimit - .mockResolvedValueOnce([{ userId: 'member-1' }]) - .mockResolvedValueOnce([{ billedOverageThisPeriod: '0' }]) - .mockResolvedValueOnce([{ creditBalance: '0', departedMemberUsage: '25' }]) - .mockResolvedValueOnce([ + ], + lockedOwnerRows: [{ userId: 'member-1' }], + lockedMemberUsageRows: [ { userId: 'owner-1', role: 'member', @@ -880,12 +655,18 @@ describe('checkAndBillOverageThreshold', () => { currentPeriodCost: '25', departedMemberUsage: '25', }, - ]) + ], + }) + mockComputeOrgOverageAmount.mockResolvedValue({ + totalOverage: 250, + baseSubscriptionAmount: 100, + effectiveUsage: 350, + }) await checkAndBillOverageThreshold('user-1') - expect(mockDbTransaction).toHaveBeenCalled() + expect(dbChainMockFns.transaction).toHaveBeenCalled() expect(mockEnqueueOutboxEvent).not.toHaveBeenCalled() - expect(mockTxUpdate).not.toHaveBeenCalled() + expect(dbChainMockFns.update).not.toHaveBeenCalled() }) }) diff --git a/apps/sim/lib/billing/webhooks/enterprise-reconciliation-lease.test.ts b/apps/sim/lib/billing/webhooks/enterprise-reconciliation-lease.test.ts index 2b9dcb08069..8fcff1ee94e 100644 --- a/apps/sim/lib/billing/webhooks/enterprise-reconciliation-lease.test.ts +++ b/apps/sim/lib/billing/webhooks/enterprise-reconciliation-lease.test.ts @@ -1,10 +1,10 @@ /** * @vitest-environment node */ -import { describe, expect, it, vi } from 'vitest' +import { dbChainMock, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' -vi.mock('@sim/db', () => ({ db: {} })) -vi.mock('@sim/db/schema', () => ({ idempotencyKey: {} })) +vi.mock('@sim/db', () => dbChainMock) import { assertEnterpriseReconciliationLeaseHeld, @@ -15,6 +15,15 @@ import { } from '@/lib/billing/webhooks/enterprise-reconciliation-lease' import type { DbOrTx } from '@/lib/db/types' +beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() +}) + +afterAll(() => { + resetDbChainMock() +}) + function inMemoryLeaseStore(): EnterpriseReconciliationLeaseStore { let held: EnterpriseReconciliationLease | null = null let nextToken = 0 diff --git a/apps/sim/lib/billing/webhooks/outbox-handlers.test.ts b/apps/sim/lib/billing/webhooks/outbox-handlers.test.ts index d1f9a9748e7..5cf73d5cb85 100644 --- a/apps/sim/lib/billing/webhooks/outbox-handlers.test.ts +++ b/apps/sim/lib/billing/webhooks/outbox-handlers.test.ts @@ -1,36 +1,25 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const { mockGetPlanByName, mockResolveDefaultPaymentMethod, queryQueue, stripeMock } = vi.hoisted( - () => { - const stripeMock = { - subscriptions: { - retrieve: vi.fn(), - update: vi.fn(), - }, - } - return { - mockGetPlanByName: vi.fn(), - mockResolveDefaultPaymentMethod: vi.fn(), - queryQueue: { value: [] as unknown[][] }, - stripeMock, - } +import { dbChainMock, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockGetPlanByName, mockResolveDefaultPaymentMethod, stripeMock } = vi.hoisted(() => { + const stripeMock = { + subscriptions: { + retrieve: vi.fn(), + update: vi.fn(), + }, } -) - -vi.mock('@sim/db', () => { - const makeChain = () => { - const chain: Record = {} - chain.from = () => chain - chain.where = () => chain - chain.limit = () => Promise.resolve(queryQueue.value.shift() ?? []) - return chain + return { + mockGetPlanByName: vi.fn(), + mockResolveDefaultPaymentMethod: vi.fn(), + stripeMock, } - return { db: { select: () => makeChain() } } }) +vi.mock('@sim/db', () => dbChainMock) + vi.mock('@/lib/billing/stripe-client', () => ({ requireStripeClient: () => stripeMock, })) @@ -76,10 +65,17 @@ function stripeItem(overrides: { } } +/** Queues the handler's pre-Stripe and re-verification subscription reads. */ +function queueSubscriptionReads(rowSets: unknown[][]) { + for (const rows of rowSets) { + queueTableRows(schemaMock.subscription, rows) + } +} + describe('stripeSyncSubscriptionSeats outbox handler', () => { beforeEach(() => { vi.clearAllMocks() - queryQueue.value = [] + resetDbChainMock() mockGetPlanByName.mockReturnValue({ priceId: 'price_team_month', annualDiscountPriceId: 'price_team_year', @@ -87,6 +83,10 @@ describe('stripeSyncSubscriptionSeats outbox handler', () => { stripeMock.subscriptions.update.mockResolvedValue({}) }) + afterAll(() => { + resetDbChainMock() + }) + it('reconciles both price and quantity for a Pro→Team conversion', async () => { const row = { plan: 'team_6000', @@ -94,7 +94,7 @@ describe('stripeSyncSubscriptionSeats outbox handler', () => { status: 'active', stripeSubscriptionId: 'stripe_sub', } - queryQueue.value = [[row], [row]] + queueSubscriptionReads([[row], [row]]) stripeMock.subscriptions.retrieve.mockResolvedValue( stripeItem({ quantity: 1, priceId: 'price_pro_month' }) ) @@ -118,7 +118,7 @@ describe('stripeSyncSubscriptionSeats outbox handler', () => { status: 'past_due', stripeSubscriptionId: 'stripe_sub', } - queryQueue.value = [[row], [row]] + queueSubscriptionReads([[row], [row]]) stripeMock.subscriptions.retrieve.mockResolvedValue( stripeItem({ quantity: 1, priceId: 'price_team_month', status: 'past_due' }) ) @@ -139,7 +139,7 @@ describe('stripeSyncSubscriptionSeats outbox handler', () => { status: 'active', stripeSubscriptionId: 'stripe_sub', } - queryQueue.value = [[row], [row]] + queueSubscriptionReads([[row], [row]]) stripeMock.subscriptions.retrieve.mockResolvedValue( stripeItem({ quantity: 1, priceId: 'price_pro_year', interval: 'year' }) ) @@ -162,7 +162,7 @@ describe('stripeSyncSubscriptionSeats outbox handler', () => { status: 'active', stripeSubscriptionId: 'stripe_sub', } - queryQueue.value = [[row], [row]] + queueSubscriptionReads([[row], [row]]) stripeMock.subscriptions.retrieve.mockResolvedValue( stripeItem({ quantity: 2, priceId: 'price_team_month' }) ) @@ -183,7 +183,7 @@ describe('stripeSyncSubscriptionSeats outbox handler', () => { status: 'active', stripeSubscriptionId: 'stripe_sub', } - queryQueue.value = [[row], [row]] + queueSubscriptionReads([[row], [row]]) stripeMock.subscriptions.retrieve.mockResolvedValue( stripeItem({ quantity: 2, priceId: 'price_team_month' }) ) @@ -194,9 +194,9 @@ describe('stripeSyncSubscriptionSeats outbox handler', () => { }) it('skips non-Team subscriptions', async () => { - queryQueue.value = [ + queueSubscriptionReads([ [{ plan: 'pro_6000', seats: 1, status: 'active', stripeSubscriptionId: 's' }], - ] + ]) await seatSyncHandler({ subscriptionId: 'sub-1' }, ctx) diff --git a/apps/sim/lib/credentials/access.test.ts b/apps/sim/lib/credentials/access.test.ts index fcba37053f8..0fc2c8dc2f8 100644 --- a/apps/sim/lib/credentials/access.test.ts +++ b/apps/sim/lib/credentials/access.test.ts @@ -1,43 +1,9 @@ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { credential, credentialMember } from '@sim/db/schema' +import { queueTableRows, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' -const { mockCheckWorkspaceAccess, dbState } = vi.hoisted(() => ({ +const { mockCheckWorkspaceAccess } = vi.hoisted(() => ({ mockCheckWorkspaceAccess: vi.fn(), - dbState: { results: [] as any[][] }, -})) - -function makeChain() { - const chain: any = {} - chain.from = vi.fn(() => chain) - chain.where = vi.fn(() => chain) - chain.limit = vi.fn(() => Promise.resolve(dbState.results.shift() ?? [])) - return chain -} - -vi.mock('@sim/db', () => ({ - db: { select: vi.fn(() => makeChain()) }, -})) - -vi.mock('@sim/db/schema', () => ({ - credentialTypeEnum: { - enumValues: ['oauth', 'env_workspace', 'env_personal', 'service_account'], - }, - credential: { - id: 'credential.id', - workspaceId: 'credential.workspaceId', - type: 'credential.type', - }, - credentialMember: { - credentialId: 'credentialMember.credentialId', - userId: 'credentialMember.userId', - status: 'credentialMember.status', - role: 'credentialMember.role', - }, -})) - -vi.mock('drizzle-orm', () => ({ - and: vi.fn((...args: unknown[]) => ({ and: args })), - eq: vi.fn((a: unknown, b: unknown) => ({ eq: [a, b] })), - inArray: vi.fn((a: unknown, b: unknown) => ({ inArray: [a, b] })), })) vi.mock('@/lib/workspaces/permissions/utils', () => ({ @@ -46,17 +12,20 @@ vi.mock('@/lib/workspaces/permissions/utils', () => ({ import { getCredentialActorContext } from '@/lib/credentials/access' +afterAll(resetDbChainMock) + const workspaceAdminAccess = { hasAccess: true, canWrite: true, canAdmin: true } const noWorkspaceAccess = { hasAccess: false, canWrite: false, canAdmin: false } describe('getCredentialActorContext', () => { beforeEach(() => { vi.clearAllMocks() - dbState.results = [] + resetDbChainMock() }) it('treats an explicit credential admin membership as admin', async () => { - dbState.results = [[{ id: 'c1', workspaceId: 'ws', type: 'oauth' }], [{ role: 'admin' }]] + queueTableRows(credential, [{ id: 'c1', workspaceId: 'ws', type: 'oauth' }]) + queueTableRows(credentialMember, [{ role: 'admin' }]) mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: true, canWrite: true, canAdmin: false }) const ctx = await getCredentialActorContext('c1', 'user1') @@ -65,7 +34,7 @@ describe('getCredentialActorContext', () => { }) it('derives credential admin from workspace admin for shared credentials', async () => { - dbState.results = [[{ id: 'c1', workspaceId: 'ws', type: 'oauth' }], []] + queueTableRows(credential, [{ id: 'c1', workspaceId: 'ws', type: 'oauth' }]) mockCheckWorkspaceAccess.mockResolvedValue(workspaceAdminAccess) const ctx = await getCredentialActorContext('c1', 'admin-user') @@ -74,7 +43,7 @@ describe('getCredentialActorContext', () => { }) it('does not derive credential admin on personal env credentials', async () => { - dbState.results = [[{ id: 'c1', workspaceId: 'ws', type: 'env_personal' }], []] + queueTableRows(credential, [{ id: 'c1', workspaceId: 'ws', type: 'env_personal' }]) mockCheckWorkspaceAccess.mockResolvedValue(workspaceAdminAccess) const ctx = await getCredentialActorContext('c1', 'admin-user') @@ -83,7 +52,7 @@ describe('getCredentialActorContext', () => { }) it('is not admin for a non-admin without membership', async () => { - dbState.results = [[{ id: 'c1', workspaceId: 'ws', type: 'oauth' }], []] + queueTableRows(credential, [{ id: 'c1', workspaceId: 'ws', type: 'oauth' }]) mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: true, canWrite: false, @@ -96,8 +65,6 @@ describe('getCredentialActorContext', () => { }) it('returns empty context when the credential does not exist', async () => { - dbState.results = [[]] - const ctx = await getCredentialActorContext('missing', 'user1') expect(ctx.credential).toBeNull() @@ -106,7 +73,7 @@ describe('getCredentialActorContext', () => { }) it('exposes workspace access flags from checkWorkspaceAccess', async () => { - dbState.results = [[{ id: 'c1', workspaceId: 'ws', type: 'oauth' }], []] + queueTableRows(credential, [{ id: 'c1', workspaceId: 'ws', type: 'oauth' }]) mockCheckWorkspaceAccess.mockResolvedValue(noWorkspaceAccess) const ctx = await getCredentialActorContext('c1', 'outsider') diff --git a/apps/sim/lib/execution/payloads/large-value-metadata.test.ts b/apps/sim/lib/execution/payloads/large-value-metadata.test.ts index b7b0b337f15..2675a1f5bd2 100644 --- a/apps/sim/lib/execution/payloads/large-value-metadata.test.ts +++ b/apps/sim/lib/execution/payloads/large-value-metadata.test.ts @@ -2,142 +2,10 @@ * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const { - mockAnd, - mockDelete, - mockEq, - mockExecute, - mockInsert, - mockOnConflictDoNothing, - mockSelect, - mockSelectFrom, - mockSelectLimit, - mockSelectWhere, - mockTransaction, - mockTxDelete, - mockTxInsert, - mockTxSelect, - mockTxSelectDistinct, - mockTxSelectFrom, - mockTxSelectLimit, - mockTxSelectWhere, - mockTxValues, - mockValues, - mockWhere, - mockTxWhere, - mockNotInArray, -} = vi.hoisted(() => { - const mockOnConflictDoNothing = vi.fn(async () => undefined) - const mockValues = vi.fn(() => ({ onConflictDoNothing: mockOnConflictDoNothing })) - const mockInsert = vi.fn(() => ({ values: mockValues })) - const mockWhere = vi.fn(async () => undefined) - const mockDelete = vi.fn(() => ({ where: mockWhere })) - const mockSelectLimit = vi.fn(async () => []) - const mockSelectWhere = vi.fn(() => ({ limit: mockSelectLimit })) - const mockSelectFrom = vi.fn(() => ({ where: mockSelectWhere })) - const mockSelect = vi.fn(() => ({ from: mockSelectFrom })) - - const mockTxValues = vi.fn(() => ({ onConflictDoNothing: mockOnConflictDoNothing })) - const mockTxInsert = vi.fn(() => ({ values: mockTxValues })) - const mockTxWhere = vi.fn(async () => undefined) - const mockTxDelete = vi.fn(() => ({ where: mockTxWhere })) - const mockTxSelectLimit = vi.fn(async () => []) - const mockTxSelectWhere = vi.fn(() => ({ limit: mockTxSelectLimit })) - const mockTxSelectFrom = vi.fn(() => ({ where: mockTxSelectWhere })) - const mockTxSelect = vi.fn(() => ({ from: mockTxSelectFrom })) - const mockTxSelectDistinct = vi.fn(() => ({ from: mockTxSelectFrom })) - - return { - mockAnd: vi.fn((...args: unknown[]) => ({ op: 'and', args })), - mockDelete, - mockEq: vi.fn((...args: unknown[]) => ({ op: 'eq', args })), - mockExecute: vi.fn(async () => [{ count: 0 }]), - mockInsert, - mockNotInArray: vi.fn((...args: unknown[]) => ({ op: 'notInArray', args })), - mockOnConflictDoNothing, - mockSelect, - mockSelectFrom, - mockSelectLimit, - mockSelectWhere, - mockTransaction: vi.fn(async (callback) => - callback({ - delete: mockTxDelete, - insert: mockTxInsert, - select: mockTxSelect, - selectDistinct: mockTxSelectDistinct, - }) - ), - mockTxDelete, - mockTxInsert, - mockTxSelect, - mockTxSelectDistinct, - mockTxSelectFrom, - mockTxSelectLimit, - mockTxSelectWhere, - mockTxValues, - mockValues, - mockWhere, - mockTxWhere, - } -}) - -vi.mock('@sim/db', () => { - const db = { - delete: mockDelete, - execute: mockExecute, - insert: mockInsert, - select: mockSelect, - transaction: mockTransaction, - } - return { - db, - // Exec-pool client shares the instance so the seeded chains still apply. - dbFor: () => db, - } -}) - -vi.mock('@sim/db/schema', () => ({ - executionLargeValueDependencies: { - childKey: 'executionLargeValueDependencies.childKey', - parentKey: 'executionLargeValueDependencies.parentKey', - workspaceId: 'executionLargeValueDependencies.workspaceId', - }, - executionLargeValueReferences: { - executionId: 'executionLargeValueReferences.executionId', - key: 'executionLargeValueReferences.key', - source: 'executionLargeValueReferences.source', - workspaceId: 'executionLargeValueReferences.workspaceId', - }, - executionLargeValues: { - key: 'executionLargeValues.key', - ownerExecutionId: 'executionLargeValues.ownerExecutionId', - workspaceId: 'executionLargeValues.workspaceId', - }, - pausedExecutions: { - executionId: 'pausedExecutions.executionId', - status: 'pausedExecutions.status', - }, - workflowExecutionLogs: { - executionId: 'workflowExecutionLogs.executionId', - }, -})) - -vi.mock('@sim/logger', () => ({ - createLogger: vi.fn(() => ({ - warn: vi.fn(), - })), -})) - -vi.mock('drizzle-orm', () => ({ - and: mockAnd, - eq: mockEq, - inArray: vi.fn((...args: unknown[]) => ({ op: 'inArray', args })), - notInArray: mockNotInArray, - sql: vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ strings, values })), -})) - +import { executionLargeValueDependencies, executionLargeValueReferences } from '@sim/db/schema' +import { dbChainMockFns, resetDbChainMock } from '@sim/testing' +import { eq, notInArray } from 'drizzle-orm' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' import { addLargeValueReference, MAX_LARGE_VALUE_REFERENCES_PER_SCOPE, @@ -150,9 +18,13 @@ function largeValueKey(id: string, executionId = 'source-execution'): string { return `execution/workspace-1/workflow-1/${executionId}/large-value-lv_${id}.json` } +afterAll(resetDbChainMock) + describe('large value metadata', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() + dbChainMockFns.execute.mockResolvedValue([{ count: 0 }]) }) it('registers valid large value owner metadata', async () => { @@ -165,15 +37,15 @@ describe('large value metadata', () => { }) expect(registered).toBe(true) - expect(mockTxInsert).toHaveBeenCalledOnce() - expect(mockTxValues).toHaveBeenCalledWith({ + expect(dbChainMockFns.insert).toHaveBeenCalledOnce() + expect(dbChainMockFns.values).toHaveBeenCalledWith({ key: 'execution/workspace-1/workflow-1/execution-1/large-value-lv_abcdefghijkl.json', workspaceId: 'workspace-1', workflowId: 'workflow-1', ownerExecutionId: 'execution-1', size: 124, }) - expect(mockOnConflictDoNothing).toHaveBeenCalledOnce() + expect(dbChainMockFns.onConflictDoNothing).toHaveBeenCalledOnce() }) it('skips malformed owner keys', async () => { @@ -186,14 +58,14 @@ describe('large value metadata', () => { }) expect(registered).toBe(false) - expect(mockTxInsert).not.toHaveBeenCalled() + expect(dbChainMockFns.insert).not.toHaveBeenCalled() }) it('records dependency closure for nested large value refs', async () => { const directKey = largeValueKey('abcdefghijkl') const transitiveKey = largeValueKey('mnopqrstuvwx', 'root-execution') const deepTransitiveKey = largeValueKey('deepqrstuvwx', 'deep-execution') - mockTxSelectLimit + dbChainMockFns.limit .mockResolvedValueOnce([{ childKey: transitiveKey }]) .mockResolvedValueOnce([{ childKey: deepTransitiveKey }]) .mockResolvedValueOnce([]) @@ -210,8 +82,8 @@ describe('large value metadata', () => { ) expect(registered).toBe(true) - expect(mockTxSelectDistinct).toHaveBeenCalledTimes(3) - expect(mockTxValues).toHaveBeenLastCalledWith([ + expect(dbChainMockFns.selectDistinct).toHaveBeenCalledTimes(3) + expect(dbChainMockFns.values).toHaveBeenLastCalledWith([ { parentKey: 'execution/workspace-1/workflow-1/execution-1/large-value-lv_zyxwvutsrqpo.json', childKey: directKey, @@ -246,9 +118,9 @@ describe('large value metadata', () => { keys ) - expect(mockTxValues).toHaveBeenCalledTimes(3) - expect(mockTxValues.mock.calls[1]?.[0]).toHaveLength(500) - expect(mockTxValues.mock.calls[2]?.[0]).toHaveLength(1) + expect(dbChainMockFns.values).toHaveBeenCalledTimes(3) + expect(dbChainMockFns.values.mock.calls[1]?.[0]).toHaveLength(500) + expect(dbChainMockFns.values.mock.calls[2]?.[0]).toHaveLength(1) }) it('rejects reference sets over the metadata cardinality limit', async () => { @@ -272,7 +144,7 @@ describe('large value metadata', () => { it('limits dependency closure reads to the remaining reference budget', async () => { const directKey = largeValueKey('a00000000000') - mockTxSelectLimit.mockResolvedValueOnce( + dbChainMockFns.limit.mockResolvedValueOnce( Array.from({ length: MAX_LARGE_VALUE_REFERENCES_PER_SCOPE }, (_, index) => ({ childKey: largeValueKey(`c${index.toString(36).padStart(11, '0')}`), })) @@ -291,7 +163,7 @@ describe('large value metadata', () => { ) ).rejects.toThrow('Large value dependency closure exceeds the limit') - expect(mockTxSelectLimit).toHaveBeenCalledWith(MAX_LARGE_VALUE_REFERENCES_PER_SCOPE) + expect(dbChainMockFns.limit).toHaveBeenCalledWith(MAX_LARGE_VALUE_REFERENCES_PER_SCOPE) }) it('filters known dependency children before applying the remaining reference budget', async () => { @@ -300,13 +172,15 @@ describe('large value metadata', () => { ) const knownChildKey = directKeys[1] const unseenChildKey = largeValueKey('unseenchild1', 'source-execution') - mockTxSelectLimit.mockImplementationOnce(async () => { - const filtersKnownChildren = mockNotInArray.mock.calls.some( - ([field, values]) => - field === 'executionLargeValueDependencies.childKey' && - Array.isArray(values) && - values.includes(knownChildKey) - ) + dbChainMockFns.limit.mockImplementationOnce(async () => { + const filtersKnownChildren = vi + .mocked(notInArray) + .mock.calls.some( + ([field, values]) => + field === executionLargeValueDependencies.childKey && + Array.isArray(values) && + values.includes(knownChildKey) + ) return [{ childKey: filtersKnownChildren ? unseenChildKey : knownChildKey }] }) @@ -323,7 +197,7 @@ describe('large value metadata', () => { ) ).rejects.toThrow('Large value dependency closure exceeds the limit') - expect(mockTxSelectLimit).toHaveBeenCalledWith(1) + expect(dbChainMockFns.limit).toHaveBeenCalledWith(1) }) it('replaces an execution reference set with same-workspace unique keys', async () => { @@ -366,10 +240,10 @@ describe('large value metadata', () => { } ) - expect(mockTransaction).toHaveBeenCalledOnce() - expect(mockTxDelete).toHaveBeenCalledOnce() - expect(mockEq).toHaveBeenCalledWith('executionLargeValueReferences.source', 'execution_log') - expect(mockTxValues).toHaveBeenCalledWith([ + expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() + expect(dbChainMockFns.delete).toHaveBeenCalledOnce() + expect(eq).toHaveBeenCalledWith(executionLargeValueReferences.source, 'execution_log') + expect(dbChainMockFns.values).toHaveBeenCalledWith([ { key: matchingKey, workspaceId: 'workspace-1', @@ -393,9 +267,9 @@ describe('large value metadata', () => { key ) - expect(mockSelectLimit).toHaveBeenCalledWith(1) - expect(mockSelectLimit).toHaveBeenCalledWith(MAX_LARGE_VALUE_REFERENCES_PER_SCOPE + 1) - expect(mockValues).toHaveBeenCalledWith({ + expect(dbChainMockFns.limit).toHaveBeenCalledWith(1) + expect(dbChainMockFns.limit).toHaveBeenCalledWith(MAX_LARGE_VALUE_REFERENCES_PER_SCOPE + 1) + expect(dbChainMockFns.values).toHaveBeenCalledWith({ key, workspaceId: 'workspace-1', workflowId: 'workflow-1', @@ -405,7 +279,7 @@ describe('large value metadata', () => { }) it('rejects materialized references once the scope reaches the reference cap', async () => { - mockSelectLimit.mockResolvedValueOnce([]).mockResolvedValueOnce( + dbChainMockFns.limit.mockResolvedValueOnce([]).mockResolvedValueOnce( Array.from({ length: MAX_LARGE_VALUE_REFERENCES_PER_SCOPE }, (_, index) => ({ key: largeValueKey(`d${index.toString(36).padStart(11, '0')}`), })) @@ -423,11 +297,11 @@ describe('large value metadata', () => { ) ).rejects.toThrow('exceeding the limit') - expect(mockInsert).not.toHaveBeenCalled() + expect(dbChainMockFns.insert).not.toHaveBeenCalled() }) it('prunes large value metadata in bounded batches', async () => { - mockExecute + dbChainMockFns.execute .mockResolvedValueOnce([{ count: 2 }]) .mockResolvedValueOnce([{ count: 3 }]) .mockResolvedValueOnce([{ count: 4 }]) @@ -454,7 +328,7 @@ describe('large value metadata', () => { maxRowsPerTable: 100, }) - const [query] = mockExecute.mock.calls[0] ?? [] + const [query] = dbChainMockFns.execute.mock.calls[0] ?? [] const sqlText = Array.isArray(query?.strings) ? query.strings.join(' ') : '' expect(sqlText).toContain("ref.source = 'execution_log'") expect(sqlText).toContain("ref.source = 'paused_snapshot'") diff --git a/apps/sim/lib/execution/preprocessing.test.ts b/apps/sim/lib/execution/preprocessing.test.ts index b8a2e8e3c24..6aff30c4ea0 100644 --- a/apps/sim/lib/execution/preprocessing.test.ts +++ b/apps/sim/lib/execution/preprocessing.test.ts @@ -22,8 +22,6 @@ const { mockResolveSystemBillingAttribution: vi.fn(), })) -vi.mock('@sim/db', () => ({ db: {} })) -vi.mock('drizzle-orm', () => ({ eq: vi.fn() })) vi.mock('@/lib/auth/ban', () => ({ getActivelyBannedUserIds: mockGetActivelyBannedUserIds, })) diff --git a/apps/sim/lib/execution/preprocessing.webhook-correlation.test.ts b/apps/sim/lib/execution/preprocessing.webhook-correlation.test.ts index 56f54cee6aa..4dc1b8bf9c0 100644 --- a/apps/sim/lib/execution/preprocessing.webhook-correlation.test.ts +++ b/apps/sim/lib/execution/preprocessing.webhook-correlation.test.ts @@ -9,8 +9,6 @@ const { mockResolveSystemBillingAttribution } = vi.hoisted(() => ({ mockResolveSystemBillingAttribution: vi.fn(), })) -vi.mock('@sim/db', () => ({ db: {} })) -vi.mock('drizzle-orm', () => ({ eq: vi.fn() })) vi.mock('@/lib/billing/calculations/usage-monitor', () => ({ checkServerSideUsageLimits: vi.fn(), })) diff --git a/apps/sim/lib/global-work/summary.test.ts b/apps/sim/lib/global-work/summary.test.ts index 4e6317e6b82..b97cf771aa5 100644 --- a/apps/sim/lib/global-work/summary.test.ts +++ b/apps/sim/lib/global-work/summary.test.ts @@ -2,15 +2,15 @@ * @vitest-environment node */ +import { dbChainMockFns, resetDbChainMock } from '@sim/testing' import type { SQL } from 'drizzle-orm' import { PgDialect } from 'drizzle-orm/pg-core' -import { describe, expect, it, vi } from 'vitest' - -const execute = vi.hoisted(() => vi.fn()) +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' vi.unmock('drizzle-orm') vi.unmock('@sim/db/schema') -vi.mock('@sim/db', () => ({ dbReplica: { execute } })) + +const execute = dbChainMockFns.execute const dialect = new PgDialect() @@ -32,7 +32,14 @@ import { getLatestCompletedGlobalWorkMonth, } from '@/lib/global-work/summary' +afterAll(resetDbChainMock) + describe('Global Work Pacific reporting windows', () => { + beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + }) + it('defaults to the latest completed Pacific month across a year boundary', () => { expect(getLatestCompletedGlobalWorkMonth(new Date('2026-01-15T12:00:00.000Z'))).toBe('2025-12') }) diff --git a/apps/sim/lib/knowledge/connectors/queue.test.ts b/apps/sim/lib/knowledge/connectors/queue.test.ts index 070ed03ea55..81c5627e456 100644 --- a/apps/sim/lib/knowledge/connectors/queue.test.ts +++ b/apps/sim/lib/knowledge/connectors/queue.test.ts @@ -1,34 +1,18 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { dbChainMock, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' -const { - mockDbChain, - mockExecuteSync, - mockIsTriggerAvailable, - mockResolveTriggerRegion, - mockTrigger, -} = vi.hoisted(() => { - const chain = { - select: vi.fn().mockReturnThis(), - from: vi.fn().mockReturnThis(), - innerJoin: vi.fn().mockReturnThis(), - where: vi.fn().mockReturnThis(), - limit: vi.fn(), - update: vi.fn().mockReturnThis(), - set: vi.fn().mockReturnThis(), - } - return { - mockDbChain: chain, +const { mockExecuteSync, mockIsTriggerAvailable, mockResolveTriggerRegion, mockTrigger } = + vi.hoisted(() => ({ mockExecuteSync: vi.fn(), mockIsTriggerAvailable: vi.fn(), mockResolveTriggerRegion: vi.fn(), mockTrigger: vi.fn(), - } -}) + })) -vi.mock('@sim/db', () => ({ db: mockDbChain })) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@trigger.dev/sdk', () => ({ tasks: { trigger: mockTrigger } })) vi.mock('@/lib/core/async-jobs/region', () => ({ resolveTriggerRegion: mockResolveTriggerRegion, @@ -66,13 +50,8 @@ const BILLING_ATTRIBUTION = { describe('connector sync queue', () => { beforeEach(() => { vi.clearAllMocks() - mockDbChain.select.mockReturnThis() - mockDbChain.from.mockReturnThis() - mockDbChain.innerJoin.mockReturnThis() - mockDbChain.where.mockReturnThis() - mockDbChain.update.mockReturnThis() - mockDbChain.set.mockReturnThis() - mockDbChain.limit.mockResolvedValue([ + resetDbChainMock() + queueTableRows(schemaMock.knowledgeConnector, [ { knowledgeBaseId: 'knowledge-base-1', connectorArchivedAt: null, @@ -86,6 +65,10 @@ describe('connector sync queue', () => { mockTrigger.mockResolvedValue({ id: 'run-1' }) }) + afterAll(() => { + resetDbChainMock() + }) + it('preserves the actor and immutable workspace payer in the queued payload', async () => { await dispatchSync('connector-1', { billingAttribution: BILLING_ATTRIBUTION, diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts index 5d30ce45dde..faa7fca4717 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts @@ -1,11 +1,11 @@ /** * @vitest-environment node */ -import { authOAuthUtilsMock, urlsMock } from '@sim/testing' +import { authOAuthUtilsMock, dbChainMock, urlsMock } from '@sim/testing' import { generateShortId } from '@sim/utils/id' import { beforeEach, describe, expect, it, vi } from 'vitest' -vi.mock('@sim/db', () => ({ db: {} })) +vi.mock('@sim/db', () => dbChainMock) vi.mock('drizzle-orm', () => ({ and: vi.fn(), eq: vi.fn(), diff --git a/apps/sim/lib/logs/execution/logger.test.ts b/apps/sim/lib/logs/execution/logger.test.ts index 4b83f955541..740986611fb 100644 --- a/apps/sim/lib/logs/execution/logger.test.ts +++ b/apps/sim/lib/logs/execution/logger.test.ts @@ -1,39 +1,10 @@ -import { envFlagsMock } from '@sim/testing' -import { beforeEach, describe, expect, test, vi } from 'vitest' +import { usageLog, workflow } from '@sim/db/schema' +import { dbChainMockFns, envFlagsMock, queueTableRows, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, test, vi } from 'vitest' import { recordUsage } from '@/lib/billing/core/usage-log' import { ExecutionLogger } from '@/lib/logs/execution/logger' -const dbSelectMock = vi.hoisted(() => vi.fn()) -const dbExecuteMock = vi.hoisted(() => vi.fn()) -const txUpdateMock = vi.hoisted(() => - vi.fn(() => ({ set: () => ({ where: () => Promise.resolve() }) })) -) - -vi.mock('@sim/db', () => { - // The reconcile runs inside db.transaction with an advisory lock. The tx - // shares dbSelectMock so the existing call-order seeding (call 1 = workflow - // row via .limit, call 2 = already-billed via .groupBy) still applies; - // tx.execute (set_config + pg_advisory_xact_lock) is a no-op; tx.update backs - // the exact cost_total refine. - const tx = { - select: dbSelectMock, - insert: vi.fn(), - update: txUpdateMock, - execute: dbExecuteMock, - } - const db = { - select: dbSelectMock, - insert: vi.fn(), - update: vi.fn(), - execute: dbExecuteMock, - transaction: vi.fn(async (cb: (txArg: typeof tx) => Promise) => cb(tx)), - } - return { - db, - // Exec-pool client shares the instance so call-order seeding still applies. - dbFor: () => db, - } -}) +afterAll(resetDbChainMock) // Mock billing modules vi.mock('@/lib/billing/core/subscription', () => ({ @@ -149,6 +120,7 @@ describe('ExecutionLogger', () => { beforeEach(() => { logger = new ExecutionLogger() vi.clearAllMocks() + resetDbChainMock() }) describe('class instantiation', () => { @@ -549,6 +521,7 @@ describe('recordExecutionUsage boundary-delta reconciliation', () => { beforeEach(() => { logger = new ExecutionLogger() as any vi.clearAllMocks() + resetDbChainMock() }) const costSummary = (overrides: Record = {}) => ({ @@ -571,22 +544,12 @@ describe('recordExecutionUsage boundary-delta reconciliation', () => { }, } - // db.select() is called twice in recordExecutionUsage: first the workflow row - // (terminated by .limit), then the already-billed usage_log rows (terminated - // by .groupBy). Return each in order. + // recordExecutionUsage reads two tables: the workflow row (from(workflow) + // ... .limit(1)), then the already-billed usage_log rows (from(usageLog) + // ... .groupBy(...)). Route each result set by its table. const mockDb = (billedRows: Array>) => { - let call = 0 - dbSelectMock.mockImplementation(() => { - call += 1 - const rows = call === 1 ? [{ id: 'workflow-1', workspaceId: 'ws-1' }] : billedRows - const chain: any = { - from: () => chain, - where: () => chain, - limit: () => Promise.resolve(rows), - groupBy: () => Promise.resolve(rows), - } - return chain - }) + queueTableRows(workflow, [{ id: 'workflow-1', workspaceId: 'ws-1' }]) + queueTableRows(usageLog, billedRows) } const run = ( @@ -639,12 +602,12 @@ describe('recordExecutionUsage boundary-delta reconciliation', () => { // Returns the amount recorded at this boundary (drives threshold-email math). expect(recorded).toBeCloseTo(1.005, 8) // cost_total is refined to the exact ledger sum inside the locked tx. - expect(txUpdateMock).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.update).toHaveBeenCalledTimes(1) }) test('leaves Mothership model spend to cumulative update-cost while ledgering ordinary models', async () => { const setCostTotalMock = vi.fn(() => ({ where: () => Promise.resolve() })) - txUpdateMock.mockImplementationOnce(() => ({ set: setCostTotalMock })) + dbChainMockFns.update.mockReturnValueOnce({ set: setCostTotalMock }) const recorded = await run( costSummary({ @@ -938,7 +901,7 @@ describe('recordExecutionUsage boundary-delta reconciliation', () => { ) // set_config('lock_timeout') + pg_advisory_xact_lock both run on the tx. - expect(dbExecuteMock).toHaveBeenCalledTimes(2) + expect(dbChainMockFns.execute).toHaveBeenCalledTimes(2) expect(recordUsage).toHaveBeenCalledTimes(1) // The ledger INSERT participates in the locked transaction. expect(vi.mocked(recordUsage).mock.calls[0][0]).toHaveProperty('tx') diff --git a/apps/sim/lib/logs/execution/logging-session.test.ts b/apps/sim/lib/logs/execution/logging-session.test.ts index d07fe618eb5..4e9bf1f1646 100644 --- a/apps/sim/lib/logs/execution/logging-session.test.ts +++ b/apps/sim/lib/logs/execution/logging-session.test.ts @@ -1,39 +1,11 @@ -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const dbMocks = vi.hoisted(() => { - const selectLimit = vi.fn() - const selectWhere = vi.fn() - const selectFrom = vi.fn() - const select = vi.fn() - const updateWhere = vi.fn() - const updateSet = vi.fn() - const update = vi.fn() - const execute = vi.fn() - const eq = vi.fn() - const and = vi.fn((...args: unknown[]) => ({ type: 'and', args })) - const sql = vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ strings, values })) - - select.mockReturnValue({ from: selectFrom }) - selectFrom.mockReturnValue({ where: selectWhere }) - selectWhere.mockReturnValue({ limit: selectLimit }) - - update.mockReturnValue({ set: updateSet }) - updateSet.mockReturnValue({ where: updateWhere }) - - return { - select, - selectFrom, - selectWhere, - selectLimit, - update, - updateSet, - updateWhere, - execute, - eq, - and, - sql, - } -}) +import { dbChainMockFns, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const dbMocks = vi.hoisted(() => ({ + eq: vi.fn(), + and: vi.fn((...args: unknown[]) => ({ type: 'and', args })), + sql: vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ strings, values })), +})) const { completeWorkflowExecutionMock, @@ -47,19 +19,6 @@ const { releaseExecutionSlotMock: vi.fn(), })) -vi.mock('@sim/db', () => { - const db = { - select: dbMocks.select, - update: dbMocks.update, - execute: dbMocks.execute, - } - return { - db, - // Exec-pool client shares the instance so the seeded chains still apply. - dbFor: () => db, - } -}) - vi.mock('drizzle-orm', () => ({ eq: dbMocks.eq, and: dbMocks.and, @@ -116,9 +75,12 @@ vi.mock('@/lib/logs/execution/logging-factory', () => ({ import { calculateCostSummary } from '@/lib/logs/execution/logging-factory' import { LoggingSession } from './logging-session' +afterAll(resetDbChainMock) + describe('LoggingSession start snapshots', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() startWorkflowExecutionMock.mockResolvedValue({}) loadWorkflowStateForExecutionMock.mockResolvedValue({ blocks: { @@ -223,9 +185,8 @@ describe('LoggingSession start snapshots', () => { describe('LoggingSession completion retries', () => { beforeEach(() => { vi.clearAllMocks() - dbMocks.selectLimit.mockResolvedValue([{ executionData: {} }]) - dbMocks.updateWhere.mockResolvedValue(undefined) - dbMocks.execute.mockResolvedValue(undefined) + resetDbChainMock() + dbChainMockFns.limit.mockResolvedValue([{ executionData: {} }]) }) it('keeps completion best-effort when a later error completion retries after full completion and fallback both fail', async () => { @@ -489,8 +450,8 @@ describe('LoggingSession completion retries', () => { await session.onBlockStart('block-1', 'Fetch', 'api', '2025-01-01T00:00:00.000Z') - expect(dbMocks.select).not.toHaveBeenCalled() - expect(dbMocks.execute).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.select).not.toHaveBeenCalled() + expect(dbChainMockFns.execute).toHaveBeenCalledTimes(1) }) it('enforces started marker monotonicity in the database write path', async () => { @@ -499,7 +460,7 @@ describe('LoggingSession completion retries', () => { await session.onBlockStart('block-1', 'Fetch', 'api', '2025-01-01T00:00:00.000Z') expect(dbMocks.sql).toHaveBeenCalled() - expect(dbMocks.execute).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.execute).toHaveBeenCalledTimes(1) }) it('allows same-millisecond started markers to replace the prior marker', async () => { @@ -522,8 +483,8 @@ describe('LoggingSession completion retries', () => { output: { value: true }, }) - expect(dbMocks.select).not.toHaveBeenCalled() - expect(dbMocks.execute).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.select).not.toHaveBeenCalled() + expect(dbChainMockFns.execute).toHaveBeenCalledTimes(1) }) it('allows same-millisecond completed markers to replace the prior marker', async () => { @@ -651,12 +612,11 @@ describe('LoggingSession completion retries', () => { describe('completeWithError cancelled-status guard', () => { beforeEach(() => { vi.clearAllMocks() - dbMocks.updateWhere.mockResolvedValue(undefined) - dbMocks.execute.mockResolvedValue(undefined) + resetDbChainMock() }) it('skips writing failed and marks session complete when DB status is already cancelled', async () => { - dbMocks.selectLimit.mockResolvedValue([{ status: 'cancelled' }]) + dbChainMockFns.limit.mockResolvedValue([{ status: 'cancelled' }]) const session = new LoggingSession('workflow-1', 'execution-1', 'api', 'req-1') await session.safeCompleteWithError({ error: { message: 'block errored mid-cancel' } }) @@ -666,7 +626,7 @@ describe('completeWithError cancelled-status guard', () => { }) it('writes failed when DB status is running (no cancel in flight)', async () => { - dbMocks.selectLimit.mockResolvedValue([{ status: 'running' }]) + dbChainMockFns.limit.mockResolvedValue([{ status: 'running' }]) completeWorkflowExecutionMock.mockResolvedValue({}) const session = new LoggingSession('workflow-1', 'execution-1', 'api', 'req-1') @@ -679,7 +639,7 @@ describe('completeWithError cancelled-status guard', () => { }) it('writes failed when no execution log exists yet', async () => { - dbMocks.selectLimit.mockResolvedValue([]) + dbChainMockFns.limit.mockResolvedValue([]) completeWorkflowExecutionMock.mockResolvedValue({}) const session = new LoggingSession('workflow-1', 'execution-1', 'api', 'req-1') @@ -691,7 +651,7 @@ describe('completeWithError cancelled-status guard', () => { }) it('deduplicates all subsequent completion attempts after guard early-return', async () => { - dbMocks.selectLimit.mockResolvedValue([{ status: 'cancelled' }]) + dbChainMockFns.limit.mockResolvedValue([{ status: 'cancelled' }]) completeWorkflowExecutionMock.mockResolvedValue({}) const session = new LoggingSession('workflow-1', 'execution-1', 'api', 'req-1') @@ -704,7 +664,7 @@ describe('completeWithError cancelled-status guard', () => { }) it('falls through to cost-only fallback when the DB check itself throws', async () => { - dbMocks.selectLimit.mockRejectedValueOnce(new Error('DB connection lost')) + dbChainMockFns.limit.mockRejectedValueOnce(new Error('DB connection lost')) completeWorkflowExecutionMock.mockResolvedValue({}) const session = new LoggingSession('workflow-1', 'execution-1', 'api', 'req-1') @@ -720,28 +680,27 @@ describe('completeWithError cancelled-status guard', () => { describe('LoggingSession.markExecutionAsFailed workflowId scoping', () => { beforeEach(() => { vi.clearAllMocks() - dbMocks.updateWhere.mockResolvedValue(undefined) + resetDbChainMock() }) it('scopes UPDATE by both executionId and workflowId', async () => { await LoggingSession.markExecutionAsFailed('exec-1', undefined, undefined, 'wf-1') - expect(dbMocks.update).toHaveBeenCalledTimes(1) - expect(dbMocks.updateSet).toHaveBeenCalledTimes(1) - expect(dbMocks.updateWhere).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.update).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.set).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.where).toHaveBeenCalledTimes(1) - const whereArgs = dbMocks.updateWhere.mock.calls[0] + const whereArgs = dbChainMockFns.where.mock.calls[0] expect(whereArgs).toBeDefined() }) it('instance markAsFailed forwards workflowId to the static method', async () => { - const updateWhereSpy = dbMocks.updateWhere - dbMocks.selectLimit.mockResolvedValue([{ executionData: {} }]) + dbChainMockFns.limit.mockResolvedValue([{ executionData: {} }]) const session = new LoggingSession('wf-42', 'exec-42', 'api', 'req-1') await session.markAsFailed('something went wrong') - expect(updateWhereSpy).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.update).toHaveBeenCalledTimes(1) expect(releaseExecutionSlotMock).toHaveBeenCalledWith('exec-42') }) @@ -800,7 +759,7 @@ describe('LoggingSession progress-marker write path', () => { loops: {}, parallels: {}, }) - dbMocks.execute.mockResolvedValue(undefined) + resetDbChainMock() }) it('writes markers to Redis (not the row) when Redis accepts the write', async () => { @@ -820,7 +779,7 @@ describe('LoggingSession progress-marker write path', () => { 'exec-redis', expect.objectContaining({ blockId: 'b1', success: true }) ) - expect(dbMocks.execute).not.toHaveBeenCalled() + expect(dbChainMockFns.execute).not.toHaveBeenCalled() }) it('falls back to the SQL UPDATE when the Redis write fails', async () => { @@ -831,6 +790,6 @@ describe('LoggingSession progress-marker write path', () => { await session.onBlockStart('b1', 'Fetch', 'api', '2026-06-27T10:00:00.000Z') expect(setLastStartedBlockMock).toHaveBeenCalled() - expect(dbMocks.execute).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.execute).toHaveBeenCalledTimes(1) }) }) diff --git a/apps/sim/lib/logs/list-logs.test.ts b/apps/sim/lib/logs/list-logs.test.ts index d8ad2cdb6cf..ff6168733fc 100644 --- a/apps/sim/lib/logs/list-logs.test.ts +++ b/apps/sim/lib/logs/list-logs.test.ts @@ -2,18 +2,12 @@ * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { jobExecutionLogs, workflowExecutionLogs } from '@sim/db/schema' +import { dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' -const { selectMock } = vi.hoisted(() => ({ selectMock: vi.fn() })) - -vi.mock('@sim/db', () => { - const instance = { select: selectMock } - return { db: instance, dbReplica: instance } -}) - -// Local drizzle-orm mock: the global mock's `sql` lacks `.as()` and the chain -// mock doesn't support `.orderBy().limit()`. We only need condition/sql builders -// to produce truthy stubs (the mocked db ignores them). +// Local drizzle-orm mock: the global mock's `sql` lacks `.as()`. We only need +// condition/sql builders to produce truthy stubs (the mocked db ignores them). vi.mock('drizzle-orm', () => { const make = (): Record => { const o: Record = {} @@ -64,15 +58,7 @@ vi.mock('@/lib/workspaces/permissions/utils', () => ({ import type { ListLogsParams } from './list-logs' import { decodeCursor, listLogs } from './list-logs' -/** A chainable, thenable query-builder stub that resolves to the given rows. */ -function builder(rows: unknown[]) { - const b: Record = {} - for (const method of ['from', 'leftJoin', 'innerJoin', 'where', 'orderBy', 'limit']) { - b[method] = () => b - } - ;(b as { then: unknown }).then = (resolve: (value: unknown) => unknown) => resolve(rows) - return b -} +afterAll(resetDbChainMock) function workflowRow(overrides: Record = {}) { return { @@ -136,12 +122,12 @@ function baseParams(overrides: Partial = {}): ListLogsParams { describe('listLogs', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() }) it('merges workflow and job rows into summaries', async () => { - selectMock - .mockReturnValueOnce(builder([workflowRow()])) - .mockReturnValueOnce(builder([jobRow()])) + queueTableRows(workflowExecutionLogs, [workflowRow()]) + queueTableRows(jobExecutionLogs, [jobRow()]) const result = await listLogs(baseParams(), 'user-1') @@ -165,14 +151,11 @@ describe('listLogs', () => { it('returns a decodable nextCursor when results exceed the limit', async () => { // limit 1, two workflow rows → page of 1, hasMore true - selectMock - .mockReturnValueOnce( - builder([ - workflowRow({ id: 'log-a', sortValue: new Date('2026-01-02T00:00:00.000Z') }), - workflowRow({ id: 'log-b', sortValue: new Date('2026-01-01T00:00:00.000Z') }), - ]) - ) - .mockReturnValueOnce(builder([])) + queueTableRows(workflowExecutionLogs, [ + workflowRow({ id: 'log-a', sortValue: new Date('2026-01-02T00:00:00.000Z') }), + workflowRow({ id: 'log-b', sortValue: new Date('2026-01-01T00:00:00.000Z') }), + ]) + queueTableRows(jobExecutionLogs, []) const result = await listLogs(baseParams({ limit: 1 }), 'user-1') @@ -183,12 +166,12 @@ describe('listLogs', () => { }) it('excludes job logs when a workflow-specific filter is present', async () => { - selectMock.mockReturnValueOnce(builder([workflowRow()])) + queueTableRows(workflowExecutionLogs, [workflowRow()]) const result = await listLogs(baseParams({ workflowIds: 'wf-1' }), 'user-1') // Only the workflow query runs; the job query is Promise.resolve([]). - expect(selectMock).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(1) expect(result.data).toHaveLength(1) expect(result.data[0].workflowId).toBe('wf-1') }) diff --git a/apps/sim/lib/mcp/oauth/revoke.test.ts b/apps/sim/lib/mcp/oauth/revoke.test.ts index 1e69c99f8d8..73583eb5926 100644 --- a/apps/sim/lib/mcp/oauth/revoke.test.ts +++ b/apps/sim/lib/mcp/oauth/revoke.test.ts @@ -8,7 +8,8 @@ * raw `fetch`. */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { dbChainMock, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const BLOCKED_ENDPOINT = 'http://169.254.170.2/v2/credentials/' const PUBLIC_SERVER_URL = 'https://mcp.attacker.com' @@ -20,14 +21,12 @@ const { mockDiscoverOAuthServerInfo, mockLoadOauthRow, mockDecryptSecret, - mockDbSelect, } = vi.hoisted(() => ({ mockUndiciFetch: vi.fn(), mockValidateMcpServerSsrf: vi.fn(), mockDiscoverOAuthServerInfo: vi.fn(), mockLoadOauthRow: vi.fn(), mockDecryptSecret: vi.fn(), - mockDbSelect: vi.fn(), })) vi.mock('@/lib/core/security/input-validation.server', () => ({ @@ -50,24 +49,22 @@ vi.mock('@/lib/mcp/oauth/storage', () => ({ vi.mock('@/lib/core/security/encryption', () => ({ decryptSecret: mockDecryptSecret, })) -vi.mock('@sim/db', () => ({ - db: { select: mockDbSelect }, -})) +vi.mock('@sim/db', () => dbChainMock) import { revokeMcpOauthTokens } from './revoke' function wireServerRow(row: Record) { - const builder = { - from: () => builder, - where: () => builder, - limit: () => Promise.resolve([row]), - } - mockDbSelect.mockReturnValue(builder) + queueTableRows(schemaMock.mcpServers, [row]) } describe('revokeMcpOauthTokens — SSRF guard', () => { + afterAll(() => { + resetDbChainMock() + }) + beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockLoadOauthRow.mockResolvedValue({ tokens: { access_token: 'access-secret', refresh_token: 'refresh-secret' }, diff --git a/apps/sim/lib/table/cell-write.test.ts b/apps/sim/lib/table/cell-write.test.ts index b4fa0fcffe4..fab52b73106 100644 --- a/apps/sim/lib/table/cell-write.test.ts +++ b/apps/sim/lib/table/cell-write.test.ts @@ -1,23 +1,18 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' import type { RowExecutionMetadata, TableDefinition, WorkflowGroup } from '@/lib/table/types' -const { mockAppendTableEvent, mockTransaction, mockUpdateRow, mockWriteExecutionsPatch } = - vi.hoisted(() => ({ - mockAppendTableEvent: vi.fn(), - mockTransaction: vi.fn(), - mockUpdateRow: vi.fn(), - mockWriteExecutionsPatch: vi.fn(), - })) - -vi.mock('@sim/db', () => ({ - db: { - transaction: mockTransaction, - }, +const { mockAppendTableEvent, mockUpdateRow, mockWriteExecutionsPatch } = vi.hoisted(() => ({ + mockAppendTableEvent: vi.fn(), + mockUpdateRow: vi.fn(), + mockWriteExecutionsPatch: vi.fn(), })) +vi.mock('@sim/db', () => dbChainMock) + vi.mock('@/lib/table/events', () => ({ appendTableEvent: mockAppendTableEvent, })) @@ -77,9 +72,13 @@ const RUNNING_STATE: RowExecutionMetadata = { } describe('writeWorkflowGroupState', () => { + afterAll(() => { + resetDbChainMock() + }) + beforeEach(() => { vi.clearAllMocks() - mockTransaction.mockImplementation(async (callback) => callback({})) + resetDbChainMock() mockWriteExecutionsPatch.mockResolvedValue('wrote') mockUpdateRow.mockResolvedValue({}) mockAppendTableEvent.mockResolvedValue(null) @@ -90,7 +89,7 @@ describe('writeWorkflowGroupState', () => { 'wrote' ) - expect(mockTransaction).toHaveBeenCalledOnce() + expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() expect(mockWriteExecutionsPatch).toHaveBeenCalledWith( expect.anything(), TABLE.id, diff --git a/apps/sim/lib/table/snapshot-cache.test.ts b/apps/sim/lib/table/snapshot-cache.test.ts index 2fe0927c537..35f5b4d1310 100644 --- a/apps/sim/lib/table/snapshot-cache.test.ts +++ b/apps/sim/lib/table/snapshot-cache.test.ts @@ -1,26 +1,18 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const { - mockDb, - mockSelectExportRowPage, - mockCreateMultipartUpload, - mockHeadObject, - mockDeleteFile, -} = vi.hoisted(() => { - const limit = vi.fn() - return { - mockDb: { limit, select: () => ({ from: () => ({ where: () => ({ limit }) }) }) }, +import { dbChainMock, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockSelectExportRowPage, mockCreateMultipartUpload, mockHeadObject, mockDeleteFile } = + vi.hoisted(() => ({ mockSelectExportRowPage: vi.fn(), mockCreateMultipartUpload: vi.fn(), mockHeadObject: vi.fn(), mockDeleteFile: vi.fn(), - } -}) + })) -vi.mock('@sim/db', () => ({ db: mockDb })) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/table/jobs/service', () => ({ selectExportRowPage: mockSelectExportRowPage })) vi.mock('@/lib/uploads/core/storage-service', () => ({ createMultipartUpload: mockCreateMultipartUpload, @@ -45,12 +37,17 @@ let lastHandle: { /** Queue the values successive `readRowsVersion` calls return. */ function versions(...values: number[]) { - for (const v of values) mockDb.limit.mockResolvedValueOnce([{ rowsVersion: v }]) + for (const v of values) queueTableRows(schemaMock.userTableDefinitions, [{ rowsVersion: v }]) } describe('getOrCreateTableSnapshot', () => { + afterAll(() => { + resetDbChainMock() + }) + beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() lastHandle = null mockDeleteFile.mockResolvedValue(undefined) mockSelectExportRowPage.mockResolvedValueOnce([ diff --git a/apps/sim/lib/webhooks/deploy.test.ts b/apps/sim/lib/webhooks/deploy.test.ts index 8409705850d..2ffbf360fe7 100644 --- a/apps/sim/lib/webhooks/deploy.test.ts +++ b/apps/sim/lib/webhooks/deploy.test.ts @@ -2,27 +2,12 @@ * @vitest-environment node */ import { credential } from '@sim/db/schema' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { resetDbChainMock } from '@sim/testing' +import { eq } from 'drizzle-orm' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' import type { SubBlockConfig } from '@/blocks/types' import type { BlockState } from '@/stores/workflows/workflow/types' -const { mockEq, mockFrom, mockLimit, mockSelect, mockWhere } = vi.hoisted(() => ({ - mockEq: vi.fn((left: unknown, right: unknown) => ({ left, right })), - mockFrom: vi.fn(), - mockLimit: vi.fn(), - mockSelect: vi.fn(), - mockWhere: vi.fn(), -})) - -vi.mock('@sim/db', () => ({ db: { select: mockSelect } })) -vi.mock('drizzle-orm', () => ({ - and: vi.fn((...conditions: unknown[]) => ({ conditions })), - eq: mockEq, - inArray: vi.fn((...args: unknown[]) => ({ args })), - isNull: vi.fn((value: unknown) => ({ value })), - or: vi.fn((...conditions: unknown[]) => ({ conditions })), -})) - // deploy.ts pulls in the trigger/block/provider registries at module load; none are exercised by // buildProviderConfig (a pure function), so stub them to keep this unit test fast and isolated. vi.mock('@/blocks', () => ({ getBlock: vi.fn() })) @@ -42,6 +27,8 @@ vi.mock('@/lib/webhooks/pending-verification', () => ({ import { buildProviderConfig, resolveTriggerCredentialId } from '@/lib/webhooks/deploy' +afterAll(resetDbChainMock) + const trigger = (subBlocks: Partial[]): { subBlocks: SubBlockConfig[] } => ({ subBlocks: subBlocks as SubBlockConfig[], }) @@ -96,10 +83,7 @@ function makeBlock( beforeEach(() => { vi.clearAllMocks() - mockSelect.mockReturnValue({ from: mockFrom }) - mockFrom.mockReturnValue({ where: mockWhere }) - mockWhere.mockReturnValue({ limit: mockLimit }) - mockLimit.mockResolvedValue([{ id: 'credential-1' }]) + resetDbChainMock() }) describe('buildProviderConfig canonical collapse', () => { @@ -200,10 +184,10 @@ describe('resolveTriggerCredentialId', () => { it('canonicalizes an OAuth service alias at the credential lookup boundary', async () => { await resolveTriggerCredentialId('credential-1', 'workspace-1', 'gmail') - expect(mockEq).toHaveBeenCalledWith(credential.workspaceId, 'workspace-1') - expect(mockEq).toHaveBeenCalledWith(credential.type, 'oauth') - expect(mockEq).toHaveBeenCalledWith(credential.providerId, 'google-email') - expect(mockEq).toHaveBeenCalledWith(credential.id, 'credential-1') - expect(mockEq).toHaveBeenCalledWith(credential.accountId, 'credential-1') + expect(eq).toHaveBeenCalledWith(credential.workspaceId, 'workspace-1') + expect(eq).toHaveBeenCalledWith(credential.type, 'oauth') + expect(eq).toHaveBeenCalledWith(credential.providerId, 'google-email') + expect(eq).toHaveBeenCalledWith(credential.id, 'credential-1') + expect(eq).toHaveBeenCalledWith(credential.accountId, 'credential-1') }) }) diff --git a/apps/sim/lib/webhooks/polling/utils.test.ts b/apps/sim/lib/webhooks/polling/utils.test.ts index 137c70b3667..48deafe868d 100644 --- a/apps/sim/lib/webhooks/polling/utils.test.ts +++ b/apps/sim/lib/webhooks/polling/utils.test.ts @@ -1,37 +1,14 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const { mockUpdate, mockSet, mockWhere, mockSelect, mockSelectRows, sqlCalls } = vi.hoisted(() => { - const mockSelectRows = vi.fn() - return { - mockUpdate: vi.fn(), - mockSet: vi.fn(), - mockWhere: vi.fn(), - mockSelectRows, - mockSelect: vi.fn(() => ({ - from: vi.fn(() => ({ - where: vi.fn(() => ({ - limit: mockSelectRows, - })), - })), - })), - sqlCalls: [] as Array<{ strings: readonly string[]; values: unknown[] }>, - } -}) +import { account } from '@sim/db/schema' +import { dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' -vi.mock('@sim/db', () => ({ db: { update: mockUpdate, select: mockSelect } })) -vi.mock('@sim/db/schema', () => ({ - webhook: { - id: 'webhook.id', - providerConfig: 'webhook.providerConfig', - updatedAt: 'webhook.updatedAt', - }, - account: {}, - workflow: {}, - workflowDeploymentVersion: {}, +const { sqlCalls } = vi.hoisted(() => ({ + sqlCalls: [] as Array<{ strings: readonly string[]; values: unknown[] }>, })) + vi.mock('drizzle-orm', () => ({ sql: (strings: readonly string[], ...values: unknown[]) => { const node = { strings, values } @@ -59,6 +36,8 @@ import { resolveOAuthAccountId, } from '@/app/api/auth/oauth/utils' +afterAll(resetDbChainMock) + const logger = { error: vi.fn() } as never function allInterpolatedValues(): unknown[] { @@ -72,10 +51,8 @@ function allSqlText(): string { describe('updateWebhookProviderConfig (atomic jsonb merge)', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() sqlCalls.length = 0 - mockWhere.mockResolvedValue(undefined) - mockSet.mockReturnValue({ where: mockWhere }) - mockUpdate.mockReturnValue({ set: mockSet }) }) it('merges defined keys (null preserved) and removes undefined keys', async () => { @@ -85,7 +62,7 @@ describe('updateWebhookProviderConfig (atomic jsonb merge)', () => { logger ) - expect(mockUpdate).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.update).toHaveBeenCalledTimes(1) expect(allInterpolatedValues()).toContain(JSON.stringify({ historyId: 'h1', nulled: null })) expect(allInterpolatedValues()).toContainEqual(['cleared']) }) @@ -114,14 +91,14 @@ describe('resolveOAuthCredential (single-credential polling)', () => { beforeEach(() => { vi.clearAllMocks() - mockSelectRows.mockResolvedValue([]) + resetDbChainMock() }) it('resolves via credentialId: account lookup then token refresh', async () => { vi.mocked(resolveOAuthAccountId).mockResolvedValue({ accountId: 'acc-1', } as Awaited>) - mockSelectRows.mockResolvedValue([{ userId: 'owner-1' }]) + queueTableRows(account, [{ userId: 'owner-1' }]) vi.mocked(refreshAccessTokenIfNeeded).mockResolvedValue('tok-abc') const token = await resolveOAuthCredential( @@ -148,7 +125,6 @@ describe('resolveOAuthCredential (single-credential polling)', () => { vi.mocked(resolveOAuthAccountId).mockResolvedValue({ accountId: 'acc-missing', } as Awaited>) - mockSelectRows.mockResolvedValue([]) await expect( resolveOAuthCredential(makeWebhook({ credentialId: 'cred-1' }), 'google-email', 'req-1') diff --git a/apps/sim/lib/webhooks/processor.test.ts b/apps/sim/lib/webhooks/processor.test.ts index 92f0e9c6927..21d5c00eee0 100644 --- a/apps/sim/lib/webhooks/processor.test.ts +++ b/apps/sim/lib/webhooks/processor.test.ts @@ -5,14 +5,18 @@ import type { webhook, workflow } from '@sim/db/schema' import { createMockRequest, + dbChainMock, envFlagsMock, executionPreprocessingMock, executionPreprocessingMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, workflowsPersistenceUtilsMock, workflowsPersistenceUtilsMockFns, } from '@sim/testing' import type { NextRequest } from 'next/server' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' import { ADMISSION_ERROR_CODE, ADMISSION_RETRY_AFTER_SECONDS, @@ -33,7 +37,6 @@ const { mockReleaseExecutionSlot, mockProviderHandler, mockShouldExecuteInline, - mockWebhookLookupResult, } = vi.hoisted(() => ({ mockGenerateId: vi.fn(), mockAdmissionRelease: vi.fn(), @@ -42,39 +45,11 @@ const { mockReleaseExecutionSlot: vi.fn(), mockProviderHandler: { current: {} as Record }, mockShouldExecuteInline: vi.fn(), - mockWebhookLookupResult: { - rows: [] as WebhookLookupRow[], - claim: [] as Array<{ workflowId: string }>, - }, })) const mockPreprocessExecution = executionPreprocessingMockFns.mockPreprocessExecution -vi.mock('@sim/db', () => { - const selectChain = { - from: () => selectChain, - innerJoin: () => selectChain, - leftJoin: () => selectChain, - where: () => ({ - then: (resolve: (rows: WebhookLookupRow[]) => void) => resolve(mockWebhookLookupResult.rows), - limit: () => Promise.resolve(mockWebhookLookupResult.claim), - }), - } - return { - db: { select: () => selectChain }, - webhook: {}, - webhookPathClaim: {}, - workflow: {}, - workflowDeploymentVersion: {}, - } -}) - -vi.mock('drizzle-orm', () => ({ - and: vi.fn(), - eq: vi.fn(), - isNull: vi.fn(), - or: vi.fn(), -})) +vi.mock('@sim/db', () => ({ ...dbChainMock, ...schemaMock })) vi.mock('@sim/utils/id', () => ({ generateId: mockGenerateId, @@ -167,6 +142,8 @@ import { processPolledWebhookEvent, } from '@/lib/webhooks/processor' +afterAll(resetDbChainMock) + function makeWebhookRecord(overrides: Partial): WebhookRecord { const now = new Date('2026-01-01T00:00:00.000Z') return { @@ -228,8 +205,7 @@ const billingAttribution = { describe('findAllWebhooksForPath cross-tenant collision', () => { beforeEach(() => { vi.clearAllMocks() - mockWebhookLookupResult.rows = [] - mockWebhookLookupResult.claim = [] + resetDbChainMock() }) const makeRow = (workflowId: string, webhookId: string, createdAt: Date) => ({ @@ -237,11 +213,16 @@ describe('findAllWebhooksForPath cross-tenant collision', () => { workflow: { id: workflowId }, }) + const queueLookup = (rows: WebhookLookupRow[], claim: Array<{ workflowId: string }> = []) => { + queueTableRows(schemaMock.webhook, rows) + queueTableRows(schemaMock.webhookPathClaim, claim) + } + it('returns all rows when they belong to a single workflow', async () => { - mockWebhookLookupResult.rows = [ + queueLookup([ makeRow('workflow-1', 'wh-a', new Date('2026-01-01')), makeRow('workflow-1', 'wh-b', new Date('2026-01-02')), - ] + ]) const results = await findAllWebhooksForPath({ requestId: 'req-1', path: 'shared-path' }) @@ -252,7 +233,7 @@ describe('findAllWebhooksForPath cross-tenant collision', () => { it('drops foreign rows when a path collides across workflows, keeping the earliest owner', async () => { const victim = makeRow('victim-workflow', 'victim-wh', new Date('2026-01-01')) const attacker = makeRow('attacker-workflow', 'attacker-wh', new Date('2026-05-01')) - mockWebhookLookupResult.rows = [attacker, victim] + queueLookup([attacker, victim]) const results = await findAllWebhooksForPath({ requestId: 'req-2', path: 'shared-path' }) @@ -264,8 +245,7 @@ describe('findAllWebhooksForPath cross-tenant collision', () => { it('prefers the path-claim owner over an earlier-created interloper', async () => { const interloper = makeRow('interloper-workflow', 'interloper-wh', new Date('2026-01-01')) const claimHolder = makeRow('claim-workflow', 'claim-wh', new Date('2026-05-01')) - mockWebhookLookupResult.rows = [interloper, claimHolder] - mockWebhookLookupResult.claim = [{ workflowId: 'claim-workflow' }] + queueLookup([interloper, claimHolder], [{ workflowId: 'claim-workflow' }]) const results = await findAllWebhooksForPath({ requestId: 'req-6', path: 'shared-path' }) @@ -276,8 +256,7 @@ describe('findAllWebhooksForPath cross-tenant collision', () => { it('falls back to earliest registration when the claim owner has no deliverable rows', async () => { const victim = makeRow('victim-workflow', 'victim-wh', new Date('2026-01-01')) const attacker = makeRow('attacker-workflow', 'attacker-wh', new Date('2026-05-01')) - mockWebhookLookupResult.rows = [attacker, victim] - mockWebhookLookupResult.claim = [{ workflowId: 'absent-workflow' }] + queueLookup([attacker, victim], [{ workflowId: 'absent-workflow' }]) const results = await findAllWebhooksForPath({ requestId: 'req-7', path: 'shared-path' }) @@ -289,7 +268,7 @@ describe('findAllWebhooksForPath cross-tenant collision', () => { const victimA = makeRow('victim-workflow', 'victim-wh-a', new Date('2026-01-01')) const victimB = makeRow('victim-workflow', 'victim-wh-b', new Date('2026-01-03')) const attacker = makeRow('attacker-workflow', 'attacker-wh', new Date('2026-05-01')) - mockWebhookLookupResult.rows = [victimB, attacker, victimA] + queueLookup([victimB, attacker, victimA]) const results = await findAllWebhooksForPath({ requestId: 'req-5', path: 'shared-path' }) @@ -299,8 +278,6 @@ describe('findAllWebhooksForPath cross-tenant collision', () => { }) it('returns an empty array when no webhooks match', async () => { - mockWebhookLookupResult.rows = [] - const results = await findAllWebhooksForPath({ requestId: 'req-3', path: 'missing' }) expect(results).toEqual([]) diff --git a/apps/sim/lib/webhooks/providers/whatsapp.test.ts b/apps/sim/lib/webhooks/providers/whatsapp.test.ts index a0ff2b8a9ba..e356365c051 100644 --- a/apps/sim/lib/webhooks/providers/whatsapp.test.ts +++ b/apps/sim/lib/webhooks/providers/whatsapp.test.ts @@ -2,13 +2,11 @@ * @vitest-environment node */ import { createHmac } from 'node:crypto' +import { dbChainMock, schemaMock } from '@sim/testing' import { NextRequest } from 'next/server' import { describe, expect, it, vi } from 'vitest' -vi.mock('@sim/db', () => ({ - db: {}, - workflowDeploymentVersion: {}, -})) +vi.mock('@sim/db', () => ({ ...dbChainMock, ...schemaMock })) import { whatsappHandler } from './whatsapp' diff --git a/apps/sim/lib/webhooks/registration-store.test.ts b/apps/sim/lib/webhooks/registration-store.test.ts index bcb54a43e79..5a0641f3304 100644 --- a/apps/sim/lib/webhooks/registration-store.test.ts +++ b/apps/sim/lib/webhooks/registration-store.test.ts @@ -1,7 +1,8 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { dbChainMockFns, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' interface Condition { kind: string @@ -10,16 +11,9 @@ interface Condition { conditions?: Condition[] } -const { mockTransaction, mockIsDeploymentOperationCurrent, mockClaimWebhookPath } = vi.hoisted( - () => ({ - mockTransaction: vi.fn(), - mockIsDeploymentOperationCurrent: vi.fn(), - mockClaimWebhookPath: vi.fn(), - }) -) - -vi.mock('@sim/db', () => ({ - db: { transaction: mockTransaction }, +const { mockIsDeploymentOperationCurrent, mockClaimWebhookPath } = vi.hoisted(() => ({ + mockIsDeploymentOperationCurrent: vi.fn(), + mockClaimWebhookPath: vi.fn(), })) vi.mock('drizzle-orm', () => ({ @@ -60,6 +54,8 @@ import { type WebhookRegistrationOperationFence, } from '@/lib/webhooks/registration-store' +afterAll(resetDbChainMock) + const FENCE: WebhookRegistrationOperationFence = { workflowId: 'workflow-1', operationId: 'operation-1', @@ -154,6 +150,7 @@ function activeRow(overrides: Record = {}) { describe('activateWebhookRegistrations', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockIsDeploymentOperationCurrent.mockResolvedValue(true) }) @@ -219,17 +216,18 @@ describe('activateWebhookRegistrations', () => { describe('prepareWebhookRegistrationIntents', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockIsDeploymentOperationCurrent.mockResolvedValue(true) mockClaimWebhookPath.mockResolvedValue('hooks/a') - mockTransaction.mockImplementation(async (callback: (tx: DbOrTx) => Promise) => { - throw new Error('mockTransaction not configured for this test') + dbChainMockFns.transaction.mockImplementation(async () => { + throw new Error('db.transaction not configured for this test') }) }) function runInTx(selectResults: unknown[][]) { const harness = createTx(selectResults) - mockTransaction.mockImplementation(async (callback: (tx: DbOrTx) => Promise) => - callback(harness.tx) + dbChainMockFns.transaction.mockImplementation( + async (callback: (tx: DbOrTx) => Promise) => callback(harness.tx) ) return harness } diff --git a/apps/sim/lib/webhooks/utils.server.test.ts b/apps/sim/lib/webhooks/utils.server.test.ts index 6668984a600..15ddf9e3d6b 100644 --- a/apps/sim/lib/webhooks/utils.server.test.ts +++ b/apps/sim/lib/webhooks/utils.server.test.ts @@ -1,60 +1,21 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' - -interface Condition { - kind: string - column?: unknown - value?: unknown - conditions?: Condition[] -} - -const { mockSelect } = vi.hoisted(() => ({ mockSelect: vi.fn() })) - -vi.mock('@sim/db', () => ({ db: { select: mockSelect } })) - -vi.mock('drizzle-orm', () => ({ - and: (...conditions: Condition[]) => ({ kind: 'and', conditions }), - eq: (column: unknown, value: unknown) => ({ kind: 'eq', column, value }), - isNull: (column: unknown) => ({ kind: 'isNull', column }), -})) - +import { webhook, webhookPathClaim } from '@sim/db/schema' +import { dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' import { findConflictingWebhookPathOwner } from '@/lib/webhooks/utils.server' -function claimLookupChain(rows: unknown[], captureCondition?: (condition: Condition) => void) { - return { - from: vi.fn(() => ({ - where: vi.fn((condition: Condition) => { - captureCondition?.(condition) - return { limit: vi.fn().mockResolvedValue(rows) } - }), - })), - } -} - -function liveRowsChain(rows: unknown[]) { - return { - from: vi.fn(() => ({ - innerJoin: vi.fn(() => ({ - where: vi.fn().mockResolvedValue(rows), - })), - })), - } -} +afterAll(resetDbChainMock) describe('findConflictingWebhookPathOwner', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() }) it('returns the claim owner while the claim holder is mid-rotation', async () => { - let claimCondition: Condition | undefined - mockSelect.mockReturnValueOnce( - claimLookupChain([{ workflowId: 'workflow-owner' }], (condition) => { - claimCondition = condition - }) - ) + queueTableRows(webhookPathClaim, [{ workflowId: 'workflow-owner' }]) const owner = await findConflictingWebhookPathOwner({ path: ' /leads/ ', @@ -62,14 +23,15 @@ describe('findConflictingWebhookPathOwner', () => { }) expect(owner).toBe('workflow-owner') - expect(mockSelect).toHaveBeenCalledTimes(1) - expect(claimCondition).toEqual(expect.objectContaining({ kind: 'eq', value: 'leads' })) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.where).toHaveBeenCalledWith( + expect.objectContaining({ type: 'eq', right: 'leads' }) + ) }) it('ignores the caller-owned claim and falls through to live rows', async () => { - mockSelect - .mockReturnValueOnce(claimLookupChain([{ workflowId: 'workflow-caller' }])) - .mockReturnValueOnce(liveRowsChain([])) + queueTableRows(webhookPathClaim, [{ workflowId: 'workflow-caller' }]) + queueTableRows(webhook, []) const owner = await findConflictingWebhookPathOwner({ path: 'leads', @@ -77,15 +39,12 @@ describe('findConflictingWebhookPathOwner', () => { }) expect(owner).toBeNull() - expect(mockSelect).toHaveBeenCalledTimes(2) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(2) }) it('returns a foreign live-row owner when no claim exists', async () => { - mockSelect - .mockReturnValueOnce(claimLookupChain([])) - .mockReturnValueOnce( - liveRowsChain([{ workflowId: 'workflow-caller' }, { workflowId: 'workflow-foreign' }]) - ) + queueTableRows(webhookPathClaim, []) + queueTableRows(webhook, [{ workflowId: 'workflow-caller' }, { workflowId: 'workflow-foreign' }]) const owner = await findConflictingWebhookPathOwner({ path: 'leads', @@ -96,7 +55,7 @@ describe('findConflictingWebhookPathOwner', () => { }) it('skips the claim lookup entirely for empty paths', async () => { - mockSelect.mockReturnValueOnce(liveRowsChain([])) + queueTableRows(webhook, []) const owner = await findConflictingWebhookPathOwner({ path: ' ', @@ -104,6 +63,6 @@ describe('findConflictingWebhookPathOwner', () => { }) expect(owner).toBeNull() - expect(mockSelect).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(1) }) }) diff --git a/apps/sim/lib/workspaces/lifecycle.test.ts b/apps/sim/lib/workspaces/lifecycle.test.ts index 1013b2280f4..671e6c73df6 100644 --- a/apps/sim/lib/workspaces/lifecycle.test.ts +++ b/apps/sim/lib/workspaces/lifecycle.test.ts @@ -1,23 +1,24 @@ /** * @vitest-environment node */ -import { permissionsMock, permissionsMockFns } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + dbChainMock, + dbChainMockFns, + permissionsMock, + permissionsMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, +} from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' -const { mockSelect, mockTransaction, mockArchiveWorkflowsForWorkspace } = vi.hoisted(() => ({ - mockSelect: vi.fn(), - mockTransaction: vi.fn(), +const { mockArchiveWorkflowsForWorkspace } = vi.hoisted(() => ({ mockArchiveWorkflowsForWorkspace: vi.fn(), })) const mockGetWorkspaceWithOwner = permissionsMockFns.mockGetWorkspaceWithOwner -vi.mock('@sim/db', () => ({ - db: { - select: mockSelect, - transaction: mockTransaction, - }, -})) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/workflows/lifecycle', () => ({ archiveWorkflowsForWorkspace: (...args: unknown[]) => mockArchiveWorkflowsForWorkspace(...args), @@ -38,6 +39,11 @@ function createUpdateChain() { describe('workspace lifecycle', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() + }) + + afterAll(() => { + resetDbChainMock() }) it('archives workspace and dependent resources', async () => { @@ -48,11 +54,7 @@ describe('workspace lifecycle', () => { archivedAt: null, }) mockArchiveWorkflowsForWorkspace.mockResolvedValue(2) - mockSelect.mockReturnValue({ - from: vi.fn().mockReturnValue({ - where: vi.fn().mockResolvedValue([{ id: 'server-1' }]), - }), - }) + queueTableRows(schemaMock.workflowMcpServer, [{ id: 'server-1' }]) const tx = { select: vi.fn().mockReturnValue({ @@ -65,8 +67,8 @@ describe('workspace lifecycle', () => { where: vi.fn().mockResolvedValue([]), })), } - mockTransaction.mockImplementation(async (callback: (trx: typeof tx) => Promise) => - callback(tx) + dbChainMockFns.transaction.mockImplementation( + async (callback: (trx: typeof tx) => Promise) => callback(tx) ) const result = await archiveWorkspace('workspace-1', { requestId: 'req-1' }) @@ -99,6 +101,6 @@ describe('workspace lifecycle', () => { expect(mockArchiveWorkflowsForWorkspace).toHaveBeenCalledWith('workspace-1', { requestId: 'req-1', }) - expect(mockTransaction).not.toHaveBeenCalled() + expect(dbChainMockFns.transaction).not.toHaveBeenCalled() }) }) diff --git a/apps/sim/lib/workspaces/organization-workspaces.test.ts b/apps/sim/lib/workspaces/organization-workspaces.test.ts index d9f2804624d..d84c3571cf7 100644 --- a/apps/sim/lib/workspaces/organization-workspaces.test.ts +++ b/apps/sim/lib/workspaces/organization-workspaces.test.ts @@ -1,101 +1,32 @@ /** * @vitest-environment node */ -import { schemaMock } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + dbChainMock, + dbChainMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, +} from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { - mockDbResults, - mockUpdateWhere, - mockUpdateReturning, - mockUpdateSet, - mockDbUpdate, - mockOnConflictDoUpdate, - mockInsertValues, - mockDbInsert, mockEnsureUserInOrganizationTx, mockSyncUsageLimitsFromSubscription, mockReapplyPaidOrgJoinBillingForExistingMemberTx, mockAcquireOrganizationMutationLock, mockAcquireInvitationMutationLocks, mockChangeWorkspaceStoragePayersInTx, - mockSelectForUpdate, -} = vi.hoisted(() => { - const mockDbResults: { value: any[] } = { value: [] } - const mockUpdateReturning = vi.fn() - const mockUpdateWhere = vi.fn().mockReturnValue({ returning: mockUpdateReturning }) - const mockUpdateSet = vi.fn().mockReturnValue({ where: mockUpdateWhere }) - const mockDbUpdate = vi.fn().mockReturnValue({ set: mockUpdateSet }) - const mockOnConflictDoUpdate = vi.fn().mockResolvedValue(undefined) - const mockInsertValues = vi.fn().mockReturnValue({ - onConflictDoUpdate: mockOnConflictDoUpdate, - }) - const mockDbInsert = vi.fn().mockReturnValue({ values: mockInsertValues }) - const mockEnsureUserInOrganizationTx = vi.fn() - const mockSyncUsageLimitsFromSubscription = vi.fn().mockResolvedValue(undefined) - const mockReapplyPaidOrgJoinBillingForExistingMemberTx = vi.fn().mockResolvedValue({ - proUsageSnapshotted: false, - proCancelledAtPeriodEnd: false, - }) - const mockAcquireOrganizationMutationLock = vi.fn() - const mockAcquireInvitationMutationLocks = vi.fn() - const mockChangeWorkspaceStoragePayersInTx = vi.fn() - const mockSelectForUpdate = vi.fn() - - return { - mockDbResults, - mockUpdateWhere, - mockUpdateReturning, - mockUpdateSet, - mockDbUpdate, - mockOnConflictDoUpdate, - mockInsertValues, - mockDbInsert, - mockEnsureUserInOrganizationTx, - mockSyncUsageLimitsFromSubscription, - mockReapplyPaidOrgJoinBillingForExistingMemberTx, - mockAcquireOrganizationMutationLock, - mockAcquireInvitationMutationLocks, - mockChangeWorkspaceStoragePayersInTx, - mockSelectForUpdate, - } -}) - -vi.mock('@sim/db', () => { - const selectImpl = vi.fn().mockImplementation(() => { - const chain: any = {} - chain.from = vi.fn().mockReturnValue(chain) - chain.where = vi.fn().mockReturnValue(chain) - chain.orderBy = vi.fn().mockReturnValue(chain) - chain.for = vi.fn().mockImplementation(() => { - mockSelectForUpdate() - return chain - }) - chain.limit = vi - .fn() - .mockImplementation(() => Promise.resolve(mockDbResults.value.shift() || [])) - chain.then = vi.fn().mockImplementation((callback: (rows: any[]) => unknown) => { - const result = mockDbResults.value.shift() || [] - return Promise.resolve(callback ? callback(result) : result) - }) - return chain - }) - const txObject = { - select: selectImpl, - update: mockDbUpdate, - insert: mockDbInsert, - } - return { - db: { - select: selectImpl, - update: mockDbUpdate, - insert: mockDbInsert, - transaction: vi.fn(async (fn: (tx: typeof txObject) => unknown) => fn(txObject)), - }, - } -}) +} = vi.hoisted(() => ({ + mockEnsureUserInOrganizationTx: vi.fn(), + mockSyncUsageLimitsFromSubscription: vi.fn(), + mockReapplyPaidOrgJoinBillingForExistingMemberTx: vi.fn(), + mockAcquireOrganizationMutationLock: vi.fn(), + mockAcquireInvitationMutationLocks: vi.fn(), + mockChangeWorkspaceStoragePayersInTx: vi.fn(), +})) -vi.mock('@sim/db/schema', () => schemaMock) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/billing/organizations/membership', () => ({ acquireOrganizationMutationLock: mockAcquireOrganizationMutationLock, @@ -129,26 +60,31 @@ import { describe('organization workspace helpers', () => { beforeEach(() => { vi.clearAllMocks() - mockDbResults.value = [] + resetDbChainMock() mockEnsureUserInOrganizationTx.mockReset() - mockUpdateReturning.mockReset() mockChangeWorkspaceStoragePayersInTx.mockReset() mockSyncUsageLimitsFromSubscription.mockResolvedValue(undefined) + mockReapplyPaidOrgJoinBillingForExistingMemberTx.mockResolvedValue({ + proUsageSnapshotted: false, + proCancelledAtPeriodEnd: false, + }) + }) + + afterAll(() => { + resetDbChainMock() }) it('attaches owned workspaces to an organization and syncs existing members', async () => { - mockDbResults.value = [ - [{ id: 'ws-1' }, { id: 'ws-2' }], - [{ id: 'ws-1' }, { id: 'ws-2' }], - [ - { id: 'ws-1', billedAccountUserId: 'user-1', organizationId: null }, - { id: 'ws-2', billedAccountUserId: 'user-1', organizationId: null }, - ], - [{ userId: 'owner-1' }], - [{ userId: 'owner-1' }, { userId: 'member-1' }], - [{ userId: 'owner-1', organizationId: 'org-1' }], - ] - mockUpdateReturning.mockResolvedValueOnce([{ id: 'ws-2' }, { id: 'ws-1' }]) + queueTableRows(schemaMock.workspace, [{ id: 'ws-1' }, { id: 'ws-2' }]) + queueTableRows(schemaMock.workspace, [{ id: 'ws-1' }, { id: 'ws-2' }]) + queueTableRows(schemaMock.workspace, [ + { id: 'ws-1', billedAccountUserId: 'user-1', organizationId: null }, + { id: 'ws-2', billedAccountUserId: 'user-1', organizationId: null }, + ]) + queueTableRows(schemaMock.member, [{ userId: 'owner-1' }]) + queueTableRows(schemaMock.permissions, [{ userId: 'owner-1' }, { userId: 'member-1' }]) + queueTableRows(schemaMock.member, [{ userId: 'owner-1', organizationId: 'org-1' }]) + dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'ws-2' }, { id: 'ws-1' }]) mockEnsureUserInOrganizationTx .mockResolvedValueOnce({ success: true, @@ -194,11 +130,11 @@ describe('organization workspace helpers', () => { 'owner-1', 'org-1' ) - expect(mockUpdateSet).toHaveBeenCalledWith( + expect(dbChainMockFns.set).toHaveBeenCalledWith( expect.objectContaining({ organizationAssignedAt: expect.any(Date) }) ) expect(mockChangeWorkspaceStoragePayersInTx).toHaveBeenCalledTimes(1) - expect(mockSelectForUpdate.mock.invocationCallOrder[0]).toBeLessThan( + expect(dbChainMockFns.for.mock.invocationCallOrder[0]).toBeLessThan( mockEnsureUserInOrganizationTx.mock.invocationCallOrder[0] ) expect(mockChangeWorkspaceStoragePayersInTx).toHaveBeenCalledWith(expect.anything(), [ @@ -221,23 +157,23 @@ describe('organization workspace helpers', () => { }, }, ]) - expect(mockDbUpdate).toHaveBeenCalledTimes(1) - expect(mockDbInsert).toHaveBeenCalledTimes(1) - expect(mockInsertValues).toHaveBeenCalledWith([ + expect(dbChainMockFns.update).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.insert).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.values).toHaveBeenCalledWith([ expect.objectContaining({ entityId: 'ws-1', userId: 'owner-1' }), expect.objectContaining({ entityId: 'ws-2', userId: 'owner-1' }), ]) }) it('fails before attaching workspaces when an existing member belongs to another organization', async () => { - mockDbResults.value = [ - [{ id: 'ws-1' }], - [{ id: 'ws-1' }], - [{ id: 'ws-1', billedAccountUserId: 'user-1', organizationId: null }], - [{ userId: 'owner-1' }], - [{ userId: 'owner-1' }, { userId: 'member-2' }], - [{ userId: 'member-2', organizationId: 'org-2' }], - ] + queueTableRows(schemaMock.workspace, [{ id: 'ws-1' }]) + queueTableRows(schemaMock.workspace, [{ id: 'ws-1' }]) + queueTableRows(schemaMock.workspace, [ + { id: 'ws-1', billedAccountUserId: 'user-1', organizationId: null }, + ]) + queueTableRows(schemaMock.member, [{ userId: 'owner-1' }]) + queueTableRows(schemaMock.permissions, [{ userId: 'owner-1' }, { userId: 'member-2' }]) + queueTableRows(schemaMock.member, [{ userId: 'member-2', organizationId: 'org-2' }]) await expect( attachOwnedWorkspacesToOrganization({ @@ -247,19 +183,19 @@ describe('organization workspace helpers', () => { ).rejects.toBeInstanceOf(WorkspaceOrganizationMembershipConflictError) expect(mockEnsureUserInOrganizationTx).not.toHaveBeenCalled() - expect(mockDbUpdate).not.toHaveBeenCalled() + expect(dbChainMockFns.update).not.toHaveBeenCalled() }) it('keeps cross-org members external and still attaches when policy is keep-external', async () => { - mockDbResults.value = [ - [{ id: 'ws-1' }], - [{ id: 'ws-1' }], - [{ id: 'ws-1', billedAccountUserId: 'user-1', organizationId: null }], - [{ userId: 'owner-1' }], - [{ userId: 'owner-1' }, { userId: 'member-2' }], - [{ userId: 'member-2', organizationId: 'org-2' }], - ] - mockUpdateReturning.mockResolvedValueOnce([{ id: 'ws-1' }]) + queueTableRows(schemaMock.workspace, [{ id: 'ws-1' }]) + queueTableRows(schemaMock.workspace, [{ id: 'ws-1' }]) + queueTableRows(schemaMock.workspace, [ + { id: 'ws-1', billedAccountUserId: 'user-1', organizationId: null }, + ]) + queueTableRows(schemaMock.member, [{ userId: 'owner-1' }]) + queueTableRows(schemaMock.permissions, [{ userId: 'owner-1' }, { userId: 'member-2' }]) + queueTableRows(schemaMock.member, [{ userId: 'member-2', organizationId: 'org-2' }]) + dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'ws-1' }]) mockEnsureUserInOrganizationTx.mockResolvedValueOnce({ success: true, alreadyMember: true, @@ -287,11 +223,13 @@ describe('organization workspace helpers', () => { expect.anything(), expect.objectContaining({ userId: 'owner-1' }) ) - expect(mockDbUpdate).toHaveBeenCalled() + expect(dbChainMockFns.update).toHaveBeenCalled() }) it('rolls back membership work when a concurrent move wins before the locked re-read', async () => { - mockDbResults.value = [[{ id: 'ws-1' }], [{ id: 'ws-1' }], []] + queueTableRows(schemaMock.workspace, [{ id: 'ws-1' }]) + queueTableRows(schemaMock.workspace, [{ id: 'ws-1' }]) + queueTableRows(schemaMock.workspace, []) const result = await attachOwnedWorkspacesToOrganization({ ownerUserId: 'user-1', @@ -309,19 +247,19 @@ describe('organization workspace helpers', () => { workspaceIds: ['ws-1'], }) expect(mockEnsureUserInOrganizationTx).not.toHaveBeenCalled() - expect(mockDbUpdate).not.toHaveBeenCalled() - expect(mockDbInsert).not.toHaveBeenCalled() + expect(dbChainMockFns.update).not.toHaveBeenCalled() + expect(dbChainMockFns.insert).not.toHaveBeenCalled() }) it('does not report a committed attachment as failed when derived usage refresh fails', async () => { - mockDbResults.value = [ - [{ id: 'ws-1' }], - [{ id: 'ws-1' }], - [{ id: 'ws-1', billedAccountUserId: 'user-1', organizationId: null }], - [{ userId: 'owner-1' }], - [{ userId: 'member-1' }], - [], - ] + queueTableRows(schemaMock.workspace, [{ id: 'ws-1' }]) + queueTableRows(schemaMock.workspace, [{ id: 'ws-1' }]) + queueTableRows(schemaMock.workspace, [ + { id: 'ws-1', billedAccountUserId: 'user-1', organizationId: null }, + ]) + queueTableRows(schemaMock.member, [{ userId: 'owner-1' }]) + queueTableRows(schemaMock.permissions, [{ userId: 'member-1' }]) + queueTableRows(schemaMock.member, []) mockEnsureUserInOrganizationTx.mockResolvedValueOnce({ success: true, alreadyMember: false, @@ -331,7 +269,7 @@ describe('organization workspace helpers', () => { proCancelledAtPeriodEnd: false, }, }) - mockUpdateReturning.mockResolvedValueOnce([{ id: 'ws-1' }]) + dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'ws-1' }]) mockSyncUsageLimitsFromSubscription.mockRejectedValueOnce(new Error('refresh failed')) await expect( @@ -343,18 +281,18 @@ describe('organization workspace helpers', () => { }) it('detaches organization workspaces into grandfathered shared mode', async () => { - mockDbResults.value = [ - [{ userId: 'owner-1' }], - [{ id: 'ws-1', ownerId: 'creator-1', billedAccountUserId: 'old-owner' }], - [{ id: 'ws-1' }], - ] + queueTableRows(schemaMock.member, [{ userId: 'owner-1' }]) + queueTableRows(schemaMock.workspace, [ + { id: 'ws-1', ownerId: 'creator-1', billedAccountUserId: 'old-owner' }, + ]) + queueTableRows(schemaMock.workspace, [{ id: 'ws-1' }]) const result = await detachOrganizationWorkspaces('org-1') expect(result.detachedWorkspaceIds).toEqual(['ws-1']) expect(result.billedAccountUserId).toBe('owner-1') expect(mockChangeWorkspaceStoragePayersInTx).toHaveBeenCalledTimes(1) - expect(mockSelectForUpdate.mock.invocationCallOrder[0]).toBeLessThan( + expect(dbChainMockFns.for.mock.invocationCallOrder[0]).toBeLessThan( mockChangeWorkspaceStoragePayersInTx.mock.invocationCallOrder[0] ) expect(mockChangeWorkspaceStoragePayersInTx).toHaveBeenCalledWith(expect.anything(), [ @@ -368,17 +306,17 @@ describe('organization workspace helpers', () => { }, }, ]) - expect(mockUpdateSet).toHaveBeenCalledWith( + expect(dbChainMockFns.set).toHaveBeenCalledWith( expect.objectContaining({ workspaceMode: 'grandfathered_shared', organizationAssignedAt: null, }) ) - expect(mockDbUpdate).toHaveBeenCalledTimes(1) - expect(mockDbInsert).toHaveBeenCalledTimes(1) - expect(mockInsertValues).toHaveBeenCalledWith([ + expect(dbChainMockFns.update).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.insert).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.values).toHaveBeenCalledWith([ expect.objectContaining({ entityId: 'ws-1', userId: 'owner-1' }), ]) - expect(mockOnConflictDoUpdate).toHaveBeenCalled() + expect(dbChainMockFns.onConflictDoUpdate).toHaveBeenCalled() }) })