diff --git a/apps/sim/app/api/folders/[id]/route.test.ts b/apps/sim/app/api/folders/[id]/route.test.ts
index 95cb3d53b05..bbe7dfeba63 100644
--- a/apps/sim/app/api/folders/[id]/route.test.ts
+++ b/apps/sim/app/api/folders/[id]/route.test.ts
@@ -7,17 +7,22 @@ import {
   auditMock,
   authMockFns,
   createMockRequest,
+  dbChainMock,
+  dbChainMockFns,
   type MockUser,
   permissionsMock,
   permissionsMockFns,
+  queueTableRows,
+  resetDbChainMock,
+  schemaMock,
   workflowsOrchestrationMock,
   workflowsOrchestrationMockFns,
   workflowsUtilsMock,
   workflowsUtilsMockFns,
 } from '@sim/testing'
-import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
-const { mockLogger, mockDbRef } = vi.hoisted(() => {
+const { mockLogger } = vi.hoisted(() => {
   const logger = {
     info: vi.fn(),
     warn: vi.fn(),
@@ -29,7 +34,6 @@ const { mockLogger, mockDbRef } = vi.hoisted(() => {
   }
   return {
     mockLogger: logger,
-    mockDbRef: { current: null as any },
   }
 })
 
@@ -45,23 +49,12 @@ vi.mock('@sim/logger', () => ({
   getRequestContext: () => undefined,
 }))
 vi.mock('@/lib/workspaces/permissions/utils', () => permissionsMock)
-vi.mock('@sim/db', () => ({
-  get db() {
-    return mockDbRef.current
-  },
-}))
+vi.mock('@sim/db', () => dbChainMock)
 vi.mock('@/lib/workflows/orchestration', () => workflowsOrchestrationMock)
 vi.mock('@/lib/workflows/utils', () => workflowsUtilsMock)
 
 import { DELETE, PUT } from '@/app/api/folders/[id]/route'
 
-interface FolderDbMockOptions {
-  folderLookupResult?: any
-  updateResult?: any[]
-  throwError?: boolean
-  circularCheckResults?: any[]
-}
-
 const TEST_USER: MockUser = {
   id: 'user-123',
   email: 'test@example.com',
@@ -80,57 +73,16 @@ const mockFolder = {
   updatedAt: new Date('2024-01-01T00:00:00Z'),
 }
 
-function createFolderDbMock(options: FolderDbMockOptions = {}) {
-  const {
-    folderLookupResult = mockFolder,
-    updateResult = [{ ...mockFolder, name: 'Updated Folder' }],
-    throwError = false,
-    circularCheckResults = [],
-  } = options
-
-  let callCount = 0
-
-  const mockSelect = vi.fn().mockImplementation(() => ({
-    from: vi.fn().mockImplementation(() => ({
-      where: vi.fn().mockImplementation(() => ({
-        then: vi.fn().mockImplementation((callback) => {
-          if (throwError) {
-            throw new Error('Database error')
-          }
-
-          callCount++
-          if (callCount === 1) {
-            const result = folderLookupResult === undefined ? [] : [folderLookupResult]
-            return Promise.resolve(callback(result))
-          }
-          if (callCount > 1 && circularCheckResults.length > 0) {
-            const index = callCount - 2
-            const result = circularCheckResults[index] ? [circularCheckResults[index]] : []
-            return Promise.resolve(callback(result))
-          }
-          return Promise.resolve(callback([]))
-        }),
-      })),
-    })),
-  }))
-
-  const mockUpdate = vi.fn().mockImplementation(() => ({
-    set: vi.fn().mockImplementation(() => ({
-      where: vi.fn().mockImplementation(() => ({
-        returning: vi.fn().mockReturnValue(updateResult),
-      })),
-    })),
-  }))
-
-  const mockDelete = vi.fn().mockImplementation(() => ({
-    where: vi.fn().mockImplementation(() => Promise.resolve()),
-  }))
+/** Queues the folder-existence lookup the route runs before authorizing. */
+function queueFolderLookup(folder: Record<string, unknown> = mockFolder) {
+  queueTableRows(schemaMock.workflowFolder, [folder])
+}
 
-  return {
-    select: mockSelect,
-    update: mockUpdate,
-    delete: mockDelete,
-  }
+/** Makes the next folder lookup throw, exercising the route's 500 path. */
+function failFolderLookup() {
+  dbChainMockFns.where.mockImplementationOnce(() => {
+    throw new Error('Database error')
+  })
 }
 
 function mockAuthenticatedUser(user?: MockUser) {
@@ -142,11 +94,15 @@ function mockUnauthenticated() {
 }
 
 describe('Individual Folder API Route', () => {
+  afterAll(() => {
+    resetDbChainMock()
+  })
+
   beforeEach(() => {
     vi.clearAllMocks()
+    resetDbChainMock()
 
     mockGetUserEntityPermissions.mockResolvedValue('admin')
-    mockDbRef.current = createFolderDbMock()
     mockPerformDeleteFolder.mockResolvedValue({
       success: true,
       deletedItems: { folders: 1, workflows: 0 },
@@ -193,6 +149,7 @@ describe('Individual Folder API Route', () => {
     it('should update folder successfully', async () => {
       mockAuthenticatedUser()
 
+      queueFolderLookup()
       const req = createMockRequest('PUT', {
         name: 'Updated Folder Name',
         color: '#FF0000',
@@ -213,6 +170,7 @@ describe('Individual Folder API Route', () => {
     it('should update parent folder successfully', async () => {
       mockAuthenticatedUser()
 
+      queueFolderLookup()
       const req = createMockRequest('PUT', {
         name: 'Updated Folder',
         parentId: 'parent-folder-1',
@@ -244,6 +202,7 @@ describe('Individual Folder API Route', () => {
       mockAuthenticatedUser()
       mockGetUserEntityPermissions.mockResolvedValue('read')
 
+      queueFolderLookup()
       const req = createMockRequest('PUT', {
         name: 'Updated Folder',
       })
@@ -261,6 +220,7 @@ describe('Individual Folder API Route', () => {
       mockAuthenticatedUser()
       mockGetUserEntityPermissions.mockResolvedValue('write')
 
+      queueFolderLookup()
       const req = createMockRequest('PUT', {
         name: 'Updated Folder',
       })
@@ -278,6 +238,7 @@ describe('Individual Folder API Route', () => {
       mockAuthenticatedUser()
       mockGetUserEntityPermissions.mockResolvedValue('admin')
 
+      queueFolderLookup()
       const req = createMockRequest('PUT', {
         name: 'Updated Folder',
       })
@@ -294,6 +255,7 @@ describe('Individual Folder API Route', () => {
     it('should return 400 when trying to set folder as its own parent', async () => {
       mockAuthenticatedUser()
 
+      queueFolderLookup()
       const req = createMockRequest('PUT', {
         name: 'Updated Folder',
         parentId: 'folder-1',
@@ -311,6 +273,7 @@ describe('Individual Folder API Route', () => {
     it('should trim folder name when updating', async () => {
       mockAuthenticatedUser()
 
+      queueFolderLookup()
       const req = createMockRequest('PUT', {
         name: '  Folder With Spaces  ',
       })
@@ -325,9 +288,7 @@ describe('Individual Folder API Route', () => {
     it('should handle database errors gracefully', async () => {
       mockAuthenticatedUser()
 
-      mockDbRef.current = createFolderDbMock({
-        throwError: true,
-      })
+      failFolderLookup()
 
       const req = createMockRequest('PUT', {
         name: 'Updated Folder',
@@ -350,6 +311,7 @@ describe('Individual Folder API Route', () => {
     it('should handle empty folder name', async () => {
       mockAuthenticatedUser()
 
+      queueFolderLookup()
       const req = createMockRequest('PUT', {
         name: '',
       })
@@ -383,13 +345,11 @@ describe('Individual Folder API Route', () => {
     it('should prevent circular references when updating parent', async () => {
       mockAuthenticatedUser()
 
-      mockDbRef.current = createFolderDbMock({
-        folderLookupResult: {
-          id: 'folder-3',
-          parentId: null,
-          name: 'Folder 3',
-          workspaceId: 'workspace-123',
-        },
+      queueFolderLookup({
+        id: 'folder-3',
+        parentId: null,
+        name: 'Folder 3',
+        workspaceId: 'workspace-123',
       })
 
       workflowsUtilsMockFns.mockCheckForCircularReference.mockResolvedValue(true)
@@ -417,9 +377,7 @@ describe('Individual Folder API Route', () => {
     it('should delete folder and all contents successfully', async () => {
       mockAuthenticatedUser()
 
-      mockDbRef.current = createFolderDbMock({
-        folderLookupResult: mockFolder,
-      })
+      queueFolderLookup()
 
       const req = createMockRequest('DELETE')
       const params = Promise.resolve({ id: 'folder-1' })
@@ -457,6 +415,7 @@ describe('Individual Folder API Route', () => {
       mockAuthenticatedUser()
       mockGetUserEntityPermissions.mockResolvedValue('read')
 
+      queueFolderLookup()
       const req = createMockRequest('DELETE')
       const params = Promise.resolve({ id: 'folder-1' })
 
@@ -472,9 +431,7 @@ describe('Individual Folder API Route', () => {
       mockAuthenticatedUser()
       mockGetUserEntityPermissions.mockResolvedValue('write')
 
-      mockDbRef.current = createFolderDbMock({
-        folderLookupResult: mockFolder,
-      })
+      queueFolderLookup()
 
       const req = createMockRequest('DELETE')
       const params = Promise.resolve({ id: 'folder-1' })
@@ -492,9 +449,7 @@ describe('Individual Folder API Route', () => {
       mockAuthenticatedUser()
       mockGetUserEntityPermissions.mockResolvedValue('admin')
 
-      mockDbRef.current = createFolderDbMock({
-        folderLookupResult: mockFolder,
-      })
+      queueFolderLookup()
 
       const req = createMockRequest('DELETE')
       const params = Promise.resolve({ id: 'folder-1' })
@@ -511,9 +466,7 @@ describe('Individual Folder API Route', () => {
     it('should handle database errors during deletion', async () => {
       mockAuthenticatedUser()
 
-      mockDbRef.current = createFolderDbMock({
-        throwError: true,
-      })
+      failFolderLookup()
 
       const req = createMockRequest('DELETE')
       const params = Promise.resolve({ id: 'folder-1' })
diff --git a/apps/sim/app/api/knowledge/[id]/documents/[documentId]/route.test.ts b/apps/sim/app/api/knowledge/[id]/documents/[documentId]/route.test.ts
index 7e08b67fa71..a512d203d43 100644
--- a/apps/sim/app/api/knowledge/[id]/documents/[documentId]/route.test.ts
+++ b/apps/sim/app/api/knowledge/[id]/documents/[documentId]/route.test.ts
@@ -3,26 +3,17 @@
  *
  * @vitest-environment node
  */
-import { auditMock, authMockFns, createMockRequest, knowledgeApiUtilsMock } from '@sim/testing'
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
-
-const { mockDbChain } = vi.hoisted(() => {
-  const mockDbChain = {
-    select: vi.fn().mockReturnThis(),
-    from: vi.fn().mockReturnThis(),
-    where: vi.fn().mockReturnThis(),
-    limit: vi.fn().mockReturnThis(),
-    update: vi.fn().mockReturnThis(),
-    set: vi.fn().mockReturnThis(),
-    delete: vi.fn().mockReturnThis(),
-    transaction: vi.fn(),
-  }
-  return { mockDbChain }
-})
+import {
+  auditMock,
+  authMockFns,
+  createMockRequest,
+  dbChainMock,
+  knowledgeApiUtilsMock,
+  resetDbChainMock,
+} from '@sim/testing'
+import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
-vi.mock('@sim/db', () => ({
-  db: mockDbChain,
-}))
+vi.mock('@sim/db', () => dbChainMock)
 
 vi.mock('@/app/api/knowledge/utils', () => knowledgeApiUtilsMock)
 
@@ -82,20 +73,9 @@ describe('Document By ID API Route', () => {
     deletedAt: null,
   }
 
-  const resetMocks = () => {
-    vi.clearAllMocks()
-    Object.values(mockDbChain).forEach((fn) => {
-      if (typeof fn === 'function') {
-        fn.mockClear().mockReset()
-        if (fn !== mockDbChain.transaction) {
-          fn.mockReturnThis()
-        }
-      }
-    })
-  }
-
   beforeEach(() => {
-    resetMocks()
+    vi.clearAllMocks()
+    resetDbChainMock()
 
     vi.stubGlobal('crypto', {
       randomUUID: vi.fn().mockReturnValue('mock-uuid-1234-5678'),
@@ -106,6 +86,10 @@ describe('Document By ID API Route', () => {
     vi.clearAllMocks()
   })
 
+  afterAll(() => {
+    resetDbChainMock()
+  })
+
   describe('GET /api/knowledge/[id]/documents/[documentId]', () => {
     const mockParams = Promise.resolve({ id: 'kb-123', documentId: 'doc-123' })
 
diff --git a/apps/sim/app/api/knowledge/[id]/documents/route.test.ts b/apps/sim/app/api/knowledge/[id]/documents/route.test.ts
index 80af0ba8cc1..26ec37a6786 100644
--- a/apps/sim/app/api/knowledge/[id]/documents/route.test.ts
+++ b/apps/sim/app/api/knowledge/[id]/documents/route.test.ts
@@ -3,29 +3,17 @@
  *
  * @vitest-environment node
  */
-import { auditMock, authMockFns, createMockRequest, knowledgeApiUtilsMock } from '@sim/testing'
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
-
-const { mockDbChain } = vi.hoisted(() => {
-  const mockDbChain = {
-    select: vi.fn().mockReturnThis(),
-    from: vi.fn().mockReturnThis(),
-    where: vi.fn().mockReturnThis(),
-    orderBy: vi.fn().mockReturnThis(),
-    limit: vi.fn().mockReturnThis(),
-    offset: vi.fn().mockReturnThis(),
-    insert: vi.fn().mockReturnThis(),
-    values: vi.fn().mockReturnThis(),
-    update: vi.fn().mockReturnThis(),
-    set: vi.fn().mockReturnThis(),
-    transaction: vi.fn(),
-  }
-  return { mockDbChain }
-})
+import {
+  auditMock,
+  authMockFns,
+  createMockRequest,
+  dbChainMock,
+  knowledgeApiUtilsMock,
+  resetDbChainMock,
+} from '@sim/testing'
+import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
-vi.mock('@sim/db', () => ({
-  db: mockDbChain,
-}))
+vi.mock('@sim/db', () => dbChainMock)
 
 vi.mock('@/app/api/knowledge/utils', () => knowledgeApiUtilsMock)
 
@@ -96,20 +84,9 @@ describe('Knowledge Base Documents API Route', () => {
     deletedAt: null,
   }
 
-  const resetMocks = () => {
-    vi.clearAllMocks()
-    Object.values(mockDbChain).forEach((fn) => {
-      if (typeof fn === 'function') {
-        fn.mockClear().mockReset()
-        if (fn !== mockDbChain.transaction) {
-          fn.mockReturnThis()
-        }
-      }
-    })
-  }
-
   beforeEach(() => {
-    resetMocks()
+    vi.clearAllMocks()
+    resetDbChainMock()
 
     vi.stubGlobal('crypto', {
       randomUUID: vi.fn().mockReturnValue('mock-uuid-1234-5678'),
@@ -120,6 +97,10 @@ describe('Knowledge Base Documents API Route', () => {
     vi.clearAllMocks()
   })
 
+  afterAll(() => {
+    resetDbChainMock()
+  })
+
   describe('GET /api/knowledge/[id]/documents', () => {
     const mockParams = Promise.resolve({ id: 'kb-123' })
 
diff --git a/apps/sim/app/api/knowledge/[id]/documents/upsert/route.test.ts b/apps/sim/app/api/knowledge/[id]/documents/upsert/route.test.ts
index 3feda21aa92..e48ef65e164 100644
--- a/apps/sim/app/api/knowledge/[id]/documents/upsert/route.test.ts
+++ b/apps/sim/app/api/knowledge/[id]/documents/upsert/route.test.ts
@@ -6,23 +6,15 @@
 import {
   auditMock,
   createMockRequest,
+  dbChainMock,
   hybridAuthMock,
   hybridAuthMockFns,
   knowledgeApiUtilsMock,
+  resetDbChainMock,
 } from '@sim/testing'
-import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
-const { mockDbChain } = vi.hoisted(() => {
-  const chain = {
-    select: vi.fn().mockReturnThis(),
-    from: vi.fn().mockReturnThis(),
-    where: vi.fn().mockReturnThis(),
-    limit: vi.fn().mockResolvedValue([]),
-  }
-  return { mockDbChain: chain }
-})
-
-vi.mock('@sim/db', () => ({ db: mockDbChain }))
+vi.mock('@sim/db', () => dbChainMock)
 vi.mock('@/lib/auth/hybrid', () => hybridAuthMock)
 vi.mock('@/app/api/knowledge/utils', () => knowledgeApiUtilsMock)
 vi.mock('@sim/audit', () => auditMock)
@@ -57,10 +49,7 @@ describe('POST /api/knowledge/[id]/documents/upsert', () => {
 
   beforeEach(() => {
     vi.clearAllMocks()
-    mockDbChain.select.mockReturnThis()
-    mockDbChain.from.mockReturnThis()
-    mockDbChain.where.mockReturnThis()
-    mockDbChain.limit.mockResolvedValue([])
+    resetDbChainMock()
 
     hybridAuthMockFns.mockCheckSessionOrInternalAuth.mockResolvedValue({
       success: true,
@@ -81,6 +70,10 @@ describe('POST /api/knowledge/[id]/documents/upsert', () => {
     vi.mocked(processDocumentsWithQueue).mockResolvedValue(undefined as any)
   })
 
+  afterAll(() => {
+    resetDbChainMock()
+  })
+
   const baseBody = {
     filename: 'note.txt',
     fileSize: 11,
diff --git a/apps/sim/app/api/knowledge/[id]/route.test.ts b/apps/sim/app/api/knowledge/[id]/route.test.ts
index 2d0dc1ce2e2..c146ee9a13c 100644
--- a/apps/sim/app/api/knowledge/[id]/route.test.ts
+++ b/apps/sim/app/api/knowledge/[id]/route.test.ts
@@ -3,24 +3,17 @@
  *
  * @vitest-environment node
  */
-import { auditMock, authMockFns, createMockRequest, knowledgeApiUtilsMock } from '@sim/testing'
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
-
-const { mockDbChain } = vi.hoisted(() => {
-  const mockDbChain = {
-    select: vi.fn().mockReturnThis(),
-    from: vi.fn().mockReturnThis(),
-    where: vi.fn().mockReturnThis(),
-    limit: vi.fn().mockReturnThis(),
-    update: vi.fn().mockReturnThis(),
-    set: vi.fn().mockReturnThis(),
-  }
-  return { mockDbChain }
-})
+import {
+  auditMock,
+  authMockFns,
+  createMockRequest,
+  dbChainMock,
+  knowledgeApiUtilsMock,
+  resetDbChainMock,
+} from '@sim/testing'
+import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
-vi.mock('@sim/db', () => ({
-  db: mockDbChain,
-}))
+vi.mock('@sim/db', () => dbChainMock)
 
 vi.mock('@sim/audit', () => auditMock)
 
@@ -64,15 +57,12 @@ describe('Knowledge Base By ID API Route', () => {
 
   const resetMocks = () => {
     vi.clearAllMocks()
-    Object.values(mockDbChain).forEach((fn) => {
-      if (typeof fn === 'function') {
-        fn.mockClear().mockReset().mockReturnThis()
-      }
-    })
+    resetDbChainMock()
   }
 
   beforeEach(() => {
     vi.clearAllMocks()
+    resetDbChainMock()
 
     vi.stubGlobal('crypto', {
       randomUUID: vi.fn().mockReturnValue('mock-uuid-1234-5678'),
@@ -83,6 +73,10 @@ describe('Knowledge Base By ID API Route', () => {
     vi.clearAllMocks()
   })
 
+  afterAll(() => {
+    resetDbChainMock()
+  })
+
   describe('GET /api/knowledge/[id]', () => {
     const mockParams = Promise.resolve({ id: 'kb-123' })
 
diff --git a/apps/sim/app/api/knowledge/route.test.ts b/apps/sim/app/api/knowledge/route.test.ts
index 4ad2aad2acf..499efa0d812 100644
--- a/apps/sim/app/api/knowledge/route.test.ts
+++ b/apps/sim/app/api/knowledge/route.test.ts
@@ -7,29 +7,15 @@ import {
   auditMock,
   authMockFns,
   createMockRequest,
+  dbChainMock,
+  dbChainMockFns,
   permissionsMock,
   permissionsMockFns,
+  resetDbChainMock,
 } from '@sim/testing'
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
-
-const { mockDbChain } = vi.hoisted(() => {
-  const mockDbChain = {
-    select: vi.fn().mockReturnThis(),
-    from: vi.fn().mockReturnThis(),
-    leftJoin: vi.fn().mockReturnThis(),
-    where: vi.fn().mockReturnThis(),
-    groupBy: vi.fn().mockReturnThis(),
-    orderBy: vi.fn().mockResolvedValue([]),
-    limit: vi.fn().mockResolvedValue([]),
-    insert: vi.fn().mockReturnThis(),
-    values: vi.fn().mockResolvedValue(undefined),
-  }
-  return { mockDbChain }
-})
+import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
-vi.mock('@sim/db', () => ({
-  db: mockDbChain,
-}))
+vi.mock('@sim/db', () => dbChainMock)
 
 vi.mock('@sim/audit', () => auditMock)
 
@@ -40,15 +26,7 @@ import { GET, POST } from '@/app/api/knowledge/route'
 describe('Knowledge Base API Route', () => {
   beforeEach(() => {
     vi.clearAllMocks()
-
-    Object.values(mockDbChain).forEach((fn) => {
-      if (typeof fn === 'function') {
-        fn.mockClear()
-        if (fn !== mockDbChain.orderBy && fn !== mockDbChain.values && fn !== mockDbChain.limit) {
-          fn.mockReturnThis()
-        }
-      }
-    })
+    resetDbChainMock()
 
     permissionsMockFns.mockGetUserEntityPermissions.mockResolvedValue('admin')
 
@@ -61,6 +39,10 @@ describe('Knowledge Base API Route', () => {
     vi.clearAllMocks()
   })
 
+  afterAll(() => {
+    resetDbChainMock()
+  })
+
   describe('GET /api/knowledge', () => {
     it('should return unauthorized for unauthenticated user', async () => {
       authMockFns.mockGetSession.mockResolvedValue(null)
@@ -77,7 +59,7 @@ describe('Knowledge Base API Route', () => {
       authMockFns.mockGetSession.mockResolvedValue({
         user: { id: 'user-123', email: 'test@example.com' },
       })
-      mockDbChain.orderBy.mockRejectedValue(new Error('Database error'))
+      dbChainMockFns.orderBy.mockRejectedValueOnce(new Error('Database error'))
 
       const req = createMockRequest('GET')
       const response = await GET(req)
@@ -113,7 +95,7 @@ describe('Knowledge Base API Route', () => {
       expect(data.success).toBe(true)
       expect(data.data.name).toBe(validKnowledgeBaseData.name)
       expect(data.data.description).toBe(validKnowledgeBaseData.description)
-      expect(mockDbChain.insert).toHaveBeenCalled()
+      expect(dbChainMockFns.insert).toHaveBeenCalled()
     })
 
     it('should return unauthorized for unauthenticated user', async () => {
@@ -169,7 +151,7 @@ describe('Knowledge Base API Route', () => {
       expect(data.error).toBe(
         'User does not have permission to create knowledge bases in this workspace'
       )
-      expect(mockDbChain.insert).not.toHaveBeenCalled()
+      expect(dbChainMockFns.insert).not.toHaveBeenCalled()
     })
 
     it('should validate chunking config constraints', async () => {
@@ -219,7 +201,7 @@ describe('Knowledge Base API Route', () => {
       authMockFns.mockGetSession.mockResolvedValue({
         user: { id: 'user-123', email: 'test@example.com' },
       })
-      mockDbChain.values.mockRejectedValue(new Error('Database error'))
+      dbChainMockFns.values.mockRejectedValueOnce(new Error('Database error'))
 
       const req = createMockRequest('POST', validKnowledgeBaseData)
       const response = await POST(req)
diff --git a/apps/sim/app/api/knowledge/search/route.test.ts b/apps/sim/app/api/knowledge/search/route.test.ts
index 385fd1a3ac1..b389990c744 100644
--- a/apps/sim/app/api/knowledge/search/route.test.ts
+++ b/apps/sim/app/api/knowledge/search/route.test.ts
@@ -8,16 +8,18 @@
 import {
   createEnvMock,
   createMockRequest,
+  dbChainMock,
+  dbChainMockFns,
   hybridAuthMockFns,
   knowledgeApiUtilsMock,
   knowledgeApiUtilsMockFns,
+  resetDbChainMock,
   workflowAuthzMockFns,
   workflowsUtilsMock,
 } from '@sim/testing'
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
 const {
-  mockDbChain,
   mockGetDocumentTagDefinitions,
   mockHandleTagOnlySearch,
   mockHandleVectorOnlySearch,
@@ -26,17 +28,6 @@ const {
   mockGenerateSearchEmbedding,
   mockGetDocumentMetadataByIds,
 } = vi.hoisted(() => ({
-  mockDbChain: {
-    select: vi.fn().mockReturnThis(),
-    from: vi.fn().mockReturnThis(),
-    where: vi.fn().mockReturnThis(),
-    orderBy: vi.fn().mockReturnThis(),
-    limit: vi.fn().mockReturnThis(),
-    innerJoin: vi.fn().mockReturnThis(),
-    leftJoin: vi.fn().mockReturnThis(),
-    groupBy: vi.fn().mockReturnThis(),
-    having: vi.fn().mockReturnThis(),
-  },
   mockGetDocumentTagDefinitions: vi.fn(),
   mockHandleTagOnlySearch: vi.fn(),
   mockHandleVectorOnlySearch: vi.fn(),
@@ -60,9 +51,7 @@ vi.mock('drizzle-orm', () => ({
   })),
 }))
 
-vi.mock('@sim/db', () => ({
-  db: mockDbChain,
-}))
+vi.mock('@sim/db', () => dbChainMock)
 
 vi.mock('@/lib/workflows/utils', () => workflowsUtilsMock)
 
@@ -142,12 +131,7 @@ describe('Knowledge Search API Route', () => {
 
   beforeEach(() => {
     vi.clearAllMocks()
-
-    Object.values(mockDbChain).forEach((fn) => {
-      if (typeof fn === 'function') {
-        fn.mockClear().mockReturnThis()
-      }
-    })
+    resetDbChainMock()
 
     mockHandleTagOnlySearch.mockClear()
     mockHandleVectorOnlySearch.mockClear()
@@ -187,6 +171,10 @@ describe('Knowledge Search API Route', () => {
     vi.clearAllMocks()
   })
 
+  afterAll(() => {
+    resetDbChainMock()
+  })
+
   describe('POST /api/knowledge/search', () => {
     const validSearchData = {
       knowledgeBaseIds: 'kb-123',
@@ -216,7 +204,7 @@ describe('Knowledge Search API Route', () => {
         },
       })
 
-      mockDbChain.limit.mockResolvedValue([])
+      dbChainMockFns.limit.mockResolvedValue([])
 
       mockHandleVectorOnlySearch.mockResolvedValue(mockSearchResults)
 
@@ -263,7 +251,7 @@ describe('Knowledge Search API Route', () => {
         .mockResolvedValueOnce({ hasAccess: true, knowledgeBase: multiKbs[0] })
         .mockResolvedValueOnce({ hasAccess: true, knowledgeBase: multiKbs[1] })
 
-      mockDbChain.limit.mockResolvedValue([])
+      dbChainMockFns.limit.mockResolvedValue([])
 
       mockHandleVectorOnlySearch.mockResolvedValue(mockSearchResults)
 
@@ -308,7 +296,7 @@ describe('Knowledge Search API Route', () => {
         },
       })
 
-      mockDbChain.limit.mockResolvedValue([])
+      dbChainMockFns.limit.mockResolvedValue([])
 
       mockHandleVectorOnlySearch.mockResolvedValue(mockSearchResults)
 
@@ -506,7 +494,7 @@ describe('Knowledge Search API Route', () => {
         },
       })
 
-      mockDbChain.limit.mockResolvedValueOnce(mockSearchResults) // Search results
+      dbChainMockFns.limit.mockResolvedValueOnce(mockSearchResults) // Search results
 
       mockFetch.mockResolvedValue({
         ok: true,
@@ -526,7 +514,7 @@ describe('Knowledge Search API Route', () => {
 
     it.concurrent('should handle OpenAI API errors', async () => {
       mockGetUserId.mockResolvedValue('user-123')
-      mockDbChain.limit.mockResolvedValueOnce(mockKnowledgeBases)
+      dbChainMockFns.limit.mockResolvedValueOnce(mockKnowledgeBases)
 
       mockGenerateSearchEmbedding.mockRejectedValueOnce(
         new Error('OpenAI API error: 401 Unauthorized - Invalid API key')
@@ -542,7 +530,7 @@ describe('Knowledge Search API Route', () => {
 
     it.concurrent('should handle missing OpenAI API key', async () => {
       mockGetUserId.mockResolvedValue('user-123')
-      mockDbChain.limit.mockResolvedValueOnce(mockKnowledgeBases)
+      dbChainMockFns.limit.mockResolvedValueOnce(mockKnowledgeBases)
 
       mockGenerateSearchEmbedding.mockRejectedValueOnce(new Error('OPENAI_API_KEY not configured'))
 
@@ -556,7 +544,7 @@ describe('Knowledge Search API Route', () => {
 
     it.concurrent('should handle database errors during search', async () => {
       mockGetUserId.mockResolvedValue('user-123')
-      mockDbChain.limit.mockResolvedValueOnce(mockKnowledgeBases)
+      dbChainMockFns.limit.mockResolvedValueOnce(mockKnowledgeBases)
 
       mockHandleVectorOnlySearch.mockRejectedValueOnce(new Error('Database error'))
 
@@ -570,7 +558,7 @@ describe('Knowledge Search API Route', () => {
 
     it.concurrent('should handle invalid OpenAI response format', async () => {
       mockGetUserId.mockResolvedValue('user-123')
-      mockDbChain.limit.mockResolvedValueOnce(mockKnowledgeBases)
+      dbChainMockFns.limit.mockResolvedValueOnce(mockKnowledgeBases)
 
       mockGenerateSearchEmbedding.mockRejectedValueOnce(
         new Error('Invalid response format from OpenAI embeddings API')
@@ -599,7 +587,7 @@ describe('Knowledge Search API Route', () => {
           },
         })
 
-        mockDbChain.limit.mockResolvedValueOnce(mockSearchResults)
+        dbChainMockFns.limit.mockResolvedValueOnce(mockSearchResults)
 
         mockFetch.mockResolvedValue({
           ok: true,
@@ -647,7 +635,7 @@ describe('Knowledge Search API Route', () => {
           },
         })
 
-        mockDbChain.limit.mockResolvedValueOnce(mockSearchResults)
+        dbChainMockFns.limit.mockResolvedValueOnce(mockSearchResults)
 
         mockFetch.mockResolvedValue({
           ok: true,
@@ -702,7 +690,7 @@ describe('Knowledge Search API Route', () => {
           },
         })
 
-        mockDbChain.limit.mockResolvedValueOnce(mockSearchResults)
+        dbChainMockFns.limit.mockResolvedValueOnce(mockSearchResults)
 
         mockFetch.mockResolvedValue({
           ok: true,
@@ -774,7 +762,7 @@ describe('Knowledge Search API Route', () => {
 
       mockGetDocumentTagDefinitions.mockResolvedValue(mockTagDefinitions)
 
-      mockDbChain.limit.mockResolvedValueOnce(mockTagDefinitions)
+      dbChainMockFns.limit.mockResolvedValueOnce(mockTagDefinitions)
 
       mockHandleTagOnlySearch.mockResolvedValue(mockTaggedResults)
 
@@ -820,7 +808,7 @@ describe('Knowledge Search API Route', () => {
 
       mockGetDocumentTagDefinitions.mockResolvedValue(mockTagDefinitions)
 
-      mockDbChain.limit.mockResolvedValueOnce(mockTagDefinitions)
+      dbChainMockFns.limit.mockResolvedValueOnce(mockTagDefinitions)
 
       mockHandleTagAndVectorSearch.mockResolvedValue(mockSearchResults)
 
@@ -957,7 +945,7 @@ describe('Knowledge Search API Route', () => {
         },
       })
 
-      mockDbChain.limit.mockResolvedValueOnce(mockSearchResults)
+      dbChainMockFns.limit.mockResolvedValueOnce(mockSearchResults)
 
       mockFetch.mockResolvedValue({
         ok: true,
@@ -1015,7 +1003,7 @@ describe('Knowledge Search API Route', () => {
 
       mockHandleTagOnlySearch.mockResolvedValue(mockTaggedResults)
 
-      mockDbChain.limit.mockResolvedValueOnce(mockTagDefinitions)
+      dbChainMockFns.limit.mockResolvedValueOnce(mockTagDefinitions)
 
       const req = createMockRequest('POST', multiKbTagData)
       const response = await POST(req)
@@ -1080,7 +1068,7 @@ describe('Knowledge Search API Route', () => {
         from: vi.fn().mockReturnThis(),
         where: vi.fn().mockResolvedValue([]),
       }
-      mockDbChain.select.mockReturnValueOnce(mockTagDefs)
+      dbChainMockFns.select.mockReturnValueOnce(mockTagDefs)
 
       const req = createMockRequest('POST', {
         knowledgeBaseIds: ['kb-123'],
@@ -1154,7 +1142,7 @@ describe('Knowledge Search API Route', () => {
           .fn()
           .mockResolvedValue([{ tagSlot: 'tag1', displayName: 'tag1', fieldType: 'text' }]),
       }
-      mockDbChain.select.mockReturnValueOnce(mockTagDefs)
+      dbChainMockFns.select.mockReturnValueOnce(mockTagDefs)
 
       const req = createMockRequest('POST', {
         knowledgeBaseIds: ['kb-123'],
@@ -1227,7 +1215,7 @@ describe('Knowledge Search API Route', () => {
           .fn()
           .mockResolvedValue([{ tagSlot: 'tag1', displayName: 'tag1', fieldType: 'text' }]),
       }
-      mockDbChain.select.mockReturnValueOnce(mockTagDefs)
+      dbChainMockFns.select.mockReturnValueOnce(mockTagDefs)
 
       const req = createMockRequest('POST', {
         knowledgeBaseIds: ['kb-123'],
diff --git a/apps/sim/app/api/knowledge/utils.test.ts b/apps/sim/app/api/knowledge/utils.test.ts
index ad9e59113cf..d61e8273ef6 100644
--- a/apps/sim/app/api/knowledge/utils.test.ts
+++ b/apps/sim/app/api/knowledge/utils.test.ts
@@ -6,7 +6,14 @@
  * This file contains unit tests for the knowledge base utility functions,
  * including access checks, document processing, and embedding generation.
  */
-import { defaultMockEnv } from '@sim/testing'
+import {
+  dbChainMock,
+  dbChainMockFns,
+  defaultMockEnv,
+  queueTableRows,
+  resetDbChainMock,
+  schemaMock,
+} from '@sim/testing'
 import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 import * as billingAttributionModule from '@/lib/billing/core/billing-attribution'
 import { env } from '@/lib/core/config/env'
@@ -20,6 +27,7 @@ afterAll(() => {
     delete (env as Record<string, unknown>)[key]
   }
   Object.assign(env, envSnapshot)
+  resetDbChainMock()
   retrySpy.mockRestore()
   vi.mocked(workspacesUtilsModule.getWorkspaceBilledAccountUserId).mockRestore()
   vi.mocked(billingAttributionModule.assertBillingAttributionSnapshot).mockRestore()
@@ -28,13 +36,6 @@ afterAll(() => {
   vi.mocked(billingAttributionModule.toBillingContext).mockRestore()
 })
 
-vi.mock('drizzle-orm', () => ({
-  and: (...args: any[]) => args,
-  eq: (...args: any[]) => args,
-  isNull: () => true,
-  sql: (strings: TemplateStringsArray, ...expr: any[]) => ({ strings, expr }),
-}))
-
 /**
  * Spy on the real documents/utils namespace instead of vi.mock: the shared
  * `@/lib/knowledge/embeddings` module may be cached bound to the real module,
@@ -110,26 +111,6 @@ vi.mock('@/lib/knowledge/documents/document-processor', () => ({
   }),
 }))
 
-const dbOps: {
-  order: string[]
-  insertRecords: any[][]
-  updatePayloads: any[]
-} = {
-  order: [],
-  insertRecords: [],
-  updatePayloads: [],
-}
-
-let kbRows: any[] = []
-let docRows: any[] = []
-let chunkRows: any[] = []
-
-function resetDatasets() {
-  kbRows = []
-  docRows = []
-  chunkRows = []
-}
-
 function createEmbeddingFetchMock() {
   return vi.fn().mockResolvedValue({
     ok: true,
@@ -145,138 +126,7 @@ function createEmbeddingFetchMock() {
 
 vi.stubGlobal('fetch', createEmbeddingFetchMock())
 
-vi.mock('@sim/db', async () => {
-  const { schemaMock } = (await import('@sim/testing')) as typeof import('@sim/testing')
-  const tableNameFor = (table: any) => {
-    if (table === schemaMock.knowledgeBase) return 'knowledge_base'
-    if (table === schemaMock.document) return 'document'
-    if (table === schemaMock.embedding) return 'embedding'
-    return ''
-  }
-  const selectBuilder = {
-    from(table: any) {
-      return {
-        where() {
-          return {
-            limit(n: number) {
-              const tableName = tableNameFor(table)
-
-              if (tableName === 'knowledge_base') {
-                return Promise.resolve(kbRows.slice(0, n))
-              }
-              if (tableName === 'document') {
-                return Promise.resolve(docRows.slice(0, n))
-              }
-              if (tableName === 'embedding') {
-                return Promise.resolve(chunkRows.slice(0, n))
-              }
-
-              return Promise.resolve([])
-            },
-          }
-        },
-        innerJoin() {
-          // document × knowledge_base context JOIN — return the first kb and
-          // doc row merged (covers processDocumentAsync's prefetch).
-          return {
-            leftJoin: () => ({
-              where: () => ({
-                limit: (n: number) =>
-                  Promise.resolve(
-                    kbRows.length > 0 && docRows.length > 0
-                      ? [
-                          { ...kbRows[0], ...docRows[0], billedAccountUserId: 'billing-user-1' },
-                        ].slice(0, n)
-                      : []
-                  ),
-              }),
-            }),
-            where: () => ({
-              limit: (n: number) =>
-                Promise.resolve(
-                  kbRows.length > 0 && docRows.length > 0
-                    ? [{ ...kbRows[0], ...docRows[0] }].slice(0, n)
-                    : []
-                ),
-            }),
-          }
-        },
-      }
-    },
-  }
-
-  return {
-    db: {
-      select: vi.fn(() => selectBuilder),
-      update: (table: any) => ({
-        set: (payload: any) => ({
-          where: () => {
-            const tableName = tableNameFor(table)
-            if (tableName === 'knowledge_base') {
-              dbOps.order.push('updateKb')
-              dbOps.updatePayloads.push(payload)
-            } else if (tableName === 'document') {
-              if (payload.processingStatus !== 'processing') {
-                dbOps.order.push('updateDoc')
-                dbOps.updatePayloads.push(payload)
-              }
-            }
-            return Promise.resolve()
-          },
-        }),
-      }),
-      delete: () => ({
-        where: () => Promise.resolve(),
-      }),
-      insert: () => ({
-        values: (records: any) => {
-          dbOps.order.push('insert')
-          dbOps.insertRecords.push(records)
-          return Promise.resolve()
-        },
-      }),
-      transaction: vi.fn(async (fn: any) => {
-        await fn({
-          select: () => ({
-            from: () => ({
-              innerJoin: () => ({
-                where: () => ({
-                  limit: () => Promise.resolve([{ id: 'doc1' }]),
-                }),
-              }),
-              where: () => ({
-                limit: () => Promise.resolve([{}]),
-              }),
-            }),
-          }),
-          delete: () => ({
-            where: () => Promise.resolve(),
-          }),
-          insert: () => ({
-            values: (records: any) => {
-              dbOps.order.push('insert')
-              dbOps.insertRecords.push(records)
-              return Promise.resolve()
-            },
-          }),
-          update: () => ({
-            set: (payload: any) => ({
-              where: () => {
-                dbOps.updatePayloads.push(payload)
-                const label = payload.processingStatus !== undefined ? 'updateDoc' : 'updateKb'
-                dbOps.order.push(label)
-                return Promise.resolve()
-              },
-            }),
-          }),
-        })
-      }),
-    },
-    document: {},
-    knowledgeBase: {},
-    embedding: {},
-  }
-})
+vi.mock('@sim/db', () => dbChainMock)
 
 import { processDocumentAsync } from '@/lib/knowledge/documents/service'
 import { generateEmbeddings } from '@/lib/knowledge/embeddings'
@@ -288,11 +138,8 @@ import {
 
 describe('Knowledge Utils', () => {
   beforeEach(() => {
-    dbOps.order.length = 0
-    dbOps.insertRecords.length = 0
-    dbOps.updatePayloads.length = 0
-    resetDatasets()
     vi.clearAllMocks()
+    resetDbChainMock()
     // `unstubGlobals: true` removes the module-scope fetch stub after the
     // first test in the worker; re-stub it per test.
     vi.stubGlobal('fetch', createEmbeddingFetchMock())
@@ -310,14 +157,19 @@ describe('Knowledge Utils', () => {
 
   describe('processDocumentAsync', () => {
     it('should insert embeddings before updating document counters', async () => {
-      kbRows.push({
-        id: 'kb1',
-        userId: 'user1',
-        workspaceId: 'workspace1',
-        embeddingModel: 'text-embedding-3-small',
-        chunkingConfig: { maxSize: 1024, minSize: 1, overlap: 200 },
-      })
-      docRows.push({ id: 'doc1', knowledgeBaseId: 'kb1' })
+      /** Context prefetch JOIN (document × knowledge_base × workspace). */
+      queueTableRows(schemaMock.document, [
+        {
+          workspaceId: 'workspace1',
+          knowledgeBaseUserId: 'user1',
+          chunkingConfig: { maxSize: 1024, minSize: 1, overlap: 200 },
+          embeddingModel: 'text-embedding-3-small',
+          billedAccountUserId: 'billing-user-1',
+          uploadedBy: null,
+        },
+      ])
+      /** In-transaction active-document recheck. */
+      queueTableRows(schemaMock.document, [{ id: 'doc1' }])
 
       await processDocumentAsync(
         'kb1',
@@ -343,25 +195,29 @@ describe('Knowledge Utils', () => {
         }
       )
 
-      // Embeddings are inserted first, then the document counter update. A
-      // usage_log billing insert (recordUsage) may trail after updateDoc and is
-      // irrelevant to this ordering invariant, so assert position rather than
-      // exact array equality.
-      expect(dbOps.order[0]).toBe('insert')
-      expect(dbOps.order.indexOf('updateDoc')).toBeGreaterThan(0)
-
-      expect(dbOps.updatePayloads[0]).toMatchObject({
+      /**
+       * Embeddings are inserted first, then the document counter update. The
+       * status→'processing' update precedes both and a usage_log billing insert
+       * (recordUsage) may trail after — assert relative order via the shared
+       * spies' invocation order rather than exact call sequences.
+       */
+      const setPayloads = dbChainMockFns.set.mock.calls.map((call) => call[0])
+      const completedIndex = setPayloads.findIndex((p) => p?.processingStatus === 'completed')
+      expect(setPayloads[completedIndex]).toMatchObject({
         processingStatus: 'completed',
         chunkCount: 2,
       })
 
-      expect(dbOps.insertRecords[0].length).toBe(2)
+      expect(dbChainMockFns.values.mock.calls[0][0]).toHaveLength(2)
+      expect(dbChainMockFns.values.mock.invocationCallOrder[0]).toBeLessThan(
+        dbChainMockFns.set.mock.invocationCallOrder[completedIndex]
+      )
     })
   })
 
   describe('checkKnowledgeBaseAccess', () => {
     it('should return success for owner', async () => {
-      kbRows.push({ id: 'kb1', userId: 'user1' })
+      queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb1', userId: 'user1' }])
       const result = await checkKnowledgeBaseAccess('kb1', 'user1')
 
       expect(result.hasAccess).toBe(true)
@@ -377,7 +233,7 @@ describe('Knowledge Utils', () => {
 
   describe('checkDocumentAccess', () => {
     it('should return unauthorized when user mismatch', async () => {
-      kbRows.push({ id: 'kb1', userId: 'owner' })
+      queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb1', userId: 'owner' }])
       const result = await checkDocumentAccess('kb1', 'doc1', 'intruder')
 
       expect(result.hasAccess).toBe(false)
@@ -389,8 +245,10 @@ describe('Knowledge Utils', () => {
 
   describe('checkChunkAccess', () => {
     it('should fail when document is not completed', async () => {
-      kbRows.push({ id: 'kb1', userId: 'user1' })
-      docRows.push({ id: 'doc1', knowledgeBaseId: 'kb1', processingStatus: 'processing' })
+      queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb1', userId: 'user1' }])
+      queueTableRows(schemaMock.document, [
+        { id: 'doc1', knowledgeBaseId: 'kb1', processingStatus: 'processing' },
+      ])
 
       const result = await checkChunkAccess('kb1', 'doc1', 'chunk1', 'user1')
 
@@ -401,9 +259,11 @@ describe('Knowledge Utils', () => {
     })
 
     it('should return success for valid access', async () => {
-      kbRows.push({ id: 'kb1', userId: 'user1' })
-      docRows.push({ id: 'doc1', knowledgeBaseId: 'kb1', processingStatus: 'completed' })
-      chunkRows.push({ id: 'chunk1', documentId: 'doc1' })
+      queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb1', userId: 'user1' }])
+      queueTableRows(schemaMock.document, [
+        { id: 'doc1', knowledgeBaseId: 'kb1', processingStatus: 'completed' },
+      ])
+      queueTableRows(schemaMock.embedding, [{ id: 'chunk1', documentId: 'doc1' }])
 
       const result = await checkChunkAccess('kb1', 'doc1', 'chunk1', 'user1')
 
diff --git a/apps/sim/app/api/organizations/[id]/invitations/route.test.ts b/apps/sim/app/api/organizations/[id]/invitations/route.test.ts
index 43f4ff6ce8b..1066801bef8 100644
--- a/apps/sim/app/api/organizations/[id]/invitations/route.test.ts
+++ b/apps/sim/app/api/organizations/[id]/invitations/route.test.ts
@@ -1,11 +1,19 @@
 /**
  * @vitest-environment node
  */
-import { auditMock, createMockRequest, createSession, loggerMock } from '@sim/testing'
-import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { member, organization, permissions, user, workspace } from '@sim/db/schema'
+import {
+  auditMock,
+  createMockRequest,
+  createSession,
+  dbChainMock,
+  loggerMock,
+  queueTableRows,
+  resetDbChainMock,
+} from '@sim/testing'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
 const {
-  mockDbState,
   mockGetSession,
   mockValidateInvitationsAllowed,
   mockValidateSeatAvailability,
@@ -14,9 +22,6 @@ const {
   mockCancelPendingInvitation,
   mockGrantWorkspaceAccessDirectly,
 } = vi.hoisted(() => ({
-  mockDbState: {
-    selectResults: [] as any[],
-  },
   mockGetSession: vi.fn(),
   mockValidateInvitationsAllowed: vi.fn(),
   mockValidateSeatAvailability: vi.fn(),
@@ -26,77 +31,7 @@ const {
   mockGrantWorkspaceAccessDirectly: vi.fn(),
 }))
 
-function createSelectChain() {
-  const chain: any = {}
-  chain.from = vi.fn().mockReturnValue(chain)
-  chain.innerJoin = vi.fn().mockReturnValue(chain)
-  chain.leftJoin = vi.fn().mockReturnValue(chain)
-  chain.where = vi.fn().mockReturnValue(chain)
-  chain.orderBy = vi.fn().mockReturnValue(chain)
-  chain.limit = vi
-    .fn()
-    .mockImplementation(() => Promise.resolve(mockDbState.selectResults.shift() ?? []))
-  chain.then = vi.fn().mockImplementation((callback: (rows: any[]) => unknown) => {
-    const rows = mockDbState.selectResults.shift() ?? []
-    return Promise.resolve(callback(rows))
-  })
-  return chain
-}
-
-vi.mock('@sim/db', () => ({
-  db: {
-    select: vi.fn().mockImplementation(() => createSelectChain()),
-  },
-}))
-
-vi.mock('@sim/db/schema', () => ({
-  invitation: {
-    id: 'invitation.id',
-    organizationId: 'invitation.organizationId',
-    status: 'invitation.status',
-    email: 'invitation.email',
-    kind: 'invitation.kind',
-    role: 'invitation.role',
-    inviterId: 'invitation.inviterId',
-    expiresAt: 'invitation.expiresAt',
-    createdAt: 'invitation.createdAt',
-  },
-  member: {
-    organizationId: 'member.organizationId',
-    userId: 'member.userId',
-    role: 'member.role',
-  },
-  organization: {
-    id: 'organization.id',
-    name: 'organization.name',
-  },
-  user: {
-    id: 'user.id',
-    name: 'user.name',
-    email: 'user.email',
-  },
-  workspace: {
-    id: 'workspace.id',
-    name: 'workspace.name',
-    organizationId: 'workspace.organizationId',
-    workspaceMode: 'workspace.workspaceMode',
-  },
-  permissions: {
-    entityId: 'permissions.entityId',
-    entityType: 'permissions.entityType',
-    userId: 'permissions.userId',
-  },
-  invitationWorkspaceGrant: {
-    invitationId: 'invitationWorkspaceGrant.invitationId',
-    workspaceId: 'invitationWorkspaceGrant.workspaceId',
-  },
-}))
-
-vi.mock('drizzle-orm', () => ({
-  and: vi.fn((...conditions: unknown[]) => ({ type: 'and', conditions })),
-  eq: vi.fn((field: unknown, value: unknown) => ({ field, value })),
-  inArray: vi.fn((field: unknown, values: unknown[]) => ({ field, values })),
-}))
+vi.mock('@sim/db', () => dbChainMock)
 
 vi.mock('@sim/logger', () => loggerMock)
 
@@ -140,10 +75,23 @@ vi.mock('@/ee/access-control/utils/permission-check', () => ({
 
 import { POST } from '@/app/api/organizations/[id]/invitations/route'
 
+/** Queues the caller's admin-role check followed by the org-name lookup. */
+function queueOwnerAndOrg() {
+  queueTableRows(member, [{ role: 'owner' }])
+  queueTableRows(organization, [{ name: 'Org One' }])
+}
+
+/** Queues the inviter-details lookup that precedes invitation/email sends. */
+function queueInviterRow() {
+  queueTableRows(user, [{ name: 'Owner', email: 'owner@example.com' }])
+}
+
+afterAll(resetDbChainMock)
+
 describe('POST /api/organizations/[id]/invitations', () => {
   beforeEach(() => {
     vi.clearAllMocks()
-    mockDbState.selectResults = []
+    resetDbChainMock()
     mockValidateInvitationsAllowed.mockResolvedValue(undefined)
     mockValidateSeatAvailability.mockResolvedValue({
       canInvite: true,
@@ -164,13 +112,11 @@ describe('POST /api/organizations/[id]/invitations', () => {
     mockGetSession.mockResolvedValue(
       createSession({ userId: 'user-1', email: 'owner@example.com', name: 'Owner' })
     )
-    mockDbState.selectResults = [
-      [{ role: 'owner' }],
-      [{ name: 'Org One' }],
-      [],
-      [],
-      [{ name: 'Owner', email: 'owner@example.com' }],
-    ]
+    queueOwnerAndOrg()
+    // Explicit empty existing-members set: the query joins `user`, so it must
+    // not fall through to the inviter row queued on the user table.
+    queueTableRows(member, [])
+    queueInviterRow()
 
     const response = await POST(
       createMockRequest(
@@ -202,17 +148,16 @@ describe('POST /api/organizations/[id]/invitations', () => {
     mockGetSession.mockResolvedValue(
       createSession({ userId: 'user-1', email: 'owner@example.com', name: 'Owner' })
     )
-    mockDbState.selectResults = [
-      [{ role: 'owner' }],
-      [{ name: 'Org One' }],
-      [{ id: 'ws-1', name: 'Workspace 1', organizationId: 'org-1', workspaceMode: 'organization' }],
-      [{ id: 'ws-2', name: 'Workspace 2', organizationId: 'org-1', workspaceMode: 'organization' }],
-      [{ userId: 'user-2', userEmail: 'member@example.com' }],
-      [],
-      [{ userId: 'user-2', workspaceId: 'ws-1' }],
-      [],
-      [{ name: 'Owner', email: 'owner@example.com' }],
-    ]
+    queueOwnerAndOrg()
+    queueTableRows(workspace, [
+      { id: 'ws-1', name: 'Workspace 1', organizationId: 'org-1', workspaceMode: 'organization' },
+    ])
+    queueTableRows(workspace, [
+      { id: 'ws-2', name: 'Workspace 2', organizationId: 'org-1', workspaceMode: 'organization' },
+    ])
+    queueTableRows(member, [{ userId: 'user-2', userEmail: 'member@example.com' }])
+    queueTableRows(permissions, [{ userId: 'user-2', workspaceId: 'ws-1' }])
+    queueInviterRow()
 
     const response = await POST(
       createMockRequest(
@@ -259,17 +204,15 @@ describe('POST /api/organizations/[id]/invitations', () => {
     mockGrantWorkspaceAccessDirectly
       .mockResolvedValueOnce({ outcome: 'added', permission: 'write' })
       .mockRejectedValueOnce(new Error('db blip'))
-    mockDbState.selectResults = [
-      [{ role: 'owner' }],
-      [{ name: 'Org One' }],
-      [{ id: 'ws-1', name: 'Workspace 1', organizationId: 'org-1', workspaceMode: 'organization' }],
-      [{ id: 'ws-2', name: 'Workspace 2', organizationId: 'org-1', workspaceMode: 'organization' }],
-      [{ userId: 'user-2', userEmail: 'member@example.com' }],
-      [],
-      [],
-      [],
-      [{ name: 'Owner', email: 'owner@example.com' }],
-    ]
+    queueOwnerAndOrg()
+    queueTableRows(workspace, [
+      { id: 'ws-1', name: 'Workspace 1', organizationId: 'org-1', workspaceMode: 'organization' },
+    ])
+    queueTableRows(workspace, [
+      { id: 'ws-2', name: 'Workspace 2', organizationId: 'org-1', workspaceMode: 'organization' },
+    ])
+    queueTableRows(member, [{ userId: 'user-2', userEmail: 'member@example.com' }])
+    queueInviterRow()
 
     const response = await POST(
       createMockRequest(
@@ -301,19 +244,15 @@ describe('POST /api/organizations/[id]/invitations', () => {
     mockGrantWorkspaceAccessDirectly
       .mockResolvedValueOnce({ outcome: 'added', permission: 'write' })
       .mockRejectedValueOnce(new Error('db blip'))
-    mockDbState.selectResults = [
-      [{ role: 'owner' }],
-      [{ name: 'Org One' }],
-      [{ id: 'ws-1', name: 'Workspace 1', organizationId: 'org-1', workspaceMode: 'organization' }],
-      [
-        { userId: 'user-a', userEmail: 'a@example.com' },
-        { userId: 'user-b', userEmail: 'b@example.com' },
-      ],
-      [],
-      [],
-      [],
-      [{ name: 'Owner', email: 'owner@example.com' }],
-    ]
+    queueOwnerAndOrg()
+    queueTableRows(workspace, [
+      { id: 'ws-1', name: 'Workspace 1', organizationId: 'org-1', workspaceMode: 'organization' },
+    ])
+    queueTableRows(member, [
+      { userId: 'user-a', userEmail: 'a@example.com' },
+      { userId: 'user-b', userEmail: 'b@example.com' },
+    ])
+    queueInviterRow()
 
     const response = await POST(
       createMockRequest(
@@ -340,15 +279,12 @@ describe('POST /api/organizations/[id]/invitations', () => {
     mockGetSession.mockResolvedValue(
       createSession({ userId: 'user-1', email: 'owner@example.com', name: 'Owner' })
     )
-    mockDbState.selectResults = [
-      [{ role: 'owner' }],
-      [{ name: 'Org One' }],
-      [{ id: 'ws-1', organizationId: 'org-1', workspaceMode: 'organization' }],
-      [{ userId: 'user-2', userEmail: 'member@example.com' }],
-      [],
-      [{ userId: 'user-2', workspaceId: 'ws-1' }],
-      [],
-    ]
+    queueOwnerAndOrg()
+    queueTableRows(workspace, [
+      { id: 'ws-1', organizationId: 'org-1', workspaceMode: 'organization' },
+    ])
+    queueTableRows(member, [{ userId: 'user-2', userEmail: 'member@example.com' }])
+    queueTableRows(permissions, [{ userId: 'user-2', workspaceId: 'ws-1' }])
 
     const response = await POST(
       createMockRequest(
@@ -373,16 +309,12 @@ describe('POST /api/organizations/[id]/invitations', () => {
     mockGetSession.mockResolvedValue(
       createSession({ userId: 'user-1', email: 'owner@example.com', name: 'Owner' })
     )
-    mockDbState.selectResults = [
-      [{ role: 'owner' }],
-      [{ name: 'Org One' }],
-      [{ id: 'ws-1', name: 'Workspace 1', organizationId: 'org-1', workspaceMode: 'organization' }],
-      [{ userId: 'user-2', userEmail: 'member@example.com' }],
-      [],
-      [],
-      [],
-      [{ name: 'Owner', email: 'owner@example.com' }],
-    ]
+    queueOwnerAndOrg()
+    queueTableRows(workspace, [
+      { id: 'ws-1', name: 'Workspace 1', organizationId: 'org-1', workspaceMode: 'organization' },
+    ])
+    queueTableRows(member, [{ userId: 'user-2', userEmail: 'member@example.com' }])
+    queueInviterRow()
 
     const response = await POST(
       createMockRequest(
@@ -427,12 +359,8 @@ describe('POST /api/organizations/[id]/invitations', () => {
     mockGetSession.mockResolvedValue(
       createSession({ userId: 'user-1', email: 'owner@example.com', name: 'Owner' })
     )
-    mockDbState.selectResults = [
-      [{ role: 'owner' }],
-      [{ name: 'Org One' }],
-      [{ userId: 'user-2', userEmail: 'member@example.com' }],
-      [],
-    ]
+    queueOwnerAndOrg()
+    queueTableRows(member, [{ userId: 'user-2', userEmail: 'member@example.com' }])
 
     const response = await POST(
       createMockRequest(
@@ -456,13 +384,11 @@ describe('POST /api/organizations/[id]/invitations', () => {
     mockGetSession.mockResolvedValue(
       createSession({ userId: 'user-1', email: 'owner@example.com', name: 'Owner' })
     )
-    mockDbState.selectResults = [
-      [{ role: 'owner' }],
-      [{ name: 'Org One' }],
-      [],
-      [],
-      [{ name: 'Owner', email: 'owner@example.com' }],
-    ]
+    queueOwnerAndOrg()
+    // Explicit empty existing-members set: the query joins `user`, so it must
+    // not fall through to the inviter row queued on the user table.
+    queueTableRows(member, [])
+    queueInviterRow()
     mockSendInvitationEmail.mockResolvedValue({ success: false, error: 'mailer unavailable' })
 
     const response = await POST(
diff --git a/apps/sim/app/api/organizations/[id]/permission-groups/utils.test.ts b/apps/sim/app/api/organizations/[id]/permission-groups/utils.test.ts
index 43f46c0beea..2e89f6db56f 100644
--- a/apps/sim/app/api/organizations/[id]/permission-groups/utils.test.ts
+++ b/apps/sim/app/api/organizations/[id]/permission-groups/utils.test.ts
@@ -1,32 +1,13 @@
 /**
  * @vitest-environment node
  */
-import { beforeEach, describe, expect, it, vi } from 'vitest'
-
-const {
-  mockIsOrganizationAdminOrOwner,
-  mockIsOrganizationOnEnterprisePlan,
-  mockConflictRows,
-  mockAllMembersRows,
-} = vi.hoisted(() => ({
+import { permissionGroup, permissionGroupMember } from '@sim/db/schema'
+import { dbChainMock, queueTableRows, resetDbChainMock } from '@sim/testing'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { mockIsOrganizationAdminOrOwner, mockIsOrganizationOnEnterprisePlan } = vi.hoisted(() => ({
   mockIsOrganizationAdminOrOwner: vi.fn<() => Promise<boolean>>(),
   mockIsOrganizationOnEnterprisePlan: vi.fn<() => Promise<boolean>>(),
-  mockConflictRows: {
-    value: [] as Array<{
-      userId: string
-      userName: string | null
-      userEmail: string | null
-      otherGroupId: string
-      otherGroupName: string
-    }>,
-  },
-  mockAllMembersRows: {
-    value: [] as Array<{
-      conflictingGroupId: string
-      conflictingGroupName: string
-      workspaceName: string
-    }>,
-  },
 }))
 
 vi.mock('@/lib/billing', () => ({
@@ -37,51 +18,20 @@ vi.mock('@/lib/workspaces/permissions/utils', () => ({
   isOrganizationAdminOrOwner: mockIsOrganizationAdminOrOwner,
 }))
 
-vi.mock('@sim/db', () => ({
-  db: {
-    select: vi.fn(() => {
-      const chain: Record<string, unknown> = {}
-      chain.from = vi.fn(() => chain)
-      chain.innerJoin = vi.fn(() => chain)
-      chain.leftJoin = vi.fn(() => chain)
-      chain.where = vi.fn(() => chain)
-      chain.orderBy = vi.fn(() => chain)
-      // findAllMembersWorkspaceConflict ends in `.limit(1)`; findScopeConflicts
-      // awaits the builder directly after `.where()`.
-      chain.limit = vi.fn(() => Promise.resolve(mockAllMembersRows.value))
-      chain.then = (onFulfilled: (rows: unknown) => unknown) =>
-        Promise.resolve(mockConflictRows.value).then(onFulfilled)
-      return chain
-    }),
-  },
-}))
-
-vi.mock('@sim/db/schema', () => ({
-  permissionGroup: {},
-  permissionGroupMember: {},
-  permissionGroupWorkspace: {},
-  user: {},
-  workspace: {},
-}))
-
-vi.mock('drizzle-orm', () => ({
-  and: vi.fn(),
-  asc: vi.fn(),
-  eq: vi.fn(),
-  inArray: vi.fn(),
-  ne: vi.fn(),
-  sql: vi.fn(),
-}))
+vi.mock('@sim/db', () => dbChainMock)
 
 import {
   authorizeOrgAccessControl,
   findAllMembersWorkspaceConflict,
   findScopeConflicts,
-} from './utils'
+} from '@/app/api/organizations/[id]/permission-groups/utils'
+
+afterAll(resetDbChainMock)
 
 describe('authorizeOrgAccessControl', () => {
   beforeEach(() => {
     vi.clearAllMocks()
+    resetDbChainMock()
   })
 
   it('returns a 403 when the user is not an organization admin/owner', async () => {
@@ -122,7 +72,7 @@ describe('authorizeOrgAccessControl', () => {
 describe('findScopeConflicts', () => {
   beforeEach(() => {
     vi.clearAllMocks()
-    mockConflictRows.value = []
+    resetDbChainMock()
   })
 
   const baseParams = {
@@ -141,7 +91,7 @@ describe('findScopeConflicts', () => {
   })
 
   it('returns no conflicts when there are no candidate users', async () => {
-    mockConflictRows.value = [conflictRow('user-1')]
+    queueTableRows(permissionGroupMember, [conflictRow('user-1')])
 
     const conflicts = await findScopeConflicts({ ...baseParams, candidateUserIds: [] })
 
@@ -149,7 +99,7 @@ describe('findScopeConflicts', () => {
   })
 
   it('returns no conflicts when there are no target workspaces', async () => {
-    mockConflictRows.value = [conflictRow('user-1')]
+    queueTableRows(permissionGroupMember, [conflictRow('user-1')])
 
     const conflicts = await findScopeConflicts({ ...baseParams, workspaceIds: [] })
 
@@ -157,7 +107,7 @@ describe('findScopeConflicts', () => {
   })
 
   it('flags a candidate already in another group that shares a workspace', async () => {
-    mockConflictRows.value = [conflictRow('user-1')]
+    queueTableRows(permissionGroupMember, [conflictRow('user-1')])
 
     const conflicts = await findScopeConflicts(baseParams)
 
@@ -166,7 +116,10 @@ describe('findScopeConflicts', () => {
   })
 
   it('returns at most one conflict per user', async () => {
-    mockConflictRows.value = [conflictRow('user-1', 'Marketing'), conflictRow('user-1', 'Sales')]
+    queueTableRows(permissionGroupMember, [
+      conflictRow('user-1', 'Marketing'),
+      conflictRow('user-1', 'Sales'),
+    ])
 
     const conflicts = await findScopeConflicts(baseParams)
 
@@ -175,8 +128,6 @@ describe('findScopeConflicts', () => {
   })
 
   it('returns no conflicts when the query finds no overlapping memberships', async () => {
-    mockConflictRows.value = []
-
     const conflicts = await findScopeConflicts(baseParams)
 
     expect(conflicts).toEqual([])
@@ -186,7 +137,7 @@ describe('findScopeConflicts', () => {
 describe('findAllMembersWorkspaceConflict', () => {
   beforeEach(() => {
     vi.clearAllMocks()
-    mockAllMembersRows.value = []
+    resetDbChainMock()
   })
 
   const baseParams = {
@@ -196,9 +147,9 @@ describe('findAllMembersWorkspaceConflict', () => {
   }
 
   it('returns null when there are no target workspaces', async () => {
-    mockAllMembersRows.value = [
+    queueTableRows(permissionGroup, [
       { conflictingGroupId: 'group-2', conflictingGroupName: 'Marketing', workspaceName: 'Acme' },
-    ]
+    ])
 
     const conflict = await findAllMembersWorkspaceConflict({ ...baseParams, workspaceIds: [] })
 
@@ -206,9 +157,9 @@ describe('findAllMembersWorkspaceConflict', () => {
   })
 
   it('returns the conflicting all-members group sharing a workspace', async () => {
-    mockAllMembersRows.value = [
+    queueTableRows(permissionGroup, [
       { conflictingGroupId: 'group-2', conflictingGroupName: 'Marketing', workspaceName: 'Acme' },
-    ]
+    ])
 
     const conflict = await findAllMembersWorkspaceConflict(baseParams)
 
@@ -220,8 +171,6 @@ describe('findAllMembersWorkspaceConflict', () => {
   })
 
   it('returns null when no other all-members group targets the workspaces', async () => {
-    mockAllMembersRows.value = []
-
     const conflict = await findAllMembersWorkspaceConflict(baseParams)
 
     expect(conflict).toBeNull()
diff --git a/apps/sim/app/api/organizations/[id]/roster/route.test.ts b/apps/sim/app/api/organizations/[id]/roster/route.test.ts
index 82a3eef2fa8..338e7a5e598 100644
--- a/apps/sim/app/api/organizations/[id]/roster/route.test.ts
+++ b/apps/sim/app/api/organizations/[id]/roster/route.test.ts
@@ -1,87 +1,29 @@
 /**
  * @vitest-environment node
  */
-import { createMockRequest, createSession, loggerMock } from '@sim/testing'
-import { beforeEach, describe, expect, it, vi } from 'vitest'
-
-const { mockDbState, mockExpireStaleInvitations, mockGetSession } = vi.hoisted(() => ({
-  mockDbState: {
-    selectResults: [] as unknown[][],
-  },
+import {
+  invitation,
+  invitationWorkspaceGrant,
+  member,
+  permissions,
+  workspace,
+} from '@sim/db/schema'
+import {
+  createMockRequest,
+  createSession,
+  dbChainMock,
+  loggerMock,
+  queueTableRows,
+  resetDbChainMock,
+} from '@sim/testing'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { mockExpireStaleInvitations, mockGetSession } = vi.hoisted(() => ({
   mockExpireStaleInvitations: vi.fn(),
   mockGetSession: vi.fn(),
 }))
 
-function createSelectChain() {
-  const chain = {
-    from: vi.fn(),
-    innerJoin: vi.fn(),
-    leftJoin: vi.fn(),
-    where: vi.fn(),
-    limit: vi.fn(),
-    then: vi.fn(),
-  }
-  chain.from.mockReturnValue(chain)
-  chain.innerJoin.mockReturnValue(chain)
-  chain.leftJoin.mockReturnValue(chain)
-  chain.where.mockReturnValue(chain)
-  chain.limit.mockImplementation(() => Promise.resolve(mockDbState.selectResults.shift() ?? []))
-  chain.then.mockImplementation((resolve: (rows: unknown[]) => unknown) =>
-    Promise.resolve(resolve(mockDbState.selectResults.shift() ?? []))
-  )
-  return chain
-}
-
-vi.mock('@sim/db', () => ({
-  db: {
-    select: vi.fn(() => createSelectChain()),
-  },
-}))
-
-vi.mock('@sim/db/schema', () => ({
-  invitation: {
-    id: 'invitation.id',
-    email: 'invitation.email',
-    role: 'invitation.role',
-    kind: 'invitation.kind',
-    membershipIntent: 'invitation.membershipIntent',
-    organizationId: 'invitation.organizationId',
-    status: 'invitation.status',
-    createdAt: 'invitation.createdAt',
-    expiresAt: 'invitation.expiresAt',
-  },
-  invitationWorkspaceGrant: {
-    invitationId: 'invitationWorkspaceGrant.invitationId',
-    workspaceId: 'invitationWorkspaceGrant.workspaceId',
-    permission: 'invitationWorkspaceGrant.permission',
-  },
-  member: {
-    id: 'member.id',
-    organizationId: 'member.organizationId',
-    userId: 'member.userId',
-    role: 'member.role',
-    createdAt: 'member.createdAt',
-  },
-  permissions: {
-    userId: 'permissions.userId',
-    entityId: 'permissions.entityId',
-    entityType: 'permissions.entityType',
-    permissionType: 'permissions.permissionType',
-    createdAt: 'permissions.createdAt',
-  },
-  user: {
-    id: 'user.id',
-    name: 'user.name',
-    email: 'user.email',
-    image: 'user.image',
-  },
-  workspace: {
-    id: 'workspace.id',
-    name: 'workspace.name',
-    organizationId: 'workspace.organizationId',
-    archivedAt: 'workspace.archivedAt',
-  },
-}))
+vi.mock('@sim/db', () => dbChainMock)
 
 vi.mock('@sim/logger', () => loggerMock)
 
@@ -89,14 +31,6 @@ vi.mock('@sim/platform-authz/workspace', () => ({
   isOrgAdminRole: (role: string | null | undefined) => role === 'owner' || role === 'admin',
 }))
 
-vi.mock('drizzle-orm', () => ({
-  and: vi.fn((...conditions: unknown[]) => ({ type: 'and', conditions })),
-  eq: vi.fn((field: unknown, value: unknown) => ({ field, value })),
-  inArray: vi.fn((field: unknown, values: unknown[]) => ({ field, values })),
-  isNull: vi.fn((field: unknown) => ({ type: 'isNull', field })),
-  sql: vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ strings, values })),
-}))
-
 vi.mock('@/lib/auth', () => ({
   getSession: mockGetSession,
 }))
@@ -128,16 +62,19 @@ const MEMBER_ROWS = [
   },
 ]
 
+afterAll(resetDbChainMock)
+
 describe('GET /api/organizations/[id]/roster', () => {
   beforeEach(() => {
     vi.clearAllMocks()
-    mockDbState.selectResults = []
+    resetDbChainMock()
     mockExpireStaleInvitations.mockResolvedValue(undefined)
   })
 
   it('returns a redacted roster to a target-organization member', async () => {
     mockGetSession.mockResolvedValue(createSession({ userId: 'user-reader' }))
-    mockDbState.selectResults = [[{ role: 'member' }], MEMBER_ROWS]
+    queueTableRows(member, [{ role: 'member' }])
+    queueTableRows(member, MEMBER_ROWS)
 
     const response = await GET(
       createMockRequest('GET', undefined, {}, 'http://localhost/api/organizations/org-1/roster'),
@@ -179,7 +116,6 @@ describe('GET /api/organizations/[id]/roster', () => {
 
   it('denies a workspace collaborator who is not a target-organization member', async () => {
     mockGetSession.mockResolvedValue(createSession({ userId: 'external-user' }))
-    mockDbState.selectResults = [[]]
 
     const response = await GET(
       createMockRequest('GET', undefined, {}, 'http://localhost/api/organizations/org-1/roster'),
@@ -195,37 +131,39 @@ describe('GET /api/organizations/[id]/roster', () => {
 
   it('preserves the full management roster for organization admins', async () => {
     mockGetSession.mockResolvedValue(createSession({ userId: 'user-admin' }))
-    mockDbState.selectResults = [
-      [{ role: 'admin' }],
-      MEMBER_ROWS,
-      [{ id: 'workspace-1', name: 'Workspace One' }],
-      [{ userId: 'user-reader', workspaceId: 'workspace-1', permission: 'write' }],
-      [
-        {
-          userId: 'external-user',
-          userName: 'External User',
-          userEmail: 'external@example.com',
-          userImage: null,
-          workspaceId: 'workspace-1',
-          permission: 'read',
-          createdAt: new Date('2026-03-01T00:00:00.000Z'),
-        },
-      ],
-      [
-        {
-          id: 'invitation-1',
-          email: 'pending@example.com',
-          role: 'member',
-          kind: 'workspace',
-          membershipIntent: 'external',
-          createdAt: new Date('2026-04-01T00:00:00.000Z'),
-          expiresAt: new Date('2026-04-08T00:00:00.000Z'),
-          inviteeName: null,
-          inviteeImage: null,
-        },
-      ],
-      [{ invitationId: 'invitation-1', workspaceId: 'workspace-1', permission: 'read' }],
-    ]
+    queueTableRows(member, [{ role: 'admin' }])
+    queueTableRows(member, MEMBER_ROWS)
+    queueTableRows(workspace, [{ id: 'workspace-1', name: 'Workspace One' }])
+    queueTableRows(permissions, [
+      { userId: 'user-reader', workspaceId: 'workspace-1', permission: 'write' },
+    ])
+    queueTableRows(permissions, [
+      {
+        userId: 'external-user',
+        userName: 'External User',
+        userEmail: 'external@example.com',
+        userImage: null,
+        workspaceId: 'workspace-1',
+        permission: 'read',
+        createdAt: new Date('2026-03-01T00:00:00.000Z'),
+      },
+    ])
+    queueTableRows(invitation, [
+      {
+        id: 'invitation-1',
+        email: 'pending@example.com',
+        role: 'member',
+        kind: 'workspace',
+        membershipIntent: 'external',
+        createdAt: new Date('2026-04-01T00:00:00.000Z'),
+        expiresAt: new Date('2026-04-08T00:00:00.000Z'),
+        inviteeName: null,
+        inviteeImage: null,
+      },
+    ])
+    queueTableRows(invitationWorkspaceGrant, [
+      { invitationId: 'invitation-1', workspaceId: 'workspace-1', permission: 'read' },
+    ])
 
     const response = await GET(
       createMockRequest('GET', undefined, {}, 'http://localhost/api/organizations/org-1/roster'),
diff --git a/apps/sim/app/api/organizations/route.test.ts b/apps/sim/app/api/organizations/route.test.ts
index a3a17208d84..04f0146bff6 100644
--- a/apps/sim/app/api/organizations/route.test.ts
+++ b/apps/sim/app/api/organizations/route.test.ts
@@ -1,11 +1,18 @@
 /**
  * @vitest-environment node
  */
-import { auditMock, createSession, loggerMock } from '@sim/testing'
-import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { member, subscription } from '@sim/db/schema'
+import {
+  auditMock,
+  createSession,
+  dbChainMock,
+  loggerMock,
+  queueTableRows,
+  resetDbChainMock,
+} from '@sim/testing'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
 const {
-  mockDbState,
   mockGetSession,
   mockSetActiveOrganizationForCurrentSession,
   mockCreateOrganizationForTeamPlan,
@@ -13,9 +20,6 @@ const {
   mockAttachOwnedWorkspacesToOrganization,
   WorkspaceOrganizationMembershipConflictError,
 } = vi.hoisted(() => ({
-  mockDbState: {
-    selectResults: [] as any[],
-  },
   mockGetSession: vi.fn(),
   mockSetActiveOrganizationForCurrentSession: vi.fn().mockResolvedValue(undefined),
   mockCreateOrganizationForTeamPlan: vi.fn(),
@@ -24,50 +28,7 @@ const {
   WorkspaceOrganizationMembershipConflictError: class WorkspaceOrganizationMembershipConflictError extends Error {},
 }))
 
-vi.mock('@sim/db', () => ({
-  db: {
-    select: vi.fn().mockImplementation(() => {
-      const chain: any = {}
-      chain.from = vi.fn().mockReturnValue(chain)
-      chain.where = vi.fn().mockReturnValue(chain)
-      chain.limit = vi
-        .fn()
-        .mockImplementation(() => Promise.resolve(mockDbState.selectResults.shift() ?? []))
-      chain.then = vi
-        .fn()
-        .mockImplementation((callback: (rows: any[]) => any) =>
-          Promise.resolve(callback(mockDbState.selectResults.shift() ?? []))
-        )
-      return chain
-    }),
-  },
-}))
-
-vi.mock('@sim/db/schema', () => ({
-  member: {
-    organizationId: 'member.organizationId',
-    role: 'member.role',
-    userId: 'member.userId',
-  },
-  organization: {
-    id: 'organization.id',
-    name: 'organization.name',
-  },
-  subscription: {
-    id: 'subscription.id',
-    plan: 'subscription.plan',
-    referenceId: 'subscription.referenceId',
-    status: 'subscription.status',
-    seats: 'subscription.seats',
-  },
-}))
-
-vi.mock('drizzle-orm', () => ({
-  and: vi.fn((...conditions: unknown[]) => ({ type: 'and', conditions })),
-  eq: vi.fn((field: unknown, value: unknown) => ({ field, value })),
-  inArray: vi.fn((field: unknown, value: unknown[]) => ({ field, value })),
-  or: vi.fn((...conditions: unknown[]) => ({ type: 'or', conditions })),
-}))
+vi.mock('@sim/db', () => dbChainMock)
 
 vi.mock('@sim/logger', () => loggerMock)
 
@@ -106,10 +67,12 @@ vi.mock('@/lib/workspaces/organization-workspaces', () => ({
 
 import { POST } from '@/app/api/organizations/route'
 
+afterAll(resetDbChainMock)
+
 describe('POST /api/organizations', () => {
   beforeEach(() => {
     vi.clearAllMocks()
-    mockDbState.selectResults = []
+    resetDbChainMock()
   })
 
   it('recovers an owner org when the subscription was already moved onto the organization', async () => {
@@ -120,10 +83,10 @@ describe('POST /api/organizations', () => {
         name: 'Owner',
       })
     )
-    mockDbState.selectResults = [
-      [{ organizationId: 'legacy-org-id', role: 'owner' }],
-      [{ id: 'sub-1', plan: 'team', referenceId: 'legacy-org-id', status: 'active', seats: 5 }],
-    ]
+    queueTableRows(member, [{ organizationId: 'legacy-org-id', role: 'owner' }])
+    queueTableRows(subscription, [
+      { id: 'sub-1', plan: 'team', referenceId: 'legacy-org-id', status: 'active', seats: 5 },
+    ])
 
     const response = await POST(
       new Request('http://localhost/api/organizations', {
@@ -164,10 +127,10 @@ describe('POST /api/organizations', () => {
       status: 'active',
       seats: 5,
     })
-    mockDbState.selectResults = [
-      [{ organizationId: 'legacy-org-id', role: 'owner' }],
-      [{ id: 'sub-1', plan: 'team', referenceId: 'user-1', status: 'active', seats: 5 }],
-    ]
+    queueTableRows(member, [{ organizationId: 'legacy-org-id', role: 'owner' }])
+    queueTableRows(subscription, [
+      { id: 'sub-1', plan: 'team', referenceId: 'user-1', status: 'active', seats: 5 },
+    ])
 
     const response = await POST(
       new Request('http://localhost/api/organizations', {
@@ -202,7 +165,7 @@ describe('POST /api/organizations', () => {
         name: 'Member',
       })
     )
-    mockDbState.selectResults = [[{ organizationId: 'org-1', role: 'member' }]]
+    queueTableRows(member, [{ organizationId: 'org-1', role: 'member' }])
 
     const response = await POST(
       new Request('http://localhost/api/organizations', {
@@ -230,10 +193,10 @@ describe('POST /api/organizations', () => {
         name: 'Owner',
       })
     )
-    mockDbState.selectResults = [
-      [{ organizationId: 'legacy-org-id', role: 'owner' }],
-      [{ id: 'sub-1', plan: 'team', referenceId: 'legacy-org-id', status: 'active', seats: 5 }],
-    ]
+    queueTableRows(member, [{ organizationId: 'legacy-org-id', role: 'owner' }])
+    queueTableRows(subscription, [
+      { id: 'sub-1', plan: 'team', referenceId: 'legacy-org-id', status: 'active', seats: 5 },
+    ])
     mockAttachOwnedWorkspacesToOrganization.mockRejectedValueOnce(
       new WorkspaceOrganizationMembershipConflictError([
         { userId: 'user-2', organizationId: 'org-2' },
diff --git a/apps/sim/app/api/workflows/[id]/route.test.ts b/apps/sim/app/api/workflows/[id]/route.test.ts
index 2315712fdb5..56b0ffd8b11 100644
--- a/apps/sim/app/api/workflows/[id]/route.test.ts
+++ b/apps/sim/app/api/workflows/[id]/route.test.ts
@@ -7,8 +7,11 @@
 
 import {
   auditMock,
+  dbChainMock,
+  dbChainMockFns,
   envMock,
   hybridAuthMockFns,
+  resetDbChainMock,
   telemetryMock,
   workflowAuthzMockFns,
   workflowsOrchestrationMock,
@@ -19,7 +22,7 @@ import {
   workflowsUtilsMockFns,
 } 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'
 import { getWorkflowResponseDataSchema } from '@/lib/api/contracts/workflows'
 
 const mockLoadWorkflowFromNormalizedTables =
@@ -30,12 +33,6 @@ const mockAuthorizeWorkflowByWorkspacePermission =
 const mockPerformDeleteWorkflow = workflowsOrchestrationMockFns.mockPerformDeleteWorkflow
 const mockPerformUpdateWorkflow = workflowsOrchestrationMockFns.mockPerformUpdateWorkflow
 
-const { mockDbUpdate, mockDbSelect, mockDbTransaction } = vi.hoisted(() => ({
-  mockDbUpdate: vi.fn(),
-  mockDbSelect: vi.fn(),
-  mockDbTransaction: vi.fn(),
-}))
-
 /**
  * Helper to set mock auth state consistently across getSession and hybrid auth.
  */
@@ -67,20 +64,18 @@ vi.mock('@/lib/workflows/utils', () => workflowsUtilsMock)
 
 vi.mock('@/lib/workflows/orchestration', () => workflowsOrchestrationMock)
 
-vi.mock('@sim/db', () => ({
-  db: {
-    update: () => mockDbUpdate(),
-    select: () => mockDbSelect(),
-    transaction: mockDbTransaction,
-  },
-  workflow: {},
-}))
+vi.mock('@sim/db', () => dbChainMock)
 
 import { DELETE, GET, PUT } from './route'
 
 describe('Workflow By ID API Route', () => {
+  afterAll(() => {
+    resetDbChainMock()
+  })
+
   beforeEach(() => {
     vi.clearAllMocks()
+    resetDbChainMock()
 
     vi.stubGlobal('crypto', {
       randomUUID: vi.fn().mockReturnValue('mock-request-id-12345678'),
@@ -103,18 +98,6 @@ describe('Workflow By ID API Route', () => {
         archivedAt: null,
       },
     }))
-    mockDbTransaction.mockImplementation(async (callback) =>
-      callback({
-        execute: vi.fn().mockResolvedValue(undefined),
-        select: vi.fn().mockReturnValue({
-          from: vi.fn().mockReturnValue({
-            where: vi.fn().mockReturnValue({
-              limit: vi.fn().mockResolvedValue([]),
-            }),
-          }),
-        }),
-      })
-    )
   })
 
   describe('GET /api/workflows/[id]', () => {
@@ -515,16 +498,6 @@ describe('Workflow By ID API Route', () => {
   })
 
   describe('PUT /api/workflows/[id]', () => {
-    function mockDuplicateCheck(results: Array<{ id: string }> = []) {
-      mockDbSelect.mockReturnValue({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockReturnValue({
-            limit: vi.fn().mockResolvedValue(results),
-          }),
-        }),
-      })
-    }
-
     it('should allow user with write permission to update workflow', async () => {
       const mockWorkflow = {
         id: 'workflow-123',
@@ -534,8 +507,6 @@ describe('Workflow By ID API Route', () => {
       }
 
       const updateData = { name: 'Updated Workflow' }
-      const updatedWorkflow = { ...mockWorkflow, ...updateData, updatedAt: new Date() }
-
       mockGetSession({ user: { id: 'user-123' } })
 
       mockGetWorkflowById.mockResolvedValue(mockWorkflow)
@@ -546,16 +517,6 @@ describe('Workflow By ID API Route', () => {
         workspacePermission: 'write',
       })
 
-      mockDuplicateCheck([])
-
-      mockDbUpdate.mockReturnValue({
-        set: vi.fn().mockReturnValue({
-          where: vi.fn().mockReturnValue({
-            returning: vi.fn().mockResolvedValue([updatedWorkflow]),
-          }),
-        }),
-      })
-
       const req = new NextRequest('http://localhost:3000/api/workflows/workflow-123', {
         method: 'PUT',
         body: JSON.stringify(updateData),
@@ -578,8 +539,6 @@ describe('Workflow By ID API Route', () => {
       }
 
       const updateData = { name: 'Updated Workflow' }
-      const updatedWorkflow = { ...mockWorkflow, ...updateData, updatedAt: new Date() }
-
       mockGetSession({ user: { id: 'user-123' } })
 
       mockGetWorkflowById.mockResolvedValue(mockWorkflow)
@@ -590,16 +549,6 @@ describe('Workflow By ID API Route', () => {
         workspacePermission: 'write',
       })
 
-      mockDuplicateCheck([])
-
-      mockDbUpdate.mockReturnValue({
-        set: vi.fn().mockReturnValue({
-          where: vi.fn().mockReturnValue({
-            returning: vi.fn().mockResolvedValue([updatedWorkflow]),
-          }),
-        }),
-      })
-
       const req = new NextRequest('http://localhost:3000/api/workflows/workflow-123', {
         method: 'PUT',
         body: JSON.stringify(updateData),
@@ -761,8 +710,6 @@ describe('Workflow By ID API Route', () => {
         workspaceId: 'workspace-456',
       }
 
-      const updatedWorkflow = { ...mockWorkflow, name: 'Unique Name', updatedAt: new Date() }
-
       mockGetSession({ user: { id: 'user-123' } })
       mockGetWorkflowById.mockResolvedValue(mockWorkflow)
       mockAuthorizeWorkflowByWorkspacePermission.mockResolvedValue({
@@ -772,16 +719,6 @@ describe('Workflow By ID API Route', () => {
         workspacePermission: 'write',
       })
 
-      mockDuplicateCheck([])
-
-      mockDbUpdate.mockReturnValue({
-        set: vi.fn().mockReturnValue({
-          where: vi.fn().mockReturnValue({
-            returning: vi.fn().mockResolvedValue([updatedWorkflow]),
-          }),
-        }),
-      })
-
       const req = new NextRequest('http://localhost:3000/api/workflows/workflow-123', {
         method: 'PUT',
         body: JSON.stringify({ name: 'Unique Name' }),
@@ -804,8 +741,6 @@ describe('Workflow By ID API Route', () => {
         workspaceId: 'workspace-456',
       }
 
-      const updatedWorkflow = { ...mockWorkflow, folderId: 'folder-2', updatedAt: new Date() }
-
       mockGetSession({ user: { id: 'user-123' } })
       mockGetWorkflowById.mockResolvedValue(mockWorkflow)
       mockAuthorizeWorkflowByWorkspacePermission.mockResolvedValue({
@@ -815,17 +750,6 @@ describe('Workflow By ID API Route', () => {
         workspacePermission: 'write',
       })
 
-      // No duplicate in target folder
-      mockDuplicateCheck([])
-
-      mockDbUpdate.mockReturnValue({
-        set: vi.fn().mockReturnValue({
-          where: vi.fn().mockReturnValue({
-            returning: vi.fn().mockResolvedValue([updatedWorkflow]),
-          }),
-        }),
-      })
-
       const req = new NextRequest('http://localhost:3000/api/workflows/workflow-123', {
         method: 'PUT',
         body: JSON.stringify({ folderId: 'folder-2' }),
@@ -883,12 +807,6 @@ describe('Workflow By ID API Route', () => {
         workspaceId: 'workspace-456',
       }
 
-      const updatedWorkflow = {
-        ...mockWorkflow,
-        description: 'Updated description',
-        updatedAt: new Date(),
-      }
-
       mockGetSession({ user: { id: 'user-123' } })
       mockGetWorkflowById.mockResolvedValue(mockWorkflow)
       mockAuthorizeWorkflowByWorkspacePermission.mockResolvedValue({
@@ -898,14 +816,6 @@ describe('Workflow By ID API Route', () => {
         workspacePermission: 'write',
       })
 
-      mockDbUpdate.mockReturnValue({
-        set: vi.fn().mockReturnValue({
-          where: vi.fn().mockReturnValue({
-            returning: vi.fn().mockResolvedValue([updatedWorkflow]),
-          }),
-        }),
-      })
-
       const req = new NextRequest('http://localhost:3000/api/workflows/workflow-123', {
         method: 'PUT',
         body: JSON.stringify({ description: 'Updated description' }),
@@ -915,8 +825,7 @@ describe('Workflow By ID API Route', () => {
       const response = await PUT(req, { params })
 
       expect(response.status).toBe(200)
-      // db.select should NOT have been called since no name/folder change
-      expect(mockDbSelect).not.toHaveBeenCalled()
+      expect(dbChainMockFns.select).not.toHaveBeenCalled()
     })
 
     it('should deny forkSyncExcluded update for non-admin users', async () => {
diff --git a/apps/sim/app/api/workflows/route.test.ts b/apps/sim/app/api/workflows/route.test.ts
index 261b8bf4b5d..732d4931200 100644
--- a/apps/sim/app/api/workflows/route.test.ts
+++ b/apps/sim/app/api/workflows/route.test.ts
@@ -4,43 +4,28 @@
 import {
   auditMock,
   createMockRequest,
+  dbChainMock,
+  dbChainMockFns,
   hybridAuthMockFns,
   permissionsMock,
   permissionsMockFns,
+  queueTableRows,
+  resetDbChainMock,
+  schemaMock,
   workflowAuthzMockFns,
   workflowsApiUtilsMock,
   workflowsPersistenceUtilsMock,
   workflowsPersistenceUtilsMockFns,
 } from '@sim/testing'
-import { drizzleOrmMock } from '@sim/testing/mocks'
-import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
-const { mockWorkflowCreated, mockDbSelect, mockDbInsert } = vi.hoisted(() => ({
+const { mockWorkflowCreated } = vi.hoisted(() => ({
   mockWorkflowCreated: vi.fn(),
-  mockDbSelect: vi.fn(),
-  mockDbInsert: vi.fn(),
 }))
 
 const mockGetUserEntityPermissions = permissionsMockFns.mockGetUserEntityPermissions
 
-vi.mock('drizzle-orm', () => ({
-  ...drizzleOrmMock,
-  min: vi.fn((field) => ({ type: 'min', field })),
-}))
-
-vi.mock('@sim/db', () => ({
-  db: {
-    select: (...args: unknown[]) => mockDbSelect(...args),
-    insert: (...args: unknown[]) => mockDbInsert(...args),
-    transaction: vi.fn(async (fn: (tx: Record<string, unknown>) => Promise<void>) => {
-      const tx = {
-        select: (...args: unknown[]) => mockDbSelect(...args),
-        insert: (...args: unknown[]) => mockDbInsert(...args),
-      }
-      await fn(tx)
-    }),
-  },
-}))
+vi.mock('@sim/db', () => dbChainMock)
 
 vi.mock('@sim/audit', () => auditMock)
 
@@ -67,8 +52,13 @@ vi.mock('@/lib/workflows/persistence/utils', () => workflowsPersistenceUtilsMock
 import { POST } from '@/app/api/workflows/route'
 
 describe('Workflows API Route - POST ordering', () => {
+  afterAll(() => {
+    resetDbChainMock()
+  })
+
   beforeEach(() => {
     vi.clearAllMocks()
+    resetDbChainMock()
 
     vi.stubGlobal('crypto', {
       randomUUID: vi.fn().mockReturnValue('workflow-new-id'),
@@ -102,33 +92,13 @@ describe('Workflows API Route - POST ordering', () => {
 
     const response = await POST(req)
     expect(response.status).toBe(423)
-    expect(mockDbInsert).not.toHaveBeenCalled()
+    expect(dbChainMockFns.insert).not.toHaveBeenCalled()
   })
 
   it('uses top insertion against mixed siblings (folders + workflows)', async () => {
-    const minResultsQueue: Array<Array<{ minOrder: number }>> = [
-      [],
-      [{ minOrder: 5 }],
-      [{ minOrder: 2 }],
-    ]
-
-    mockDbSelect.mockImplementation(() => ({
-      from: vi.fn().mockReturnValue({
-        where: vi.fn().mockImplementation(() => ({
-          limit: vi.fn().mockImplementation(() => Promise.resolve(minResultsQueue.shift() ?? [])),
-          then: (onFulfilled: (value: Array<{ minOrder: number }>) => unknown) =>
-            Promise.resolve(minResultsQueue.shift() ?? []).then(onFulfilled),
-        })),
-      }),
-    }))
-
-    let insertedValues: Record<string, unknown> | null = null
-    mockDbInsert.mockReturnValue({
-      values: vi.fn().mockImplementation((values: Record<string, unknown>) => {
-        insertedValues = values
-        return Promise.resolve(undefined)
-      }),
-    })
+    queueTableRows(schemaMock.workflow, [])
+    queueTableRows(schemaMock.workflow, [{ minOrder: 5 }])
+    queueTableRows(schemaMock.workflowFolder, [{ minOrder: 2 }])
 
     const req = createMockRequest('POST', {
       name: 'New Workflow',
@@ -141,31 +111,10 @@ describe('Workflows API Route - POST ordering', () => {
     const data = await response.json()
     expect(response.status).toBe(200)
     expect(data.sortOrder).toBe(1)
-    expect(insertedValues).not.toBeNull()
-    expect(insertedValues?.sortOrder).toBe(1)
+    expect(dbChainMockFns.values).toHaveBeenCalledWith(expect.objectContaining({ sortOrder: 1 }))
   })
 
   it('defaults to sortOrder 0 when there are no siblings', async () => {
-    const minResultsQueue: Array<Array<{ minOrder: number }>> = [[], [], []]
-
-    mockDbSelect.mockImplementation(() => ({
-      from: vi.fn().mockReturnValue({
-        where: vi.fn().mockImplementation(() => ({
-          limit: vi.fn().mockImplementation(() => Promise.resolve(minResultsQueue.shift() ?? [])),
-          then: (onFulfilled: (value: Array<{ minOrder: number }>) => unknown) =>
-            Promise.resolve(minResultsQueue.shift() ?? []).then(onFulfilled),
-        })),
-      }),
-    }))
-
-    let insertedValues: Record<string, unknown> | null = null
-    mockDbInsert.mockReturnValue({
-      values: vi.fn().mockImplementation((values: Record<string, unknown>) => {
-        insertedValues = values
-        return Promise.resolve(undefined)
-      }),
-    })
-
     const req = createMockRequest('POST', {
       name: 'New Workflow',
       description: 'desc',
@@ -177,6 +126,6 @@ describe('Workflows API Route - POST ordering', () => {
     const data = await response.json()
     expect(response.status).toBe(200)
     expect(data.sortOrder).toBe(0)
-    expect(insertedValues?.sortOrder).toBe(0)
+    expect(dbChainMockFns.values).toHaveBeenCalledWith(expect.objectContaining({ sortOrder: 0 }))
   })
 })
diff --git a/apps/sim/background/cleanup-logs.test.ts b/apps/sim/background/cleanup-logs.test.ts
index 9df5ec1ab77..1b0f77fce94 100644
--- a/apps/sim/background/cleanup-logs.test.ts
+++ b/apps/sim/background/cleanup-logs.test.ts
@@ -2,7 +2,9 @@
  * @vitest-environment node
  */
 
-import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { dbChainMock, dbChainMockFns, resetDbChainMock, schemaMock } from '@sim/testing'
+import { drizzleOrmMock } from '@sim/testing/mocks'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
 interface CleanupRow {
   id: string
@@ -18,143 +20,34 @@ interface CapturedBatchDeleteOptions {
 }
 
 const {
-  mockAnd,
   mockBatchDeleteByWorkspaceAndTimestamp,
   mockChunkedBatchDelete,
   mockDeleteFileMetadata,
   mockDeleteFiles,
-  mockEq,
-  mockExecute,
-  mockFrom,
-  mockInArray,
-  mockIsNull,
-  mockLeftJoin,
-  mockLimit,
-  mockLt,
   mockMarkLargeValuesDeleted,
-  mockNotInArray,
-  mockOr,
-  mockOrderBy,
   mockPruneLargeValueMetadata,
-  mockSelect,
   mockTask,
-  mockWhere,
-} = vi.hoisted(() => {
-  const mockLimit = vi.fn(async () => [])
-  const mockOrderBy = vi.fn(() => ({ limit: mockLimit }))
-  const mockWhere = vi.fn(() => ({ limit: mockLimit, orderBy: mockOrderBy }))
-  const mockLeftJoin = vi.fn(() => ({ where: mockWhere }))
-  const mockFrom = vi.fn(() => ({ leftJoin: mockLeftJoin, where: mockWhere }))
-  const mockSelect = vi.fn(() => ({ from: mockFrom }))
-
-  return {
-    mockAnd: vi.fn((...args: unknown[]) => ({ op: 'and', args })),
-    mockBatchDeleteByWorkspaceAndTimestamp: vi.fn(async () => ({
-      table: 'job',
-      deleted: 0,
-      failed: 0,
-    })),
-    mockChunkedBatchDelete: vi.fn(),
-    mockDeleteFileMetadata: vi.fn(async () => true),
-    mockDeleteFiles: vi.fn(async () => ({ deleted: 2, failed: [] })),
-    mockEq: vi.fn((...args: unknown[]) => ({ op: 'eq', args })),
-    mockExecute: vi.fn(),
-    mockFrom,
-    mockInArray: vi.fn((...args: unknown[]) => ({ op: 'inArray', args })),
-    mockIsNull: vi.fn((...args: unknown[]) => ({ op: 'isNull', args })),
-    mockLeftJoin,
-    mockLimit,
-    mockLt: vi.fn((...args: unknown[]) => ({ op: 'lt', args })),
-    mockMarkLargeValuesDeleted: vi.fn(async () => undefined),
-    mockNotInArray: vi.fn((...args: unknown[]) => ({ op: 'notInArray', args })),
-    mockOr: vi.fn((...args: unknown[]) => ({ op: 'or', args })),
-    mockOrderBy,
-    mockPruneLargeValueMetadata: vi.fn(async () => ({
-      referencesDeleted: 0,
-      dependenciesDeleted: 0,
-      tombstonesDeleted: 0,
-    })),
-    mockSelect,
-    mockTask: vi.fn((config: unknown) => config),
-    mockWhere,
-  }
-})
-
-vi.mock('@sim/db', () => {
-  const db = {
-    execute: mockExecute,
-    select: mockSelect,
-  }
-  return {
-    db,
-    // Cleanup-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',
-  },
-  executionLargeValues: {
-    createdAt: 'executionLargeValues.createdAt',
-    deletedAt: 'executionLargeValues.deletedAt',
-    key: 'executionLargeValues.key',
-    workspaceId: 'executionLargeValues.workspaceId',
-  },
-  jobExecutionLogs: {
-    startedAt: 'jobExecutionLogs.startedAt',
-    workspaceId: 'jobExecutionLogs.workspaceId',
-  },
-  pausedExecutions: {
-    executionId: 'pausedExecutions.executionId',
-    status: 'pausedExecutions.status',
-  },
-  workspaceFiles: {
-    context: 'workspaceFiles.context',
-    deletedAt: 'workspaceFiles.deletedAt',
-    key: 'workspaceFiles.key',
-    uploadedAt: 'workspaceFiles.uploadedAt',
-    workspaceId: 'workspaceFiles.workspaceId',
-  },
-  workflowExecutionLogs: {
-    executionData: 'workflowExecutionLogs.executionData',
-    executionId: 'workflowExecutionLogs.executionId',
-    files: 'workflowExecutionLogs.files',
-    id: 'workflowExecutionLogs.id',
-    startedAt: 'workflowExecutionLogs.startedAt',
-    workspaceId: 'workflowExecutionLogs.workspaceId',
-  },
-}))
-
-vi.mock('@sim/logger', () => ({
-  createLogger: vi.fn(() => ({
-    error: vi.fn(),
-    info: vi.fn(),
-    warn: vi.fn(),
+} = vi.hoisted(() => ({
+  mockBatchDeleteByWorkspaceAndTimestamp: vi.fn(async () => ({
+    table: 'job',
+    deleted: 0,
+    failed: 0,
+  })),
+  mockChunkedBatchDelete: vi.fn(),
+  mockDeleteFileMetadata: vi.fn(async () => true),
+  mockDeleteFiles: vi.fn(async () => ({ deleted: 2, failed: [] })),
+  mockMarkLargeValuesDeleted: vi.fn(async () => undefined),
+  mockPruneLargeValueMetadata: vi.fn(async () => ({
+    referencesDeleted: 0,
+    dependenciesDeleted: 0,
+    tombstonesDeleted: 0,
   })),
+  mockTask: vi.fn((config: unknown) => config),
 }))
 
-vi.mock('@trigger.dev/sdk', () => ({ task: mockTask }))
+vi.mock('@sim/db', () => dbChainMock)
 
-vi.mock('drizzle-orm', () => ({
-  and: mockAnd,
-  asc: vi.fn((column: unknown) => ({ op: 'asc', column })),
-  eq: mockEq,
-  inArray: mockInArray,
-  isNull: mockIsNull,
-  lt: mockLt,
-  notInArray: mockNotInArray,
-  or: mockOr,
-  sql: vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ strings, values })),
-}))
+vi.mock('@trigger.dev/sdk', () => ({ task: mockTask }))
 
 vi.mock('@/lib/cleanup/batch-delete', () => ({
   batchDeleteByWorkspaceAndTimestamp: mockBatchDeleteByWorkspaceAndTimestamp,
@@ -195,8 +88,13 @@ vi.mock('@/lib/uploads/server/metadata', () => ({
 import { cleanupLogsTask, runCleanupLogs } from '@/background/cleanup-logs'
 
 describe('cleanup logs worker', () => {
+  afterAll(() => {
+    resetDbChainMock()
+  })
+
   beforeEach(() => {
     vi.clearAllMocks()
+    resetDbChainMock()
     mockChunkedBatchDelete.mockImplementation(async (options: CapturedBatchDeleteOptions) => {
       await options.selectChunk(['workspace-1'], 500)
       await options.onBatch?.([
@@ -228,11 +126,11 @@ describe('cleanup logs worker', () => {
         totalRowLimit: 25_000,
       })
     )
-    expect(mockSelect).toHaveBeenCalledWith({
-      id: 'workflowExecutionLogs.id',
-      files: 'workflowExecutionLogs.files',
+    expect(dbChainMockFns.select).toHaveBeenCalledWith({
+      id: schemaMock.workflowExecutionLogs.id,
+      files: schemaMock.workflowExecutionLogs.files,
     })
-    expect(mockExecute).not.toHaveBeenCalled()
+    expect(dbChainMockFns.execute).not.toHaveBeenCalled()
     expect(mockDeleteFiles).toHaveBeenCalledWith(
       ['execution-file-a', 'execution-file-b'],
       'execution'
@@ -247,7 +145,7 @@ describe('cleanup logs worker', () => {
   it('does not count large values as deleted when deleted_at marking fails', async () => {
     const largeValueKey =
       'execution/workspace-1/workflow-1/execution-1/large-value-lv_abcdefghijkl.json'
-    mockLimit.mockResolvedValueOnce([]).mockResolvedValueOnce([{ key: largeValueKey }])
+    dbChainMockFns.limit.mockResolvedValueOnce([]).mockResolvedValueOnce([{ key: largeValueKey }])
     mockDeleteFiles
       .mockResolvedValueOnce({ deleted: 2, failed: [] })
       .mockResolvedValueOnce({ deleted: 1, failed: [] })
@@ -267,7 +165,7 @@ describe('cleanup logs worker', () => {
   it('cleans legacy large values from file metadata without selecting execution_data', async () => {
     const legacyKey =
       'execution/workspace-1/workflow-1/execution-1/large-value-lv_abcdefghijkl.json'
-    mockLimit
+    dbChainMockFns.limit
       .mockResolvedValueOnce([])
       .mockResolvedValueOnce([])
       .mockResolvedValueOnce([{ key: legacyKey }])
@@ -282,14 +180,14 @@ describe('cleanup logs worker', () => {
       workspaceIds: ['workspace-1'],
     })
 
-    expect(mockSelect).toHaveBeenCalledWith({
-      id: 'workflowExecutionLogs.id',
-      files: 'workflowExecutionLogs.files',
+    expect(dbChainMockFns.select).toHaveBeenCalledWith({
+      id: schemaMock.workflowExecutionLogs.id,
+      files: schemaMock.workflowExecutionLogs.files,
     })
-    expect(mockSelect).not.toHaveBeenCalledWith(
+    expect(dbChainMockFns.select).not.toHaveBeenCalledWith(
       expect.objectContaining({ executionData: expect.anything() })
     )
-    const legacyWhereArgs = mockAnd.mock.calls
+    const legacyWhereArgs = drizzleOrmMock.and.mock.calls
       .flat()
       .filter((arg): arg is { strings: string[] } => {
         return (
diff --git a/apps/sim/background/cleanup-soft-deletes.test.ts b/apps/sim/background/cleanup-soft-deletes.test.ts
index 7c0e29603bf..ea63df6a859 100644
--- a/apps/sim/background/cleanup-soft-deletes.test.ts
+++ b/apps/sim/background/cleanup-soft-deletes.test.ts
@@ -2,119 +2,38 @@
  * @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 {
   mockBatchDeleteByWorkspaceAndTimestamp,
   mockChunkedBatchDelete,
   mockDecrementStorageUsageForBillingContextInTx,
-  mockDelete,
-  mockDeleteReturning,
-  mockDeleteWhere,
   mockDeleteFileMetadata,
   mockDeleteFiles,
   mockDeleteRowsById,
   mockHardDeleteDocuments,
   mockIsUsingCloudStorage,
   mockKnowledgeBaseContainerDelete,
-  mockLimit,
-  mockOrderBy,
   mockPrepareChatCleanup,
   mockResolveStorageBillingContext,
-  mockSelect,
   mockSelectRowsByIdChunks,
-  mockTask,
-  mockTransaction,
-  mockWhere,
-} = vi.hoisted(() => {
-  const mockLimit = vi.fn(async () => [] as Array<{ key: string }>)
-  const mockOrderBy = vi.fn(() => ({ limit: mockLimit }))
-  const mockWhere = vi.fn(() => ({ orderBy: mockOrderBy, limit: mockLimit }))
-  const mockFrom = vi.fn(() => ({
-    where: mockWhere,
-    leftJoin: vi.fn(() => ({ where: mockWhere })),
-  }))
-  const mockSelect = vi.fn(() => ({ from: mockFrom }))
-  const mockDeleteReturning = vi.fn(async () => [] as Array<{ id: string; size?: number }>)
-  const mockDeleteWhere = vi.fn(() => ({ returning: mockDeleteReturning }))
-  const mockDelete = vi.fn(() => ({ where: mockDeleteWhere }))
-  const mockKnowledgeBaseContainerDelete = vi.fn()
-  const mockChunkedBatchDelete = vi.fn(async () => ({ deleted: 0, failed: 0 }))
-
-  return {
-    mockBatchDeleteByWorkspaceAndTimestamp: vi.fn(async () => ({ deleted: 0, failed: 0 })),
-    mockChunkedBatchDelete,
-    mockDecrementStorageUsageForBillingContextInTx: vi.fn(async () => undefined),
-    mockDelete,
-    mockDeleteReturning,
-    mockDeleteWhere,
-    mockDeleteFileMetadata: vi.fn(async () => true),
-    mockDeleteFiles: vi.fn(async () => ({ deleted: 0, failed: [] as Array<{ key: string }> })),
-    mockDeleteRowsById: vi.fn(async () => ({ deleted: 0, failed: 0 })),
-    mockHardDeleteDocuments: vi.fn(async (ids: string[]) => ids.length),
-    mockIsUsingCloudStorage: vi.fn(() => true),
-    mockKnowledgeBaseContainerDelete,
-    mockLimit,
-    mockOrderBy,
-    mockPrepareChatCleanup: vi.fn(async () => ({ execute: vi.fn(async () => undefined) })),
-    mockResolveStorageBillingContext: vi.fn(),
-    mockSelect,
-    mockSelectRowsByIdChunks: vi.fn(async () => [] as unknown[]),
-    mockTask: vi.fn((config: unknown) => config),
-    mockTransaction: vi.fn(),
-    mockWhere,
-  }
-})
-
-vi.mock('@sim/db', () => {
-  const db = {
-    delete: mockDelete,
-    select: mockSelect,
-    transaction: mockTransaction,
-  }
-  return {
-    db,
-    // Cleanup-pool client shares the instance so the seeded chains still apply.
-    dbFor: () => db,
-  }
-})
-
-vi.mock('@sim/db/schema', () => {
-  const table = (cols: string[]) =>
-    Object.fromEntries(cols.map((c) => [c, `col.${c}`])) as Record<string, string>
-  const wsFileCols = ['id', 'key', 'context', 'size', 'workspaceId', 'deletedAt', 'uploadedAt']
-  const softCols = ['id', 'archivedAt', 'deletedAt', 'workspaceId']
-  return {
-    copilotChats: table(['id', 'workflowId']),
-    document: table(['id', 'storageKey', 'knowledgeBaseId']),
-    knowledgeBase: table(softCols),
-    mcpServers: table(softCols),
-    memory: table(softCols),
-    userTableDefinitions: table(softCols),
-    workflow: table(softCols),
-    workflowFolder: table(softCols),
-    workflowMcpServer: table(softCols),
-    workspaceFile: table(wsFileCols),
-    workspaceFiles: table(wsFileCols),
-  }
-})
-
-vi.mock('@sim/logger', () => ({
-  createLogger: vi.fn(() => ({ error: vi.fn(), info: vi.fn(), warn: vi.fn() })),
+} = vi.hoisted(() => ({
+  mockBatchDeleteByWorkspaceAndTimestamp: vi.fn(async () => ({ deleted: 0, failed: 0 })),
+  mockChunkedBatchDelete: vi.fn(async () => ({ deleted: 0, failed: 0 })),
+  mockDecrementStorageUsageForBillingContextInTx: vi.fn(async () => undefined),
+  mockDeleteFileMetadata: vi.fn(async () => true),
+  mockDeleteFiles: vi.fn(async () => ({ deleted: 0, failed: [] as Array<{ key: string }> })),
+  mockDeleteRowsById: vi.fn(async () => ({ deleted: 0, failed: 0 })),
+  mockHardDeleteDocuments: vi.fn(async (ids: string[]) => ids.length),
+  mockIsUsingCloudStorage: vi.fn(() => true),
+  mockKnowledgeBaseContainerDelete: vi.fn(),
+  mockPrepareChatCleanup: vi.fn(async () => ({ execute: vi.fn(async () => undefined) })),
+  mockResolveStorageBillingContext: vi.fn(),
+  mockSelectRowsByIdChunks: vi.fn(async () => [] as unknown[]),
 }))
 
-vi.mock('@trigger.dev/sdk', () => ({ task: mockTask }))
-
-vi.mock('drizzle-orm', () => ({
-  and: vi.fn((...args: unknown[]) => ({ op: 'and', args })),
-  asc: vi.fn((column: unknown) => ({ op: 'asc', column })),
-  eq: vi.fn((...args: unknown[]) => ({ op: 'eq', args })),
-  inArray: vi.fn((...args: unknown[]) => ({ op: 'inArray', args })),
-  isNotNull: vi.fn((...args: unknown[]) => ({ op: 'isNotNull', args })),
-  isNull: vi.fn((...args: unknown[]) => ({ op: 'isNull', args })),
-  lt: vi.fn((...args: unknown[]) => ({ op: 'lt', args })),
-  sql: vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ strings, values })),
-}))
+vi.mock('@sim/db', () => dbChainMock)
 
 vi.mock('@/lib/cleanup/batch-delete', () => ({
   batchDeleteByWorkspaceAndTimestamp: mockBatchDeleteByWorkspaceAndTimestamp,
@@ -159,14 +78,17 @@ const basePayload = {
 }
 
 describe('cleanup soft deletes', () => {
+  afterAll(() => {
+    resetDbChainMock()
+  })
+
   beforeEach(() => {
     vi.clearAllMocks()
+    resetDbChainMock()
     mockIsUsingCloudStorage.mockReturnValue(true)
-    mockLimit.mockReset().mockResolvedValue([])
     mockSelectRowsByIdChunks.mockReset().mockResolvedValue([])
     mockDeleteFiles.mockReset().mockResolvedValue({ deleted: 0, failed: [] })
     mockChunkedBatchDelete.mockReset().mockResolvedValue({ deleted: 0, failed: 0 })
-    mockDeleteReturning.mockReset().mockResolvedValue([])
     mockResolveStorageBillingContext.mockResolvedValue({
       workspaceId: 'ws-1',
       billedAccountUserId: 'user-1',
@@ -174,11 +96,6 @@ describe('cleanup soft deletes', () => {
       plan: 'free',
       customStorageLimitGB: null,
     })
-    mockTransaction
-      .mockReset()
-      .mockImplementation(async (callback: (tx: { delete: typeof mockDelete }) => unknown) =>
-        callback({ delete: mockDelete })
-      )
   })
 
   it('keeps metadata rows whose object deletion failed', async () => {
@@ -201,8 +118,8 @@ describe('cleanup soft deletes', () => {
 
     await runCleanupSoftDeletes(basePayload)
 
-    expect(mockTransaction).not.toHaveBeenCalled()
-    expect(mockDelete).not.toHaveBeenCalled()
+    expect(dbChainMockFns.transaction).not.toHaveBeenCalled()
+    expect(dbChainMockFns.delete).not.toHaveBeenCalled()
     expect(
       mockDeleteRowsById.mock.calls.some(([, , ids]) => (ids as string[]).includes('file-failed'))
     ).toBe(false)
@@ -229,18 +146,18 @@ describe('cleanup soft deletes', () => {
         },
       ])
     mockDeleteFiles.mockResolvedValueOnce({ deleted: 2, failed: [] })
-    mockDeleteReturning.mockResolvedValueOnce([{ id: 'file-deleted', size: 7 }])
+    dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'file-deleted', size: 7 }])
 
     await runCleanupSoftDeletes(basePayload)
 
     expect(mockResolveStorageBillingContext).toHaveBeenCalledOnce()
     expect(mockDecrementStorageUsageForBillingContextInTx).toHaveBeenCalledWith(
-      expect.objectContaining({ delete: mockDelete }),
+      dbChainMock.db,
       expect.objectContaining({ workspaceId: 'ws-1' }),
       7
     )
     expect(mockDeleteFiles.mock.invocationCallOrder[0]).toBeLessThan(
-      mockTransaction.mock.invocationCallOrder[0]
+      dbChainMockFns.transaction.mock.invocationCallOrder[0]
     )
   })
 
@@ -258,12 +175,12 @@ describe('cleanup soft deletes', () => {
         },
       ])
     mockDeleteFiles.mockResolvedValueOnce({ deleted: 1, failed: [] })
-    mockDeleteReturning.mockResolvedValueOnce([{ id: 'chat-file' }])
+    dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'chat-file' }])
 
     await runCleanupSoftDeletes(basePayload)
 
     expect(mockDeleteFiles).toHaveBeenCalledWith(['mothership/chat-file'], 'mothership')
-    expect(mockDelete).toHaveBeenCalled()
+    expect(dbChainMockFns.delete).toHaveBeenCalled()
     expect(mockResolveStorageBillingContext).not.toHaveBeenCalled()
     expect(mockDecrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled()
   })
@@ -280,7 +197,7 @@ describe('cleanup soft deletes', () => {
         return { deleted: 1, failed: 0 }
       }
     )
-    mockLimit
+    dbChainMockFns.limit
       .mockResolvedValueOnce([{ id: 'doc-1' }])
       .mockResolvedValueOnce([])
       .mockResolvedValueOnce([])
@@ -294,7 +211,7 @@ describe('cleanup soft deletes', () => {
   })
 
   it('soft-deletes abandoned KB bindings and removes their storage objects', async () => {
-    mockLimit
+    dbChainMockFns.limit
       .mockResolvedValueOnce([{ key: 'kb/orphan-1' }, { key: 'kb/orphan-2' }])
       .mockResolvedValueOnce([])
 
@@ -307,7 +224,7 @@ describe('cleanup soft deletes', () => {
   })
 
   it('keeps an orphan KB binding when its object deletion fails', async () => {
-    mockLimit.mockResolvedValueOnce([{ key: 'kb/orphan-retry' }])
+    dbChainMockFns.limit.mockResolvedValueOnce([{ key: 'kb/orphan-retry' }])
     mockDeleteFiles.mockResolvedValueOnce({
       deleted: 0,
       failed: [{ key: 'kb/orphan-retry', error: 'storage unavailable' }],
@@ -320,7 +237,7 @@ describe('cleanup soft deletes', () => {
 
   it('still removes bindings but skips object deletion without cloud storage', async () => {
     mockIsUsingCloudStorage.mockReturnValue(false)
-    mockLimit.mockResolvedValueOnce([{ key: 'kb/orphan-1' }]).mockResolvedValueOnce([])
+    dbChainMockFns.limit.mockResolvedValueOnce([{ key: 'kb/orphan-1' }]).mockResolvedValueOnce([])
 
     await runCleanupSoftDeletes(basePayload)
 
@@ -329,7 +246,7 @@ describe('cleanup soft deletes', () => {
   })
 
   it('stops the batch loop when binding deletion makes no progress', async () => {
-    mockLimit.mockResolvedValue([{ key: 'kb/stuck' }])
+    dbChainMockFns.limit.mockResolvedValue([{ key: 'kb/stuck' }])
     mockDeleteFileMetadata.mockRejectedValue(new Error('db down'))
 
     await runCleanupSoftDeletes(basePayload)
@@ -341,7 +258,7 @@ describe('cleanup soft deletes', () => {
   it('does not run the sweep when there are no workspaces', async () => {
     await runCleanupSoftDeletes({ ...basePayload, workspaceIds: [] })
 
-    expect(mockSelect).not.toHaveBeenCalled()
+    expect(dbChainMockFns.select).not.toHaveBeenCalled()
     expect(mockDeleteFiles).not.toHaveBeenCalled()
     expect(mockDeleteFileMetadata).not.toHaveBeenCalled()
   })
diff --git a/apps/sim/background/tiktok-webhook-targets.test.ts b/apps/sim/background/tiktok-webhook-targets.test.ts
index d8acfbf1299..76f7ed76cbc 100644
--- a/apps/sim/background/tiktok-webhook-targets.test.ts
+++ b/apps/sim/background/tiktok-webhook-targets.test.ts
@@ -2,51 +2,17 @@
  * @vitest-environment node
  */
 
-import { beforeEach, describe, expect, it, vi } from 'vitest'
-
-const {
-  mockAsc,
-  mockCredentialExpression,
-  mockEq,
-  mockGt,
-  mockLike,
-  mockLimit,
-  mockOrderBy,
-  mockSelect,
-  queryRows,
-} = vi.hoisted(() => ({
+import { dbChainMock, dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { mockAsc, mockCredentialExpression, mockEq, mockGt, mockLike, tables } = vi.hoisted(() => ({
   mockAsc: vi.fn((value: unknown) => ({ asc: value })),
   mockCredentialExpression: vi.fn(() => 'webhook.credentialId'),
   mockEq: vi.fn((left: unknown, right: unknown) => ({ left, right })),
   mockGt: vi.fn((left: unknown, right: unknown) => ({ gt: [left, right] })),
   mockLike: vi.fn((left: unknown, right: unknown) => ({ left, right })),
-  mockLimit: vi.fn(),
-  mockOrderBy: vi.fn(),
-  mockSelect: vi.fn(),
-  queryRows: {
-    rows: [] as Array<{
-      accountId: string
-      webhookId: string
-      webhook: Record<string, unknown>
-      workflow: Record<string, unknown>
-    }>,
-  },
-}))
-
-vi.mock('@sim/db', () => {
-  const chain = {
-    from: vi.fn(() => chain),
-    innerJoin: vi.fn(() => chain),
-    leftJoin: vi.fn(() => chain),
-    where: vi.fn(() => chain),
-    orderBy: mockOrderBy,
-    limit: mockLimit,
-  }
-  mockOrderBy.mockImplementation(() => chain)
-  mockLimit.mockImplementation((limit: number) => Promise.resolve(queryRows.rows.slice(0, limit)))
-  mockSelect.mockImplementation(() => chain)
-
-  return {
+  /** Table-qualified column names keep the eq/like assertions unambiguous. */
+  tables: {
     account: {
       id: 'account.id',
       accountId: 'account.accountId',
@@ -59,8 +25,6 @@ vi.mock('@sim/db', () => {
       type: 'credential.type',
       workspaceId: 'credential.workspaceId',
     },
-    db: { select: mockSelect },
-    webhookCredentialIdExpression: mockCredentialExpression,
     webhook: {
       deploymentVersionId: 'webhook.deploymentVersionId',
       isActive: 'webhook.isActive',
@@ -80,8 +44,14 @@ vi.mock('@sim/db', () => {
       workflowId: 'workflowDeploymentVersion.workflowId',
       isActive: 'workflowDeploymentVersion.isActive',
     },
-  }
-})
+  },
+}))
+
+vi.mock('@sim/db', () => ({
+  ...dbChainMock,
+  ...tables,
+  webhookCredentialIdExpression: mockCredentialExpression,
+}))
 
 vi.mock('drizzle-orm', () => ({
   and: vi.fn((...conditions: unknown[]) => conditions),
@@ -100,14 +70,30 @@ import {
 
 const ACCOUNT_UUID = '11111111-2222-3333-4444-555555555555'
 
+/** Queues one page of joined rows, clipped like the SQL LIMIT would. */
+function queuePageRows(
+  rows: Array<{
+    accountId: string
+    webhookId: string
+    webhook: Record<string, unknown>
+    workflow: Record<string, unknown>
+  }>
+) {
+  queueTableRows(tables.account, rows.slice(0, TIKTOK_WEBHOOK_TARGET_PAGE_SIZE))
+}
+
 describe('findTikTokWebhookTargetPage', () => {
+  afterAll(() => {
+    resetDbChainMock()
+  })
+
   beforeEach(() => {
     vi.clearAllMocks()
-    queryRows.rows = []
+    resetDbChainMock()
   })
 
   it('returns only rows whose stored account ID exactly matches user_openid', async () => {
-    queryRows.rows = [
+    queuePageRows([
       {
         accountId: `act.user-${ACCOUNT_UUID}`,
         webhookId: 'webhook-1',
@@ -120,7 +106,7 @@ describe('findTikTokWebhookTargetPage', () => {
         webhook: { id: 'webhook-2' },
         workflow: { id: 'workflow-2' },
       },
-    ]
+    ])
 
     const page = await findTikTokWebhookTargetPage('act.user', 'request-1')
 
@@ -151,20 +137,22 @@ describe('findTikTokWebhookTargetPage', () => {
 
     expect(mockGt).toHaveBeenCalledWith('webhook.id', 'webhook-100')
     expect(mockAsc).toHaveBeenCalledWith('webhook.id')
-    expect(mockOrderBy).toHaveBeenCalledWith({ asc: 'webhook.id' })
-    expect(mockLimit).toHaveBeenCalledWith(TIKTOK_WEBHOOK_TARGET_PAGE_SIZE)
+    expect(dbChainMockFns.orderBy).toHaveBeenCalledWith({ asc: 'webhook.id' })
+    expect(dbChainMockFns.limit).toHaveBeenCalledWith(TIKTOK_WEBHOOK_TARGET_PAGE_SIZE)
   })
 
   it('returns a continuation cursor when the fixed-size page is full', async () => {
-    queryRows.rows = Array.from({ length: TIKTOK_WEBHOOK_TARGET_PAGE_SIZE + 1 }, (_, index) => {
-      const webhookId = `webhook-${String(index).padStart(3, '0')}`
-      return {
-        accountId: `act.user-${ACCOUNT_UUID}`,
-        webhookId,
-        webhook: { id: webhookId },
-        workflow: { id: `workflow-${index}` },
-      }
-    })
+    queuePageRows(
+      Array.from({ length: TIKTOK_WEBHOOK_TARGET_PAGE_SIZE + 1 }, (_, index) => {
+        const webhookId = `webhook-${String(index).padStart(3, '0')}`
+        return {
+          accountId: `act.user-${ACCOUNT_UUID}`,
+          webhookId,
+          webhook: { id: webhookId },
+          workflow: { id: `workflow-${index}` },
+        }
+      })
+    )
 
     const page = await findTikTokWebhookTargetPage('act.user', 'request-4')
 
@@ -188,6 +176,6 @@ describe('findTikTokWebhookTargetPage', () => {
       nextCursor: null,
       targets: [],
     })
-    expect(mockSelect).not.toHaveBeenCalled()
+    expect(dbChainMockFns.select).not.toHaveBeenCalled()
   })
 })
diff --git a/apps/sim/lib/admin/dashboard-credit-grant.test.ts b/apps/sim/lib/admin/dashboard-credit-grant.test.ts
index 03dba00dbcf..57e9f4677f8 100644
--- a/apps/sim/lib/admin/dashboard-credit-grant.test.ts
+++ b/apps/sim/lib/admin/dashboard-credit-grant.test.ts
@@ -1,13 +1,12 @@
 /**
  * @vitest-environment node
  */
-import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { member, organization, subscription, user, userStats, workspace } from '@sim/db/schema'
+import { dbChainMock, dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
 const mocks = vi.hoisted(() => ({
-  rows: [] as unknown[][],
-  returningRows: [] as unknown[][],
   billingSubscriptions: [] as unknown[],
-  updateSets: [] as Record<string, unknown>[],
   idempotencyCalls: [] as { namespace: string; requestFingerprint: string }[],
   recordAudit: vi.fn(),
   acquireLock: vi.fn(),
@@ -25,43 +24,7 @@ vi.mock('@sim/audit', () => ({
   AuditResourceType: { BILLING: 'billing' },
   recordAudit: mocks.recordAudit,
 }))
-vi.mock('@sim/db', () => {
-  const selectChain = () => {
-    const chain: Record<string, unknown> = {}
-    chain.from = () => chain
-    chain.where = () => chain
-    chain.orderBy = () => chain
-    chain.for = () => chain
-    chain.limit = () => Promise.resolve(mocks.rows.shift() ?? [])
-    chain.then = (resolve: (value: unknown[]) => unknown) =>
-      Promise.resolve(mocks.rows.shift() ?? []).then(resolve)
-    return chain
-  }
-  const update = () => {
-    const chain: Record<string, unknown> = {}
-    chain.set = (values: Record<string, unknown>) => {
-      mocks.updateSets.push(values)
-      return chain
-    }
-    chain.where = () => chain
-    chain.returning = () => Promise.resolve(mocks.returningRows.shift() ?? [])
-    chain.then = (resolve: (value: unknown[]) => unknown) => Promise.resolve([]).then(resolve)
-    return chain
-  }
-  const insert = () => {
-    const chain: Record<string, unknown> = {}
-    chain.values = () => chain
-    chain.onConflictDoNothing = () => Promise.resolve([])
-    return chain
-  }
-  const tx = { select: () => selectChain(), update, insert }
-  return {
-    db: {
-      select: () => selectChain(),
-      transaction: async (operation: (executor: typeof tx) => Promise<unknown>) => operation(tx),
-    },
-  }
-})
+vi.mock('@sim/db', () => dbChainMock)
 vi.mock('@/lib/core/idempotency/transaction', () => ({
   executeTransactionallyIdempotent: async (
     _tx: unknown,
@@ -113,23 +76,26 @@ import {
   grantDashboardUserBalance,
 } from '@/lib/admin/dashboard'
 
+/** The values object passed to the nth `update(...).set(...)` call. */
+const updateSetValues = (index = 0): Record<string, unknown> =>
+  dbChainMockFns.set.mock.calls[index]?.[0] as Record<string, unknown>
+
+afterAll(resetDbChainMock)
+
 describe('grantDashboardOrganizationBalance', () => {
   beforeEach(() => {
     vi.clearAllMocks()
-    mocks.rows = []
-    mocks.returningRows = []
+    resetDbChainMock()
     mocks.billingSubscriptions = []
-    mocks.updateSets = []
     mocks.idempotencyCalls = []
   })
 
   it('SQL-adds the grant to both fields without absorbing it into a custom limit', async () => {
-    mocks.rows = [
-      [{ id: 'org-1', creditBalance: '0.001', orgUsageLimit: '100' }],
-      [],
-      [{ value: 0 }],
-    ]
-    mocks.returningRows = [[{ creditBalance: '0.006', orgUsageLimit: '100.005' }]]
+    queueTableRows(organization, [{ id: 'org-1', creditBalance: '0.001', orgUsageLimit: '100' }])
+    queueTableRows(member, [{ value: 0 }])
+    dbChainMockFns.returning.mockResolvedValueOnce([
+      { creditBalance: '0.006', orgUsageLimit: '100.005' },
+    ])
 
     const result = await grantDashboardOrganizationBalance(
       'org-1',
@@ -139,10 +105,10 @@ describe('grantDashboardOrganizationBalance', () => {
       { id: 'admin-1', name: 'Admin', email: 'admin@sim.ai' }
     )
 
-    expect(mocks.updateSets[0].creditBalance).toBeDefined()
-    expect(mocks.updateSets[0].creditBalance).not.toBe('0.005')
-    expect(mocks.updateSets[0].orgUsageLimit).toBeDefined()
-    expect(mocks.updateSets).toHaveLength(1)
+    expect(dbChainMockFns.set).toHaveBeenCalledTimes(1)
+    expect(updateSetValues().creditBalance).toBeDefined()
+    expect(updateSetValues().creditBalance).not.toBe('0.005')
+    expect(updateSetValues().orgUsageLimit).toBeDefined()
     expect(mocks.idempotencyCalls[0]?.namespace).toBe('admin-credit-grant')
     expect(result).toEqual({ prepaidBalanceDollars: 0.006, usageLimitDollars: 100.005 })
     expect(mocks.recordAudit).toHaveBeenCalledWith(
@@ -154,22 +120,18 @@ describe('grantDashboardOrganizationBalance', () => {
 describe('grantDashboardUserBalance', () => {
   beforeEach(() => {
     vi.clearAllMocks()
-    mocks.rows = []
-    mocks.returningRows = []
+    resetDbChainMock()
     mocks.billingSubscriptions = []
-    mocks.updateSets = []
     mocks.idempotencyCalls = []
   })
 
   it('resets a free account to free-plus-prepaid before adding the grant', async () => {
-    mocks.rows = [
-      [{ id: 'user-1' }],
-      [],
-      [{ creditBalance: '0.001', currentUsageLimit: '100' }],
-      [],
-    ]
+    queueTableRows(user, [{ id: 'user-1' }])
+    queueTableRows(userStats, [{ creditBalance: '0.001', currentUsageLimit: '100' }])
     mocks.billingSubscriptions = [null, null]
-    mocks.returningRows = [[{ creditBalance: '0.006', currentUsageLimit: '5.006' }]]
+    dbChainMockFns.returning.mockResolvedValueOnce([
+      { creditBalance: '0.006', currentUsageLimit: '5.006' },
+    ])
 
     const result = await grantDashboardUserBalance(
       'user-1',
@@ -179,12 +141,12 @@ describe('grantDashboardUserBalance', () => {
       { id: 'admin-1', name: 'Admin', email: 'admin@sim.ai' }
     )
 
-    expect(mocks.updateSets).toHaveLength(1)
+    expect(dbChainMockFns.set).toHaveBeenCalledTimes(1)
     expect(mocks.acquireUserLock).toHaveBeenCalledWith(expect.anything(), 'user-1')
     expect(mocks.idempotencyCalls[0]?.namespace).toBe('admin-credit-grant')
-    expect(mocks.updateSets[0].creditBalance).toBeDefined()
-    expect(mocks.updateSets[0].currentUsageLimit).toBeDefined()
-    expect(JSON.stringify(mocks.updateSets[0].currentUsageLimit)).not.toContain('greatest')
+    expect(updateSetValues().creditBalance).toBeDefined()
+    expect(updateSetValues().currentUsageLimit).toBeDefined()
+    expect(JSON.stringify(updateSetValues().currentUsageLimit)).not.toContain('greatest')
     expect(result).toEqual({ prepaidBalanceDollars: 0.006, usageLimitDollars: 5.006 })
     expect(mocks.recordAudit).toHaveBeenCalledWith(
       expect.objectContaining({
@@ -200,14 +162,12 @@ describe('grantDashboardUserBalance', () => {
       status: 'active',
       plan: 'pro',
     }
-    mocks.rows = [
-      [{ id: 'user-1' }],
-      [],
-      [{ creditBalance: '0.001', currentUsageLimit: '100' }],
-      [],
-    ]
+    queueTableRows(user, [{ id: 'user-1' }])
+    queueTableRows(userStats, [{ creditBalance: '0.001', currentUsageLimit: '100' }])
     mocks.billingSubscriptions = [personalSubscription, personalSubscription]
-    mocks.returningRows = [[{ creditBalance: '0.006', currentUsageLimit: '100.005' }]]
+    dbChainMockFns.returning.mockResolvedValueOnce([
+      { creditBalance: '0.006', currentUsageLimit: '100.005' },
+    ])
 
     const result = await grantDashboardUserBalance(
       'user-1',
@@ -217,7 +177,7 @@ describe('grantDashboardUserBalance', () => {
       { id: 'admin-1', name: 'Admin', email: 'admin@sim.ai' }
     )
 
-    expect(JSON.stringify(mocks.updateSets[0].currentUsageLimit)).toContain('greatest')
+    expect(JSON.stringify(updateSetValues().currentUsageLimit)).toContain('greatest')
     expect(result).toEqual({ prepaidBalanceDollars: 0.006, usageLimitDollars: 100.005 })
   })
 
@@ -227,12 +187,10 @@ describe('grantDashboardUserBalance', () => {
       status: 'active',
       plan: 'enterprise',
     }
-    mocks.rows = [
-      [{ id: 'user-1' }],
-      [{ organizationId: 'org-1' }],
-      [{ creditBalance: '0', currentUsageLimit: null }],
-      [{ organizationId: 'org-1' }],
-    ]
+    queueTableRows(user, [{ id: 'user-1' }])
+    queueTableRows(member, [{ organizationId: 'org-1' }])
+    queueTableRows(userStats, [{ creditBalance: '0', currentUsageLimit: null }])
+    queueTableRows(member, [{ organizationId: 'org-1' }])
     mocks.billingSubscriptions = [organizationSubscription]
 
     await expect(
@@ -243,7 +201,7 @@ describe('grantDashboardUserBalance', () => {
       })
     ).rejects.toThrow('grant prepaid balance from Organizations instead')
 
-    expect(mocks.updateSets).toHaveLength(0)
+    expect(dbChainMockFns.set).not.toHaveBeenCalled()
     expect(mocks.recordAudit).not.toHaveBeenCalled()
   })
 })
@@ -251,15 +209,13 @@ describe('grantDashboardUserBalance', () => {
 describe('addDashboardOrganizationMember', () => {
   beforeEach(() => {
     vi.clearAllMocks()
-    mocks.rows = []
-    mocks.returningRows = []
-    mocks.updateSets = []
-    mocks.transferMembership.mockReset()
-    mocks.moveWorkspace.mockReset()
+    resetDbChainMock()
+    mocks.billingSubscriptions = []
+    mocks.idempotencyCalls = []
   })
 
   it('rejects an existing member inside the transaction before touching their cap', async () => {
-    mocks.rows = [[], [{ plan: 'enterprise' }]]
+    queueTableRows(subscription, [{ plan: 'enterprise' }])
     mocks.ensureMembership.mockResolvedValue({
       success: true,
       memberId: 'member-1',
@@ -285,10 +241,8 @@ describe('addDashboardOrganizationMember', () => {
   })
 
   it('uses the canonical transfer service and reports each selected personal workspace move', async () => {
-    mocks.rows = [
-      [{ id: 'workspace-1' }, { id: 'workspace-2' }],
-      [{ id: 'member-old', organizationId: 'org-old' }],
-    ]
+    queueTableRows(workspace, [{ id: 'workspace-1' }, { id: 'workspace-2' }])
+    queueTableRows(member, [{ id: 'member-old', organizationId: 'org-old' }])
     mocks.transferMembership.mockResolvedValue({
       success: true,
       memberId: 'member-new',
diff --git a/apps/sim/lib/workflows/persistence/utils.test.ts b/apps/sim/lib/workflows/persistence/utils.test.ts
index c719913e6cb..e0fb92ded00 100644
--- a/apps/sim/lib/workflows/persistence/utils.test.ts
+++ b/apps/sim/lib/workflows/persistence/utils.test.ts
@@ -16,8 +16,13 @@ import {
   createParallelBlock,
   createStarterBlock,
   createWorkflowState,
+  dbChainMock,
+  dbChainMockFns,
+  queueTableRows,
+  resetDbChainMock,
+  schemaMock,
 } from '@sim/testing'
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 import type {
   BlockState as AppBlockState,
   WorkflowState as AppWorkflowState,
@@ -49,78 +54,7 @@ function legacySubBlocks(subBlocks: Record<string, any>): any {
   return subBlocks
 }
 
-const { mockDb, mockWorkflowBlocks, mockWorkflowEdges, mockWorkflowSubflows } = vi.hoisted(() => {
-  const mockDb = {
-    select: vi.fn(),
-    insert: vi.fn(),
-    delete: vi.fn(),
-    transaction: vi.fn(),
-  }
-
-  const mockWorkflowBlocks = {
-    workflowId: 'workflowId',
-    id: 'id',
-    type: 'type',
-    name: 'name',
-    positionX: 'positionX',
-    positionY: 'positionY',
-    enabled: 'enabled',
-    horizontalHandles: 'horizontalHandles',
-    height: 'height',
-    subBlocks: 'subBlocks',
-    outputs: 'outputs',
-    data: 'data',
-    parentId: 'parentId',
-    extent: 'extent',
-  }
-
-  const mockWorkflowEdges = {
-    workflowId: 'workflowId',
-    id: 'id',
-    sourceBlockId: 'sourceBlockId',
-    targetBlockId: 'targetBlockId',
-    sourceHandle: 'sourceHandle',
-    targetHandle: 'targetHandle',
-  }
-
-  const mockWorkflowSubflows = {
-    workflowId: 'workflowId',
-    id: 'id',
-    type: 'type',
-    config: 'config',
-  }
-
-  return { mockDb, mockWorkflowBlocks, mockWorkflowEdges, mockWorkflowSubflows }
-})
-
-vi.mock('@sim/db', () => ({
-  db: mockDb,
-  runOutsideTransactionContext: <T>(fn: () => T): T => fn(),
-  workflowBlocks: mockWorkflowBlocks,
-  workflowEdges: mockWorkflowEdges,
-  workflowSubflows: mockWorkflowSubflows,
-  workflowDeploymentVersion: {
-    id: 'id',
-    workflowId: 'workflowId',
-    version: 'version',
-    state: 'state',
-    isActive: 'isActive',
-    createdAt: 'createdAt',
-    createdBy: 'createdBy',
-    deployedBy: 'deployedBy',
-  },
-  workflow: {},
-  webhook: {},
-  workflowDeploymentOperation: {
-    workflowId: 'workflowId',
-    status: 'status',
-  },
-  workflowSchedule: {
-    workflowId: 'workflowId',
-    deploymentVersionId: 'deploymentVersionId',
-    archivedAt: 'archivedAt',
-  },
-}))
+vi.mock('@sim/db', () => ({ ...dbChainMock, ...schemaMock }))
 
 const { mockSanitizeAgentToolsInBlocks } = vi.hoisted(() => ({
   mockSanitizeAgentToolsInBlocks: vi.fn(),
@@ -128,8 +62,8 @@ const { mockSanitizeAgentToolsInBlocks } = vi.hoisted(() => ({
 
 /**
  * Default identity behavior for the mocked migration step. Re-applied in the
- * cache describe block's `beforeEach` because the outer `afterEach` calls
- * `vi.resetAllMocks()`, which clears implementations.
+ * outer `beforeEach` because `vi.clearAllMocks()` clears implementations set
+ * on the hoisted spy.
  */
 const sanitizeIdentity = (blocks: unknown) => ({ blocks })
 mockSanitizeAgentToolsInBlocks.mockImplementation(sanitizeIdentity)
@@ -142,6 +76,35 @@ import * as dbHelpers from '@/lib/workflows/persistence/utils'
 
 const mockWorkflowId = 'test-workflow-123'
 
+/**
+ * Queues the four table-routed result sets consumed by
+ * `loadWorkflowFromNormalizedTablesRaw` (blocks, edges, subflows, workflow row).
+ */
+function queueLoadFixtures(options: {
+  blocks: unknown[]
+  edges?: unknown[]
+  subflows?: unknown[]
+  workspaceId?: string
+}) {
+  queueTableRows(schemaMock.workflowBlocks, options.blocks)
+  queueTableRows(schemaMock.workflowEdges, options.edges ?? [])
+  queueTableRows(schemaMock.workflowSubflows, options.subflows ?? [])
+  queueTableRows(schemaMock.workflow, [{ workspaceId: options.workspaceId ?? 'test-workspace-id' }])
+}
+
+/**
+ * Returns the row arrays passed to `insert(table).values(rows)` for the given
+ * schema table. Insert/values chains run sequentially in the code under test,
+ * so the two spies' call lists stay index-aligned.
+ */
+function insertedRowsFor(table: unknown): Record<string, unknown>[][] {
+  return dbChainMockFns.insert.mock.calls.flatMap(([calledTable], index) =>
+    calledTable === table && Array.isArray(dbChainMockFns.values.mock.calls[index]?.[0])
+      ? [dbChainMockFns.values.mock.calls[index][0] as Record<string, unknown>[]]
+      : []
+  )
+}
+
 /**
  * Converts a BlockState to a mock database block row format.
  */
@@ -332,11 +295,12 @@ const mockWorkflowState = createWorkflowState({
 describe('Database Helpers', () => {
   beforeEach(() => {
     vi.clearAllMocks()
+    resetDbChainMock()
     mockSanitizeAgentToolsInBlocks.mockImplementation(sanitizeIdentity)
   })
 
-  afterEach(() => {
-    vi.resetAllMocks()
+  afterAll(() => {
+    resetDbChainMock()
   })
 
   describe('buildWorkflowDeploymentSnapshot', () => {
@@ -377,29 +341,11 @@ describe('Database Helpers', () => {
 
   describe('loadWorkflowFromNormalizedTables', () => {
     it('should successfully load workflow data from normalized tables', async () => {
-      vi.clearAllMocks()
-
-      let callCount = 0
-      mockDb.select.mockImplementation(() => ({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockImplementation(() => {
-            callCount++
-            if (callCount === 1) {
-              return Promise.resolve(mockBlocksFromDb)
-            }
-            if (callCount === 2) {
-              return Promise.resolve(mockEdgesFromDb)
-            }
-            if (callCount === 3) {
-              return Promise.resolve(mockSubflowsFromDb)
-            }
-            if (callCount === 4) {
-              return { limit: vi.fn().mockResolvedValue([{ workspaceId: 'test-workspace-id' }]) }
-            }
-            return Promise.resolve([])
-          }),
-        }),
-      }))
+      queueLoadFixtures({
+        blocks: mockBlocksFromDb,
+        edges: mockEdgesFromDb,
+        subflows: mockSubflowsFromDb,
+      })
 
       const result = await dbHelpers.loadWorkflowFromNormalizedTables(mockWorkflowId)
 
@@ -465,23 +411,15 @@ describe('Database Helpers', () => {
     })
 
     it('should return null when no blocks are found', async () => {
-      mockDb.select.mockReturnValue({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockResolvedValue([]),
-        }),
-      })
-
       const result = await dbHelpers.loadWorkflowFromNormalizedTables(mockWorkflowId)
 
       expect(result).toBeNull()
     })
 
     it('should return null when database query fails', async () => {
-      mockDb.select.mockReturnValue({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockRejectedValue(new Error('Database connection failed')),
-        }),
-      })
+      dbChainMockFns.where.mockImplementationOnce(() =>
+        Promise.reject(new Error('Database connection failed'))
+      )
 
       const result = await dbHelpers.loadWorkflowFromNormalizedTables(mockWorkflowId)
 
@@ -498,19 +436,10 @@ describe('Database Helpers', () => {
         },
       ]
 
-      let callCount = 0
-      mockDb.select.mockReturnValue({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockImplementation(() => {
-            callCount++
-            if (callCount === 1) return Promise.resolve(mockBlocksFromDb)
-            if (callCount === 2) return Promise.resolve(mockEdgesFromDb)
-            if (callCount === 3) return Promise.resolve(subflowsWithUnknownType)
-            if (callCount === 4)
-              return { limit: vi.fn().mockResolvedValue([{ workspaceId: 'test-workspace-id' }]) }
-            return Promise.resolve([])
-          }),
-        }),
+      queueLoadFixtures({
+        blocks: mockBlocksFromDb,
+        edges: mockEdgesFromDb,
+        subflows: subflowsWithUnknownType,
       })
 
       const result = await dbHelpers.loadWorkflowFromNormalizedTables(mockWorkflowId)
@@ -538,20 +467,7 @@ describe('Database Helpers', () => {
       malformedBlocks[0].type = null as any
       malformedBlocks[0].name = null as any
 
-      let callCount = 0
-      mockDb.select.mockReturnValue({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockImplementation(() => {
-            callCount++
-            if (callCount === 1) return Promise.resolve(malformedBlocks)
-            if (callCount === 2) return Promise.resolve([])
-            if (callCount === 3) return Promise.resolve([])
-            if (callCount === 4)
-              return { limit: vi.fn().mockResolvedValue([{ workspaceId: 'test-workspace-id' }]) }
-            return Promise.resolve([])
-          }),
-        }),
-      })
+      queueLoadFixtures({ blocks: malformedBlocks })
 
       const result = await dbHelpers.loadWorkflowFromNormalizedTables(mockWorkflowId)
 
@@ -565,11 +481,7 @@ describe('Database Helpers', () => {
       const connectionError = new Error('Connection refused')
       ;(connectionError as any).code = 'ECONNREFUSED'
 
-      mockDb.select.mockReturnValue({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockRejectedValue(connectionError),
-        }),
-      })
+      dbChainMockFns.where.mockImplementationOnce(() => Promise.reject(connectionError))
 
       const result = await dbHelpers.loadWorkflowFromNormalizedTables(mockWorkflowId)
 
@@ -579,25 +491,6 @@ describe('Database Helpers', () => {
 
   describe('saveWorkflowToNormalizedTables', () => {
     it('should successfully save workflow data to normalized tables', async () => {
-      const mockTransaction = vi.fn().mockImplementation(async (callback) => {
-        const tx = {
-          select: vi.fn().mockReturnValue({
-            from: vi.fn().mockReturnValue({
-              where: vi.fn().mockResolvedValue([]),
-            }),
-          }),
-          delete: vi.fn().mockReturnValue({
-            where: vi.fn().mockResolvedValue([]),
-          }),
-          insert: vi.fn().mockReturnValue({
-            values: vi.fn().mockResolvedValue([]),
-          }),
-        }
-        return await callback(tx)
-      })
-
-      mockDb.transaction = mockTransaction
-
       const result = await dbHelpers.saveWorkflowToNormalizedTables(
         mockWorkflowId,
         asAppState(mockWorkflowState)
@@ -605,31 +498,12 @@ describe('Database Helpers', () => {
 
       expect(result.success).toBe(true)
 
-      expect(mockTransaction).toHaveBeenCalledTimes(1)
+      expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(1)
     })
 
     it('should handle empty workflow state gracefully', async () => {
       const emptyWorkflowState = createWorkflowState()
 
-      const mockTransaction = vi.fn().mockImplementation(async (callback) => {
-        const tx = {
-          select: vi.fn().mockReturnValue({
-            from: vi.fn().mockReturnValue({
-              where: vi.fn().mockResolvedValue([]),
-            }),
-          }),
-          delete: vi.fn().mockReturnValue({
-            where: vi.fn().mockResolvedValue([]),
-          }),
-          insert: vi.fn().mockReturnValue({
-            values: vi.fn().mockResolvedValue([]),
-          }),
-        }
-        return await callback(tx)
-      })
-
-      mockDb.transaction = mockTransaction
-
       const result = await dbHelpers.saveWorkflowToNormalizedTables(
         mockWorkflowId,
         asAppState(emptyWorkflowState)
@@ -639,8 +513,7 @@ describe('Database Helpers', () => {
     })
 
     it('should return error when transaction fails', async () => {
-      const mockTransaction = vi.fn().mockRejectedValue(new Error('Transaction failed'))
-      mockDb.transaction = mockTransaction
+      dbChainMockFns.transaction.mockRejectedValueOnce(new Error('Transaction failed'))
 
       const result = await dbHelpers.saveWorkflowToNormalizedTables(
         mockWorkflowId,
@@ -655,8 +528,7 @@ describe('Database Helpers', () => {
       const constraintError = new Error('Unique constraint violation')
       ;(constraintError as any).code = '23505'
 
-      const mockTransaction = vi.fn().mockRejectedValue(constraintError)
-      mockDb.transaction = mockTransaction
+      dbChainMockFns.transaction.mockRejectedValueOnce(constraintError)
 
       const result = await dbHelpers.saveWorkflowToNormalizedTables(
         mockWorkflowId,
@@ -668,42 +540,12 @@ describe('Database Helpers', () => {
     })
 
     it('should properly format block data for database insertion', async () => {
-      let capturedBlockInserts: any[] = []
-      let capturedEdgeInserts: any[] = []
-      let capturedSubflowInserts: any[] = []
-
-      const mockTransaction = vi.fn().mockImplementation(async (callback) => {
-        const tx = {
-          select: vi.fn().mockReturnValue({
-            from: vi.fn().mockReturnValue({
-              where: vi.fn().mockResolvedValue([]),
-            }),
-          }),
-          delete: vi.fn().mockReturnValue({
-            where: vi.fn().mockResolvedValue([]),
-          }),
-          insert: vi.fn().mockReturnValue({
-            values: vi.fn().mockImplementation((data) => {
-              if (data.length > 0) {
-                if (data[0].positionX !== undefined) {
-                  capturedBlockInserts = data
-                } else if (data[0].sourceBlockId !== undefined) {
-                  capturedEdgeInserts = data
-                } else if (data[0].type === 'loop' || data[0].type === 'parallel') {
-                  capturedSubflowInserts = data
-                }
-              }
-              return Promise.resolve([])
-            }),
-          }),
-        }
-        return await callback(tx)
-      })
-
-      mockDb.transaction = mockTransaction
-
       await dbHelpers.saveWorkflowToNormalizedTables(mockWorkflowId, asAppState(mockWorkflowState))
 
+      const [capturedBlockInserts = []] = insertedRowsFor(schemaMock.workflowBlocks)
+      const [capturedEdgeInserts = []] = insertedRowsFor(schemaMock.workflowEdges)
+      const [capturedSubflowInserts = []] = insertedRowsFor(schemaMock.workflowSubflows)
+
       expect(capturedBlockInserts).toHaveLength(5)
       expect(capturedBlockInserts).toEqual(
         expect.arrayContaining([
@@ -768,38 +610,14 @@ describe('Database Helpers', () => {
     })
 
     it('should regenerate missing loop and parallel definitions from block data', async () => {
-      let capturedSubflowInserts: any[] = []
-
-      const mockTransaction = vi.fn().mockImplementation(async (callback) => {
-        const tx = {
-          select: vi.fn().mockReturnValue({
-            from: vi.fn().mockReturnValue({
-              where: vi.fn().mockResolvedValue([]),
-            }),
-          }),
-          delete: vi.fn().mockReturnValue({
-            where: vi.fn().mockResolvedValue([]),
-          }),
-          insert: vi.fn().mockReturnValue({
-            values: vi.fn().mockImplementation((data) => {
-              if (data.length > 0 && (data[0].type === 'loop' || data[0].type === 'parallel')) {
-                capturedSubflowInserts = data
-              }
-              return Promise.resolve([])
-            }),
-          }),
-        }
-        return await callback(tx)
-      })
-
-      mockDb.transaction = mockTransaction
-
       const staleWorkflowState = structuredClone(mockWorkflowState)
       staleWorkflowState.loops = {}
       staleWorkflowState.parallels = {}
 
       await dbHelpers.saveWorkflowToNormalizedTables(mockWorkflowId, asAppState(staleWorkflowState))
 
+      const [capturedSubflowInserts = []] = insertedRowsFor(schemaMock.workflowSubflows)
+
       expect(capturedSubflowInserts).toHaveLength(2)
       expect(capturedSubflowInserts).toEqual(
         expect.arrayContaining([
@@ -816,13 +634,7 @@ describe('Database Helpers', () => {
 
   describe('workflowExistsInNormalizedTables', () => {
     it('should return true when workflow exists in normalized tables', async () => {
-      mockDb.select.mockReturnValue({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockReturnValue({
-            limit: vi.fn().mockResolvedValue([{ id: 'block-1' }]),
-          }),
-        }),
-      })
+      queueTableRows(schemaMock.workflowBlocks, [{ id: 'block-1' }])
 
       const result = await dbHelpers.workflowExistsInNormalizedTables(mockWorkflowId)
 
@@ -830,27 +642,13 @@ describe('Database Helpers', () => {
     })
 
     it('should return false when workflow does not exist in normalized tables', async () => {
-      mockDb.select.mockReturnValue({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockReturnValue({
-            limit: vi.fn().mockResolvedValue([]),
-          }),
-        }),
-      })
-
       const result = await dbHelpers.workflowExistsInNormalizedTables(mockWorkflowId)
 
       expect(result).toBe(false)
     })
 
     it('should return false when database query fails', async () => {
-      mockDb.select.mockReturnValue({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockReturnValue({
-            limit: vi.fn().mockRejectedValue(new Error('Database error')),
-          }),
-        }),
-      })
+      dbChainMockFns.limit.mockImplementationOnce(() => Promise.reject(new Error('Database error')))
 
       const result = await dbHelpers.workflowExistsInNormalizedTables(mockWorkflowId)
 
@@ -859,60 +657,19 @@ describe('Database Helpers', () => {
   })
 
   describe('workflow row locking', () => {
-    function createMissingWorkflowTx() {
-      const lockFor = vi.fn().mockResolvedValue([])
-      const limit = vi.fn(() => ({ for: lockFor }))
-      const where = vi.fn(() => ({ limit }))
-      const from = vi.fn(() => ({ where }))
-      const select = vi.fn(() => ({ from }))
-      const update = vi.fn()
-
-      return {
-        tx: {
-          execute: vi.fn().mockResolvedValue([{ id: mockWorkflowId }]),
-          select,
-          update,
-        },
-        lockFor,
-        update,
-      }
-    }
-
     it('returns an error when undeploy cannot lock a workflow row', async () => {
-      const { tx, update } = createMissingWorkflowTx()
-      mockDb.transaction = vi.fn().mockImplementation(async (callback) => callback(tx))
-
       const result = await dbHelpers.undeployWorkflow({ workflowId: mockWorkflowId })
 
       expect(result).toEqual({
         success: false,
         error: 'Workflow not found',
       })
-      expect(update).not.toHaveBeenCalled()
+      expect(dbChainMockFns.update).not.toHaveBeenCalled()
     })
 
     it('supersedes in-flight operations and releases path claims during undeploy', async () => {
-      const versionRows = [{ id: 'dv-1' }, { id: 'dv-2' }]
-      const createWhereResult = () => ({
-        limit: vi.fn(() => ({
-          for: vi.fn().mockResolvedValue([{ id: mockWorkflowId }]),
-        })),
-        then: (resolve: (rows: typeof versionRows) => void) => resolve(versionRows),
-      })
-      const setCalls: unknown[] = []
-      const tx = {
-        select: vi.fn(() => ({
-          from: vi.fn(() => ({ where: vi.fn(() => createWhereResult()) })),
-        })),
-        update: vi.fn(() => ({
-          set: vi.fn((payload: unknown) => {
-            setCalls.push(payload)
-            return { where: vi.fn().mockResolvedValue([]) }
-          }),
-        })),
-        delete: vi.fn(() => ({ where: vi.fn().mockResolvedValue([]) })),
-      }
-      mockDb.transaction = vi.fn().mockImplementation(async (callback) => callback(tx))
+      queueTableRows(schemaMock.workflow, [{ id: mockWorkflowId }])
+      queueTableRows(schemaMock.workflowDeploymentVersion, [{ id: 'dv-1' }, { id: 'dv-2' }])
       const onUndeployTransaction = vi.fn().mockResolvedValue(undefined)
 
       const result = await dbHelpers.undeployWorkflow({
@@ -921,6 +678,7 @@ describe('Database Helpers', () => {
       })
 
       expect(result).toEqual({ success: true })
+      const setCalls = dbChainMockFns.set.mock.calls.map(([payload]) => payload)
       expect(setCalls[0]).toEqual(expect.objectContaining({ status: 'superseded' }))
       expect(setCalls).toEqual(
         expect.arrayContaining([
@@ -928,8 +686,8 @@ describe('Database Helpers', () => {
           expect.objectContaining({ isDeployed: false, deployedAt: null }),
         ])
       )
-      expect(tx.delete).toHaveBeenCalledTimes(2)
-      expect(onUndeployTransaction).toHaveBeenCalledWith(tx, {
+      expect(dbChainMockFns.delete).toHaveBeenCalledTimes(2)
+      expect(onUndeployTransaction).toHaveBeenCalledWith(dbChainMock.db, {
         deploymentVersionIds: ['dv-1', 'dv-2'],
       })
     })
@@ -960,25 +718,6 @@ describe('Database Helpers', () => {
 
       const largeWorkflowState = createWorkflowState({ blocks, edges })
 
-      const mockTransaction = vi.fn().mockImplementation(async (callback) => {
-        const tx = {
-          select: vi.fn().mockReturnValue({
-            from: vi.fn().mockReturnValue({
-              where: vi.fn().mockResolvedValue([]),
-            }),
-          }),
-          delete: vi.fn().mockReturnValue({
-            where: vi.fn().mockResolvedValue([]),
-          }),
-          insert: vi.fn().mockReturnValue({
-            values: vi.fn().mockResolvedValue([]),
-          }),
-        }
-        return await callback(tx)
-      })
-
-      mockDb.transaction = mockTransaction
-
       const result = await dbHelpers.saveWorkflowToNormalizedTables(
         mockWorkflowId,
         asAppState(largeWorkflowState)
@@ -1015,22 +754,7 @@ describe('Database Helpers', () => {
       testBlocks[0].advancedMode = true
       testBlocks[1].advancedMode = false
 
-      vi.clearAllMocks()
-
-      let callCount = 0
-      mockDb.select.mockImplementation(() => ({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockImplementation(() => {
-            callCount++
-            if (callCount === 1) return Promise.resolve(testBlocks)
-            if (callCount === 2) return Promise.resolve([])
-            if (callCount === 3) return Promise.resolve([])
-            if (callCount === 4)
-              return { limit: vi.fn().mockResolvedValue([{ workspaceId: 'test-workspace-id' }]) }
-            return Promise.resolve([])
-          }),
-        }),
-      }))
+      queueLoadFixtures({ blocks: testBlocks })
 
       const result = await dbHelpers.loadWorkflowFromNormalizedTables(mockWorkflowId)
 
@@ -1056,20 +780,7 @@ describe('Database Helpers', () => {
         ),
       ]
 
-      vi.clearAllMocks()
-
-      let callCount = 0
-      mockDb.select.mockImplementation(() => ({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockImplementation(() => {
-            callCount++
-            if (callCount === 1) return Promise.resolve(blocksWithDefaultValues)
-            if (callCount === 4)
-              return { limit: vi.fn().mockResolvedValue([{ workspaceId: 'test-workspace-id' }]) }
-            return Promise.resolve([])
-          }),
-        }),
-      }))
+      queueLoadFixtures({ blocks: blocksWithDefaultValues })
 
       const result = await dbHelpers.loadWorkflowFromNormalizedTables(mockWorkflowId)
 
@@ -1125,22 +836,7 @@ describe('Database Helpers', () => {
       )
       duplicatedBlock.advancedMode = true
 
-      vi.clearAllMocks()
-
-      let callCount = 0
-      mockDb.select.mockImplementation(() => ({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockImplementation(() => {
-            callCount++
-            if (callCount === 1) return Promise.resolve([originalBlock, duplicatedBlock])
-            if (callCount === 2) return Promise.resolve([])
-            if (callCount === 3) return Promise.resolve([])
-            if (callCount === 4)
-              return { limit: vi.fn().mockResolvedValue([{ workspaceId: 'test-workspace-id' }]) }
-            return Promise.resolve([])
-          }),
-        }),
-      }))
+      queueLoadFixtures({ blocks: [originalBlock, duplicatedBlock] })
 
       const loadedState = await dbHelpers.loadWorkflowFromNormalizedTables(mockWorkflowId)
       expect(loadedState).toBeDefined()
@@ -1154,44 +850,19 @@ describe('Database Helpers', () => {
         parallels: {},
       }
 
-      const mockTransaction = vi.fn().mockImplementation(async (callback) => {
-        const mockTx = {
-          select: vi.fn().mockReturnValue({
-            from: vi.fn().mockReturnValue({
-              where: vi.fn().mockResolvedValue([]),
-            }),
-          }),
-          delete: vi.fn().mockReturnValue({
-            where: vi.fn().mockResolvedValue(undefined),
-          }),
-          insert: vi.fn().mockImplementation((_table) => ({
-            values: vi.fn().mockImplementation((values) => {
-              if (Array.isArray(values)) {
-                values.forEach((blockInsert) => {
-                  if (blockInsert.id === 'agent-original') {
-                    expect(blockInsert.advancedMode).toBe(true)
-                  }
-                  if (blockInsert.id === 'agent-duplicate') {
-                    expect(blockInsert.advancedMode).toBe(true)
-                  }
-                })
-              }
-              return Promise.resolve()
-            }),
-          })),
-        }
-        return await callback(mockTx)
-      })
-
-      mockDb.transaction = mockTransaction
-
       const saveResult = await dbHelpers.saveWorkflowToNormalizedTables(
         mockWorkflowId,
         workflowState
       )
       expect(saveResult.success).toBe(true)
 
-      expect(mockTransaction).toHaveBeenCalled()
+      expect(dbChainMockFns.transaction).toHaveBeenCalled()
+
+      const [blockInserts = []] = insertedRowsFor(schemaMock.workflowBlocks)
+      const savedOriginal = blockInserts.find((row) => row.id === 'agent-original')
+      const savedDuplicate = blockInserts.find((row) => row.id === 'agent-duplicate')
+      expect(savedOriginal?.advancedMode).toBe(true)
+      expect(savedDuplicate?.advancedMode).toBe(true)
     })
 
     it('should handle mixed advancedMode states correctly', async () => {
@@ -1224,20 +895,7 @@ describe('Database Helpers', () => {
       )
       advancedBlock.advancedMode = true
 
-      vi.clearAllMocks()
-
-      let callCount = 0
-      mockDb.select.mockImplementation(() => ({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockImplementation(() => {
-            callCount++
-            if (callCount === 1) return Promise.resolve([basicBlock, advancedBlock])
-            if (callCount === 4)
-              return { limit: vi.fn().mockResolvedValue([{ workspaceId: 'test-workspace-id' }]) }
-            return Promise.resolve([])
-          }),
-        }),
-      }))
+      queueLoadFixtures({ blocks: [basicBlock, advancedBlock] })
 
       const loadedState = await dbHelpers.loadWorkflowFromNormalizedTables(mockWorkflowId)
       expect(loadedState).toBeDefined()
@@ -1263,67 +921,36 @@ describe('Database Helpers', () => {
         },
       })
 
-      const mockTransaction = vi.fn().mockImplementation(async (callback) => {
-        const mockTx = {
-          select: vi.fn().mockReturnValue({
-            from: vi.fn().mockReturnValue({
-              where: vi.fn().mockResolvedValue([]),
-            }),
-          }),
-          delete: vi.fn().mockReturnValue({
-            where: vi.fn().mockResolvedValue(undefined),
-          }),
-          insert: vi.fn().mockReturnValue({
-            values: vi.fn().mockResolvedValue(undefined),
-          }),
-        }
-        return await callback(mockTx)
-      })
-
-      mockDb.transaction = mockTransaction
-
       const saveResult = await dbHelpers.saveWorkflowToNormalizedTables(
         mockWorkflowId,
         asAppState(testWorkflowState)
       )
       expect(saveResult.success).toBe(true)
 
-      vi.clearAllMocks()
-      let callCount = 0
-      mockDb.select.mockImplementation(() => ({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockImplementation(() => {
-            callCount++
-            if (callCount === 1) {
-              return Promise.resolve([
-                {
-                  id: 'block-1',
-                  workflowId: mockWorkflowId,
-                  type: 'agent',
-                  name: 'Test Agent',
-                  positionX: 100,
-                  positionY: 100,
-                  enabled: true,
-                  horizontalHandles: true,
-                  advancedMode: true,
-                  height: 200,
-                  subBlocks: {
-                    systemPrompt: { id: 'systemPrompt', type: 'textarea', value: 'System' },
-                    model: { id: 'model', type: 'select', value: 'gpt-4o' },
-                  },
-                  outputs: {},
-                  data: {},
-                  parentId: null,
-                  extent: null,
-                },
-              ])
-            }
-            if (callCount === 4)
-              return { limit: vi.fn().mockResolvedValue([{ workspaceId: 'test-workspace-id' }]) }
-            return Promise.resolve([])
-          }),
-        }),
-      }))
+      queueLoadFixtures({
+        blocks: [
+          {
+            id: 'block-1',
+            workflowId: mockWorkflowId,
+            type: 'agent',
+            name: 'Test Agent',
+            positionX: 100,
+            positionY: 100,
+            enabled: true,
+            horizontalHandles: true,
+            advancedMode: true,
+            height: 200,
+            subBlocks: {
+              systemPrompt: { id: 'systemPrompt', type: 'textarea', value: 'System' },
+              model: { id: 'model', type: 'select', value: 'gpt-4o' },
+            },
+            outputs: {},
+            data: {},
+            parentId: null,
+            extent: null,
+          },
+        ],
+      })
 
       const loadedState = await dbHelpers.loadWorkflowFromNormalizedTables(mockWorkflowId)
       expect(loadedState).toBeDefined()
@@ -1651,30 +1278,23 @@ describe('Database Helpers', () => {
     }
 
     /**
-     * Wires `db.select` to return a single active deployment-version row for the
-     * given id. Returns the inner `where` spy so tests can assert how many times
-     * the active-version SELECT ran.
+     * Queues one active deployment-version row for the next active-version
+     * SELECT; call once per expected `loadDeployedWorkflowState` invocation.
+     * Tests assert SELECT counts on `dbChainMockFns.where`.
      */
-    function mockActiveVersionSelect(versionId: string, state: unknown) {
-      const where = vi.fn().mockReturnValue({
-        orderBy: vi.fn().mockReturnValue({
-          limit: vi.fn().mockResolvedValue([{ id: versionId, state, createdAt: new Date() }]),
-        }),
-      })
-      mockDb.select.mockReturnValue({
-        from: vi.fn().mockReturnValue({ where }),
-      })
-      return where
+    function queueActiveVersion(versionId: string, state: unknown) {
+      queueTableRows(schemaMock.workflowDeploymentVersion, [
+        { id: versionId, state, createdAt: new Date() },
+      ])
     }
 
     beforeEach(() => {
-      vi.clearAllMocks()
-      mockSanitizeAgentToolsInBlocks.mockImplementation(sanitizeIdentity)
       dbHelpers.invalidateDeployedStateCache()
     })
 
     it('serves a cache HIT, skipping migrations on the second call for the same active version', async () => {
-      const where = mockActiveVersionSelect('dv-hit', buildDeployedState())
+      queueActiveVersion('dv-hit', buildDeployedState())
+      queueActiveVersion('dv-hit', buildDeployedState())
 
       const first = await dbHelpers.loadDeployedWorkflowState('wf-1', 'workspace-1')
       const second = await dbHelpers.loadDeployedWorkflowState('wf-1', 'workspace-1')
@@ -1682,20 +1302,22 @@ describe('Database Helpers', () => {
       expect(first).toBeDefined()
       expect(second).toBeDefined()
       expect(mockSanitizeAgentToolsInBlocks).toHaveBeenCalledTimes(1)
-      expect(where).toHaveBeenCalledTimes(2)
+      expect(dbChainMockFns.where).toHaveBeenCalledTimes(2)
     })
 
     it('still runs the active-version SELECT on every call so rollback/redeploy stays observable', async () => {
-      const where = mockActiveVersionSelect('dv-active', buildDeployedState())
+      queueActiveVersion('dv-active', buildDeployedState())
+      queueActiveVersion('dv-active', buildDeployedState())
 
       await dbHelpers.loadDeployedWorkflowState('wf-2', 'workspace-1')
       await dbHelpers.loadDeployedWorkflowState('wf-2', 'workspace-1')
 
-      expect(where).toHaveBeenCalledTimes(2)
+      expect(dbChainMockFns.where).toHaveBeenCalledTimes(2)
     })
 
     it('deep-clones on read: mutating the first result does not corrupt the cached copy', async () => {
-      mockActiveVersionSelect('dv-clone', buildDeployedState())
+      queueActiveVersion('dv-clone', buildDeployedState())
+      queueActiveVersion('dv-clone', buildDeployedState())
 
       const first = await dbHelpers.loadDeployedWorkflowState('wf-3', 'workspace-1')
       ;(first.blocks['block-1'] as any).name = 'MUTATED'
@@ -1715,22 +1337,18 @@ describe('Database Helpers', () => {
     })
 
     it('keys the cache by deploymentVersionId: a different active id triggers a fresh build', async () => {
-      mockActiveVersionSelect('dv-old', buildDeployedState())
+      queueActiveVersion('dv-old', buildDeployedState())
       await dbHelpers.loadDeployedWorkflowState('wf-4', 'workspace-1')
       expect(mockSanitizeAgentToolsInBlocks).toHaveBeenCalledTimes(1)
 
-      mockActiveVersionSelect('dv-new', buildDeployedState())
+      queueActiveVersion('dv-new', buildDeployedState())
       await dbHelpers.loadDeployedWorkflowState('wf-4', 'workspace-1')
       expect(mockSanitizeAgentToolsInBlocks).toHaveBeenCalledTimes(2)
     })
 
     it('loads an admitted immutable deployment version even after a later cutover', async () => {
       const state = buildDeployedState()
-      const limit = vi.fn().mockResolvedValue([{ id: 'dv-admitted', state }])
-      const where = vi.fn().mockReturnValue({ limit })
-      mockDb.select.mockReturnValue({
-        from: vi.fn().mockReturnValue({ where }),
-      })
+      queueTableRows(schemaMock.workflowDeploymentVersion, [{ id: 'dv-admitted', state }])
 
       const result = await dbHelpers.loadWorkflowDeploymentVersionState(
         'wf-admitted',
@@ -1740,11 +1358,13 @@ describe('Database Helpers', () => {
 
       expect(result.deploymentVersionId).toBe('dv-admitted')
       expect(result.blocks).toEqual(state.blocks)
-      expect(where).toHaveBeenCalledTimes(1)
+      expect(dbChainMockFns.where).toHaveBeenCalledTimes(1)
     })
 
     it('invalidateDeployedStateCache(id) forces a rebuild on the next call', async () => {
-      mockActiveVersionSelect('dv-inv', buildDeployedState())
+      queueActiveVersion('dv-inv', buildDeployedState())
+      queueActiveVersion('dv-inv', buildDeployedState())
+      queueActiveVersion('dv-inv', buildDeployedState())
 
       await dbHelpers.loadDeployedWorkflowState('wf-5', 'workspace-1')
       await dbHelpers.loadDeployedWorkflowState('wf-5', 'workspace-1')
@@ -1757,15 +1377,6 @@ describe('Database Helpers', () => {
     })
 
     it('throws when there is no active deployment and does not cache the failure', async () => {
-      const where = vi.fn().mockReturnValue({
-        orderBy: vi.fn().mockReturnValue({
-          limit: vi.fn().mockResolvedValue([]),
-        }),
-      })
-      mockDb.select.mockReturnValue({
-        from: vi.fn().mockReturnValue({ where }),
-      })
-
       await expect(dbHelpers.loadDeployedWorkflowState('wf-6', 'workspace-1')).rejects.toThrow(
         'Workflow wf-6 has no active deployment'
       )
diff --git a/packages/testing/src/mocks/database.mock.test.ts b/packages/testing/src/mocks/database.mock.test.ts
index c32fc940156..6e5f58fa178 100644
--- a/packages/testing/src/mocks/database.mock.test.ts
+++ b/packages/testing/src/mocks/database.mock.test.ts
@@ -97,6 +97,14 @@ describe('database mock', () => {
     ).resolves.toEqual([{ id: 'from-row' }])
   })
 
+  it('supports the .limit(n).for(mode) row-lock chain', async () => {
+    queueTableRows(workflowTable, [{ id: 'locked' }])
+    await expect(db.select().from(workflowTable).where({}).limit(1).for('update')).resolves.toEqual(
+      [{ id: 'locked' }]
+    )
+    expect(dbChainMockFns.for).toHaveBeenCalledWith('update')
+  })
+
   it('never lets mutation chains consume select queues', async () => {
     queueTableRows(workflowTable, [{ id: 'kept' }])
     await expect(db.update(workflowTable).set({}).where({})).resolves.toEqual([])
diff --git a/packages/testing/src/mocks/database.mock.ts b/packages/testing/src/mocks/database.mock.ts
index a432bc43972..897d2c5f0d1 100644
--- a/packages/testing/src/mocks/database.mock.ts
+++ b/packages/testing/src/mocks/database.mock.ts
@@ -73,6 +73,12 @@ export function createMockSqlOperators() {
  *
  * The queue is keyed by table object identity, so pass the same schema-mock
  * table object the code under test passes to `.from()` / the join.
+ *
+ * Footgun: because a chain falls back to its JOIN tables when the `.from()`
+ * table has nothing queued, a `from(A).innerJoin(B)` chain you expect to
+ * resolve empty will consume a set queued for a LATER select on `B`. When a
+ * suite queues `B` for a subsequent query, queue an explicit empty set on `A`
+ * first (`queueTableRows(A, [])`) so the joined chain consumes that instead.
  */
 const tableRowQueues = new Map<unknown, unknown[][]>()
 
@@ -221,10 +227,12 @@ const lazyRowsThenable = (getRows: RowsSupplier): any => ({
 })
 
 // `.limit()` returns a builder that is awaitable and also exposes `.offset()`
-// for keyset/OFFSET paging (`.limit(n).offset(m)`).
+// for keyset/OFFSET paging (`.limit(n).offset(m)`) and `.for()` for drizzle's
+// `.limit(1).for('update')` row-lock form.
 const limitBuilder = (getRows: RowsSupplier) => {
   const thenable = lazyRowsThenable(getRows)
   thenable.offset = spyOrDefault(offset, () => lazyRowsThenable(getRows))
+  thenable.for = spyOrDefault(forClause, () => limitBuilder(getRows))
   return thenable
 }