improvement(tests): db-mock migration tranche 2 — copilot, mothership, workspaces, connectors, mcp (#5863)

* improvement(tests): db-mock migration tranche 2 — copilot, mothership, workspaces, connectors, mcp

* fix(testing): drain unconsumed ...Once overrides in resetDbChainMock

vi.clearAllMocks clears call history only — a ...Once override queued by a
previous test but never consumed survived into the next test. resetDbChainMock
now mockReset()s every shared spy and stable wrapper, which restores the
original implementation AND drains once-queues.

Author: waleedlatif1

apps/sim/app/api/copilot/api-keys/validate/route.test.ts

@@ -1,11 +1,10 @@
 /**
  * @vitest-environment node
  */
-import { createMockRequest } from '@sim/testing'
-import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { createMockRequest, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
 const {
-  mockDbLimit,
   mockCheckInternalApiKey,
   mockCheckAttributedUsageLimits,
   mockCheckServerSideUsageLimits,
@@ -21,7 +20,6 @@ const {
   mockGetWorkspaceBillingSettings,
   mockFlags,
 } = vi.hoisted(() => ({
-  mockDbLimit: vi.fn(),
   mockCheckInternalApiKey: vi.fn(),
   mockCheckAttributedUsageLimits: vi.fn(),
   mockCheckServerSideUsageLimits: vi.fn(),
@@ -77,12 +75,6 @@ const OLD_GO_OPAQUE_WORKSPACE_VALIDATE_BODY = {
   workspaceId: 'local-self-hosted-workspace',
 } as const
 
-vi.mock('@sim/db', () => ({
-  db: {
-    select: () => ({ from: () => ({ where: () => ({ limit: mockDbLimit }) }) }),
-  },
-}))
-
 vi.mock('@/lib/billing/core/billing-attribution', () => ({
   BILLING_ACCOUNT_DECISION_HEADER: 'x-sim-billing-account-decision',
   BILLING_ATTRIBUTION_HEADER: 'x-sim-billing-attribution',
@@ -151,9 +143,10 @@ function request(body: Record<string, unknown>, headers: Record<string, string>
 describe('POST /api/copilot/api-keys/validate billing protocols', () => {
   beforeEach(() => {
     vi.clearAllMocks()
+    resetDbChainMock()
     mockFlags.isCopilotBillingProtocolRequired = false
     mockCheckInternalApiKey.mockReturnValue({ success: true })
-    mockDbLimit.mockResolvedValue([{ id: 'user-1' }])
+    queueTableRows(schemaMock.user, [{ id: 'user-1' }])
     mockResolveBillingAttribution.mockResolvedValue(ATTRIBUTION)
     mockResolveLegacyV0BillingAttribution.mockResolvedValue(ATTRIBUTION)
     mockSerializeBillingAttributionHeader.mockReturnValue('serialized-attribution')
@@ -193,6 +186,10 @@ describe('POST /api/copilot/api-keys/validate billing protocols', () => {
     })
   })
 
+  afterAll(() => {
+    resetDbChainMock()
+  })
+
   it('keeps the exact old-Go validate bodies contract-compatible', () => {
     expect(validateCopilotApiKeyBodySchema.safeParse(OLD_GO_HOSTED_VALIDATE_BODY).success).toBe(
       true

apps/sim/app/api/copilot/chat/update-messages/route.test.ts

@@ -3,52 +3,24 @@
  *
  * @vitest-environment node
  */
-import { authMockFns } from '@sim/testing'
+import {
+  authMockFns,
+  dbChainMockFns,
+  queueTableRows,
+  resetDbChainMock,
+  schemaMock,
+} from '@sim/testing'
 import { NextRequest } from 'next/server'
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
-
-const {
-  mockSelect,
-  mockFrom,
-  mockWhere,
-  mockLimit,
-  mockUpdate,
-  mockSet,
-  mockUpdateWhere,
-  mockReturning,
-  mockReplaceCopilotChatMessages,
-} = vi.hoisted(() => ({
-  mockSelect: vi.fn(),
-  mockFrom: vi.fn(),
-  mockWhere: vi.fn(),
-  mockLimit: vi.fn(),
-  mockUpdate: vi.fn(),
-  mockSet: vi.fn(),
-  mockUpdateWhere: vi.fn(),
-  mockReturning: vi.fn(),
-  mockReplaceCopilotChatMessages: vi.fn(),
-}))
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
-vi.mock('@sim/db', () => ({
-  db: {
-    select: mockSelect,
-    update: mockUpdate,
-    transaction: async (
-      cb: (tx: { update: typeof mockUpdate; select: typeof mockSelect }) => unknown
-    ) => cb({ update: mockUpdate, select: mockSelect }),
-  },
+const { mockReplaceCopilotChatMessages } = vi.hoisted(() => ({
+  mockReplaceCopilotChatMessages: vi.fn(),
 }))
 
 vi.mock('@/lib/copilot/chat/messages-store', () => ({
   replaceCopilotChatMessages: mockReplaceCopilotChatMessages,
 }))
 
-vi.mock('drizzle-orm', () => ({
-  and: vi.fn((...conditions: unknown[]) => ({ conditions, type: 'and' })),
-  eq: vi.fn((field: unknown, value: unknown) => ({ field, value, type: 'eq' })),
-  isNull: vi.fn((field: unknown) => ({ field, type: 'isNull' })),
-}))
-
 import { POST } from '@/app/api/copilot/chat/update-messages/route'
 
 function createMockRequest(method: string, body: Record<string, unknown>): NextRequest {
@@ -62,21 +34,15 @@ function createMockRequest(method: string, body: Record<string, unknown>): NextR
 describe('Copilot Chat Update Messages API Route', () => {
   beforeEach(() => {
     vi.clearAllMocks()
+    resetDbChainMock()
 
     authMockFns.mockGetSession.mockResolvedValue(null)
 
-    mockSelect.mockReturnValue({ from: mockFrom })
-    mockFrom.mockReturnValue({ where: mockWhere })
-    mockWhere.mockReturnValue({ limit: mockLimit })
-    mockLimit.mockResolvedValue([])
-    mockUpdate.mockReturnValue({ set: mockSet })
-    mockSet.mockReturnValue({ where: mockUpdateWhere })
-    mockUpdateWhere.mockReturnValue({ returning: mockReturning })
-    mockReturning.mockResolvedValue([{ model: 'gpt-4' }])
+    dbChainMockFns.returning.mockResolvedValue([{ model: 'gpt-4' }])
   })
 
-  afterEach(() => {
-    vi.restoreAllMocks()
+  afterAll(() => {
+    resetDbChainMock()
   })
 
   describe('POST', () => {
@@ -181,7 +147,7 @@ describe('Copilot Chat Update Messages API Route', () => {
     it('should return 404 when chat is not found', async () => {
       authMockFns.mockGetSession.mockResolvedValue({ user: { id: 'user-123' } })
 
-      mockLimit.mockResolvedValueOnce([])
+      queueTableRows(schemaMock.copilotChats, [])
 
       const req = createMockRequest('POST', {
         chatId: 'non-existent-chat',
@@ -205,7 +171,7 @@ describe('Copilot Chat Update Messages API Route', () => {
     it('should return 404 when chat belongs to different user', async () => {
       authMockFns.mockGetSession.mockResolvedValue({ user: { id: 'user-123' } })
 
-      mockLimit.mockResolvedValueOnce([])
+      queueTableRows(schemaMock.copilotChats, [])
 
       const req = createMockRequest('POST', {
         chatId: 'other-user-chat',
@@ -234,7 +200,7 @@ describe('Copilot Chat Update Messages API Route', () => {
         userId: 'user-123',
         messages: [],
       }
-      mockLimit.mockResolvedValueOnce([existingChat])
+      queueTableRows(schemaMock.copilotChats, [existingChat])
 
       const messages = [
         {
@@ -265,9 +231,9 @@ describe('Copilot Chat Update Messages API Route', () => {
         messageCount: 2,
       })
 
-      expect(mockSelect).toHaveBeenCalled()
-      expect(mockUpdate).toHaveBeenCalled()
-      expect(mockSet).toHaveBeenCalledWith({ updatedAt: expect.any(Date) })
+      expect(dbChainMockFns.select).toHaveBeenCalled()
+      expect(dbChainMockFns.update).toHaveBeenCalled()
+      expect(dbChainMockFns.set).toHaveBeenCalledWith({ updatedAt: expect.any(Date) })
       expect(mockReplaceCopilotChatMessages).toHaveBeenCalledWith(
         'chat-123',
         messages,
@@ -284,7 +250,7 @@ describe('Copilot Chat Update Messages API Route', () => {
         userId: 'user-123',
         messages: [],
       }
-      mockLimit.mockResolvedValueOnce([existingChat])
+      queueTableRows(schemaMock.copilotChats, [existingChat])
 
       const messages = [
         {
@@ -328,7 +294,7 @@ describe('Copilot Chat Update Messages API Route', () => {
         messageCount: 2,
       })
 
-      expect(mockSet).toHaveBeenCalledWith({ updatedAt: expect.any(Date) })
+      expect(dbChainMockFns.set).toHaveBeenCalledWith({ updatedAt: expect.any(Date) })
       expect(mockReplaceCopilotChatMessages).toHaveBeenCalledWith(
         'chat-456',
         [
@@ -373,7 +339,7 @@ describe('Copilot Chat Update Messages API Route', () => {
         userId: 'user-123',
         messages: [],
       }
-      mockLimit.mockResolvedValueOnce([existingChat])
+      queueTableRows(schemaMock.copilotChats, [existingChat])
 
       const req = createMockRequest('POST', {
         chatId: 'chat-789',
@@ -389,7 +355,7 @@ describe('Copilot Chat Update Messages API Route', () => {
         messageCount: 0,
       })
 
-      expect(mockSet).toHaveBeenCalledWith({ updatedAt: expect.any(Date) })
+      expect(dbChainMockFns.set).toHaveBeenCalledWith({ updatedAt: expect.any(Date) })
       expect(mockReplaceCopilotChatMessages).toHaveBeenCalledWith(
         'chat-789',
         [],
@@ -401,7 +367,7 @@ describe('Copilot Chat Update Messages API Route', () => {
     it('should handle database errors during chat lookup', async () => {
       authMockFns.mockGetSession.mockResolvedValue({ user: { id: 'user-123' } })
 
-      mockLimit.mockRejectedValueOnce(new Error('Database connection failed'))
+      dbChainMockFns.limit.mockRejectedValueOnce(new Error('Database connection failed'))
 
       const req = createMockRequest('POST', {
         chatId: 'chat-123',
@@ -430,11 +396,9 @@ describe('Copilot Chat Update Messages API Route', () => {
         userId: 'user-123',
         messages: [],
       }
-      mockLimit.mockResolvedValueOnce([existingChat])
+      queueTableRows(schemaMock.copilotChats, [existingChat])
 
-      mockSet.mockReturnValueOnce({
-        where: vi.fn().mockRejectedValue(new Error('Update operation failed')),
-      })
+      dbChainMockFns.returning.mockRejectedValueOnce(new Error('Update operation failed'))
 
       const req = createMockRequest('POST', {
         chatId: 'chat-123',
@@ -481,7 +445,7 @@ describe('Copilot Chat Update Messages API Route', () => {
         userId: 'user-123',
         messages: [],
       }
-      mockLimit.mockResolvedValueOnce([existingChat])
+      queueTableRows(schemaMock.copilotChats, [existingChat])
 
       const messages = Array.from({ length: 100 }, (_, i) => ({
         id: `msg-${i + 1}`,
@@ -504,7 +468,7 @@ describe('Copilot Chat Update Messages API Route', () => {
         messageCount: 100,
       })
 
-      expect(mockSet).toHaveBeenCalledWith({ updatedAt: expect.any(Date) })
+      expect(dbChainMockFns.set).toHaveBeenCalledWith({ updatedAt: expect.any(Date) })
       expect(mockReplaceCopilotChatMessages).toHaveBeenCalledWith(
         'chat-large',
         messages,
@@ -521,7 +485,7 @@ describe('Copilot Chat Update Messages API Route', () => {
         userId: 'user-123',
         messages: [],
       }
-      mockLimit.mockResolvedValueOnce([existingChat])
+      queueTableRows(schemaMock.copilotChats, [existingChat])
 
       const messages = [
         {

apps/sim/app/api/copilot/chats/route.test.ts

@@ -3,32 +3,15 @@
  *
  * @vitest-environment node
  */
-import { copilotHttpMock, copilotHttpMockFns } from '@sim/testing'
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
-
-const { mockSelectDistinctOn, mockFrom, mockLeftJoin, mockWhere, mockOrderBy } = vi.hoisted(() => ({
-  mockSelectDistinctOn: vi.fn(),
-  mockFrom: vi.fn(),
-  mockLeftJoin: vi.fn(),
-  mockWhere: vi.fn(),
-  mockOrderBy: vi.fn(),
-}))
-
-vi.mock('@sim/db', () => ({
-  db: {
-    selectDistinctOn: mockSelectDistinctOn,
-  },
-}))
-
-vi.mock('drizzle-orm', () => ({
-  and: vi.fn((...conditions: unknown[]) => ({ conditions, type: 'and' })),
-  eq: vi.fn((field: unknown, value: unknown) => ({ field, value, type: 'eq' })),
-  or: vi.fn((...conditions: unknown[]) => ({ conditions, type: 'or' })),
-  inArray: vi.fn((field: unknown, values: unknown) => ({ field, values, type: 'inArray' })),
-  isNull: vi.fn((field: unknown) => ({ field, type: 'isNull' })),
-  desc: vi.fn((field: unknown) => ({ field, type: 'desc' })),
-  sql: vi.fn(),
-}))
+import {
+  copilotHttpMock,
+  copilotHttpMockFns,
+  dbChainMockFns,
+  queueTableRows,
+  resetDbChainMock,
+  schemaMock,
+} from '@sim/testing'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
 vi.mock('@/lib/copilot/request/http', () => copilotHttpMock)
 
@@ -45,16 +28,11 @@ import { GET } from '@/app/api/copilot/chats/route'
 describe('Copilot Chats List API Route', () => {
   beforeEach(() => {
     vi.clearAllMocks()
-
-    mockSelectDistinctOn.mockReturnValue({ from: mockFrom })
-    mockFrom.mockReturnValue({ leftJoin: mockLeftJoin })
-    mockLeftJoin.mockReturnValue({ leftJoin: mockLeftJoin, where: mockWhere })
-    mockWhere.mockReturnValue({ orderBy: mockOrderBy })
-    mockOrderBy.mockResolvedValue([])
+    resetDbChainMock()
   })
 
-  afterEach(() => {
-    vi.restoreAllMocks()
+  afterAll(() => {
+    resetDbChainMock()
   })
 
   describe('GET', () => {
@@ -78,7 +56,7 @@ describe('Copilot Chats List API Route', () => {
         isAuthenticated: true,
       })
 
-      mockOrderBy.mockResolvedValueOnce([])
+      queueTableRows(schemaMock.copilotChats, [])
 
       const request = new Request('http://localhost:3000/api/copilot/chats')
       const response = await GET(request as any)
@@ -111,7 +89,7 @@ describe('Copilot Chats List API Route', () => {
           updatedAt: new Date('2024-01-01'),
         },
       ]
-      mockOrderBy.mockResolvedValueOnce(mockChats)
+      queueTableRows(schemaMock.copilotChats, mockChats)
 
       const request = new Request('http://localhost:3000/api/copilot/chats')
       const response = await GET(request as any)
@@ -151,7 +129,7 @@ describe('Copilot Chats List API Route', () => {
           updatedAt: new Date('2024-01-01'),
         },
       ]
-      mockOrderBy.mockResolvedValueOnce(mockChats)
+      queueTableRows(schemaMock.copilotChats, mockChats)
 
       const request = new Request('http://localhost:3000/api/copilot/chats')
       const response = await GET(request as any)
@@ -176,7 +154,7 @@ describe('Copilot Chats List API Route', () => {
           updatedAt: new Date('2024-01-01'),
         },
       ]
-      mockOrderBy.mockResolvedValueOnce(mockChats)
+      queueTableRows(schemaMock.copilotChats, mockChats)
 
       const request = new Request('http://localhost:3000/api/copilot/chats')
       const response = await GET(request as any)
@@ -192,7 +170,7 @@ describe('Copilot Chats List API Route', () => {
         isAuthenticated: true,
       })
 
-      mockOrderBy.mockRejectedValueOnce(new Error('Database connection failed'))
+      dbChainMockFns.orderBy.mockRejectedValueOnce(new Error('Database connection failed'))
 
       const request = new Request('http://localhost:3000/api/copilot/chats')
       const response = await GET(request as any)
@@ -216,13 +194,13 @@ describe('Copilot Chats List API Route', () => {
           updatedAt: new Date('2024-01-01'),
         },
       ]
-      mockOrderBy.mockResolvedValueOnce(mockChats)
+      queueTableRows(schemaMock.copilotChats, mockChats)
 
       const request = new Request('http://localhost:3000/api/copilot/chats')
       await GET(request as any)
 
-      expect(mockSelectDistinctOn).toHaveBeenCalled()
-      expect(mockWhere).toHaveBeenCalled()
+      expect(dbChainMockFns.selectDistinctOn).toHaveBeenCalled()
+      expect(dbChainMockFns.where).toHaveBeenCalled()
     })
 
     it('should return 401 when userId is null despite isAuthenticated being true', async () => {