improvement(tests): db-mock migration tranche 4 — billing, webhooks/execution/logs, routes/misc (final) (#5866)

* improvement(tests): db-mock migration tranche 4 — billing, webhooks/execution/logs, routes/misc (final)

* fix(tests): route agent-handler MCP server rows through queueTableRows

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<Record<string, unknown>>) {
+  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<Record<string, unknown>>) {
+  queueTableRows(schemaMock.ssoProvider, rows)
+  queueTableRows(schemaMock.ssoProvider, rows)
+}
 
 vi.mock('@/lib/auth', () => ({
   getSession: mockGetSession,
@@ -109,8 +88,7 @@ function request(body: Record<string, unknown>) {
 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) = <normalized 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' }

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<string, unknown>) => 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<string, unknown>) => {
+    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',

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<string, unknown>) => {
+    resetDbChainMock()
+    dbChainMockFns.select.mockImplementation((selection: Record<string, unknown>) => {
       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()

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<Record<string, unknown>> },
   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<string, unknown> = {}
-      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' }))
 

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'
+import { afterAll, 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(),
-  }
-})
+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<string, unknown>) => 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')
 

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'
+import {
+  dbChainMock,
+  dbChainMockFns,
+  queueTableRows,
+  resetDbChainMock,
+  schemaMock,
+} from '@sim/testing'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
-const {
-  mockIsOrganizationBillingBlocked,
-  mockMemberWhere,
-  mockMembersWhere,
-  mockSelect,
-  mockSubscriptionWhere,
-} = vi.hoisted(() => ({
+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')

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)

apps/sim/executor/handlers/agent/skills-resolver.test.ts

@@ -1,22 +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 { 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(() => ({})),
-}))
+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()
   })
 })

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)

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)

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

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

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('<html></html>')),
@@ -25,26 +26,7 @@ const {
   isOrgAdminRoleMock: vi.fn(() => true),
 }))
 
-vi.mock('@sim/db', () => {
-  const updateBuilder: Record<string, unknown> = {
-    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<string, unknown> = {
-    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()
   })
 })

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<string, unknown> = {}) {
   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',

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<string, unknown> = {}
-      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)
 

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<Record<string, unknown>> },
-  globalTransactionMock: vi.fn(),
 }))
 
-vi.mock('@sim/db', () => {
-  const update = () => ({
-    set: (values: Record<string, unknown>) => {
-      updateCalls.value.push(values)
-      return { where: () => Promise.resolve([]) }
-    },
-  })
-  const txMock = { update }
-  globalTransactionMock.mockImplementation(async (cb: (tx: typeof txMock) => Promise<unknown>) =>
-    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(

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<string, unknown> = {}
-    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()
 

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'
+import {
+  auditMock,
+  dbChainMock,
+  dbChainMockFns,
+  queueTableRows,
+  resetDbChainMock,
+  schemaMock,
+} from '@sim/testing'
+import { afterAll, 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 },
-  }))
+const { mockSyncSubscriptionUsageLimits, enqueueMock, mockFeatureFlags } = vi.hoisted(() => ({
+  mockSyncSubscriptionUsageLimits: vi.fn(),
+  enqueueMock: vi.fn(),
+  mockFeatureFlags: { isBillingEnabled: true },
+}))
 
-vi.mock('@sim/db', () => {
-  const makeSelectChain = () => {
-    const chain: Record<string, unknown> = {}
-    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<string, unknown>) => {
-      setMock(values)
-      return { where: () => Promise.resolve([]) }
-    },
-  })
-  const txMock = { select: () => makeSelectChain(), update }
-  const dbMock = {
-    select: () => makeSelectChain(),
-    update,
-    transaction: async (cb: (tx: typeof txMock) => Promise<unknown>) => cb(txMock),
-  }
-  return { db: dbMock }
-})
+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',

apps/sim/lib/billing/storage/limits.test.ts

@@ -1,31 +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 {
-  mockEq,
-  mockFlags,
-  mockFrom,
-  mockGetHighestPrioritySubscription,
-  mockLimit,
-  mockSelect,
-  mockWhere,
-} = vi.hoisted(() => ({
+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)
   })

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<T>(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<string, unknown> = 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<string, unknown> = {}) {
   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<string, unknown>) {
+  queueTableRows(schemaMock.userStats, [row])
 }
 
-function buildPersonalSnapshotSelectChain({
-  currentPeriodCost = '0',
-  proPeriodCostSnapshot = '0',
-  proPeriodCostSnapshotAt = null,
-  lastPeriodCost = '0',
+const orgMemberUsageRow = {
+  userId: 'owner-1',
+  role: 'owner',
+  currentPeriodCost: '350',
+  departedMemberUsage: '25',
+}
+
+/**
+ * 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,
 }: {
-  currentPeriodCost?: string
-  proPeriodCostSnapshot?: string
-  proPeriodCostSnapshotAt?: Date | null
-  lastPeriodCost?: string
-}) {
-  return buildSelectChain([
-    {
-      currentPeriodCost,
-      proPeriodCostSnapshot,
-      proPeriodCostSnapshotAt,
-      lastPeriodCost,
-    },
-  ])
+  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 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),
-    })),
-  }
-}
-
-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<void>) =>
-      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,161 +580,93 @@ 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',
         },
-      ])
-    )
+      ],
+      lockedOwnerRows: [{ userId: 'member-1' }],
+      lockedMemberUsageRows: [
+        {
+          userId: 'owner-1',
+          role: 'member',
+          currentPeriodCost: '350',
+          departedMemberUsage: '25',
+        },
+        {
+          userId: 'member-1',
+          role: 'owner',
+          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([
-        {
-          userId: 'owner-1',
-          role: 'member',
-          currentPeriodCost: '350',
-          departedMemberUsage: '25',
-        },
-        {
-          userId: 'member-1',
-          role: 'owner',
-          currentPeriodCost: '25',
-          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()
   })
 })

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

apps/sim/lib/billing/webhooks/outbox-handlers.test.ts

@@ -1,36 +1,25 @@
 /**
  * @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 { 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,
-    }
+const { mockGetPlanByName, mockResolveDefaultPaymentMethod, stripeMock } = vi.hoisted(() => {
+  const stripeMock = {
+    subscriptions: {
+      retrieve: vi.fn(),
+      update: vi.fn(),
+    },
   }
-)
-
-vi.mock('@sim/db', () => {
-  const makeChain = () => {
-    const chain: Record<string, unknown> = {}
-    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)
 

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')

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'")

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

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(),
 }))

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

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,

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(),

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<unknown>) => 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<string, unknown> = {}) => ({
@@ -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<Record<string, unknown>>) => {
-    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')

apps/sim/lib/logs/execution/logging-session.test.ts

@@ -1,39 +1,11 @@
-import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { dbChainMockFns, resetDbChainMock } from '@sim/testing'
+import { afterAll, 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,
-  }
-})
+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)
   })
 })

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<string, unknown> => {
     const o: Record<string, unknown> = {}
@@ -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<string, unknown> = {}
-  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<string, unknown> = {}) {
   return {
@@ -136,12 +122,12 @@ function baseParams(overrides: Partial<ListLogsParams> = {}): 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')
   })

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<string, unknown>) {
-  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' },

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,

apps/sim/lib/table/snapshot-cache.test.ts

@@ -1,26 +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 {
-  mockDb,
-  mockSelectExportRowPage,
-  mockCreateMultipartUpload,
-  mockHeadObject,
-  mockDeleteFile,
-} = vi.hoisted(() => {
-  const limit = vi.fn()
-  return {
-    mockDb: { limit, select: () => ({ from: () => ({ where: () => ({ limit }) }) }) },
+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([

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<SubBlockConfig>[]): { 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')
   })
 })

apps/sim/lib/webhooks/polling/utils.test.ts

@@ -1,37 +1,14 @@
 /**
  * @vitest-environment node
  */
-import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { account } from '@sim/db/schema'
+import { dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing'
+import { afterAll, 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[] }>,
-  }
-})
-
-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<ReturnType<typeof resolveOAuthAccountId>>)
-    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<ReturnType<typeof resolveOAuthAccountId>>)
-    mockSelectRows.mockResolvedValue([])
 
     await expect(
       resolveOAuthCredential(makeWebhook({ credentialId: 'cred-1' }), 'google-email', 'req-1')

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<string, unknown> },
   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>): 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([])

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'
 

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<string, unknown> = {}) {
 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<unknown>) => {
-      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<unknown>) =>
-      callback(harness.tx)
+    dbChainMockFns.transaction.mockImplementation(
+      async (callback: (tx: DbOrTx) => Promise<unknown>) => callback(harness.tx)
     )
     return harness
   }

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

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<void>) =>
-      callback(tx)
+    dbChainMockFns.transaction.mockImplementation(
+      async (callback: (trx: typeof tx) => Promise<void>) => 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()
   })
 })

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()
+} = vi.hoisted(() => ({
+  mockEnsureUserInOrganizationTx: vi.fn(),
+  mockSyncUsageLimitsFromSubscription: vi.fn(),
+  mockReapplyPaidOrgJoinBillingForExistingMemberTx: vi.fn(),
+  mockAcquireOrganizationMutationLock: vi.fn(),
+  mockAcquireInvitationMutationLocks: vi.fn(),
+  mockChangeWorkspaceStoragePayersInTx: 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.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()
   })
 })