From 182f6d057f49b1f9b0737391af203632d80688ed Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 22 Jul 2026 14:55:13 -0700 Subject: [PATCH 1/2] =?UTF-8?q?improvement(tests):=20db-mock=20migration?= =?UTF-8?q?=20tranche=202=20=E2=80=94=20copilot,=20mothership,=20workspace?= =?UTF-8?q?s,=20connectors,=20mcp?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../copilot/api-keys/validate/route.test.ts | 19 ++- .../chat/update-messages/route.test.ts | 96 ++++--------- apps/sim/app/api/copilot/chats/route.test.ts | 62 +++----- .../copilot/checkpoints/revert/route.test.ts | 132 +++++++----------- .../app/api/copilot/checkpoints/route.test.ts | 84 ++++------- .../[connectorId]/documents/route.test.ts | 56 +++----- .../connectors/[connectorId]/route.test.ts | 72 ++++------ .../[connectorId]/sync/route.test.ts | 49 +++---- .../knowledge/[id]/connectors/route.test.ts | 53 +++---- .../mcp/servers/[id]/refresh/route.test.ts | 66 ++++----- apps/sim/app/api/mcp/servers/route.test.ts | 14 +- .../chats/[chatId]/fork/route.test.ts | 122 ++++------------ .../chats/[chatId]/restore/route.test.ts | 79 +++-------- .../mothership/chats/[chatId]/route.test.ts | 59 ++------ .../api/mothership/chats/read/route.test.ts | 49 ++----- .../app/api/mothership/chats/route.test.ts | 71 +++------- .../workspaces/[id]/api-keys/route.test.ts | 25 ++-- .../workspaces/[id]/byok-keys/route.test.ts | 125 +++++++---------- .../fork/excluded-workflows/route.test.ts | 39 +++--- .../api/workspaces/invitations/route.test.ts | 41 ++---- .../tools/handlers/deployment/manage.test.ts | 50 ++----- .../copilot/tools/handlers/resources.test.ts | 6 - .../copilot/tools/handlers/vfs-mutate.test.ts | 29 ++-- .../tools/handlers/workflow/mutations.test.ts | 7 +- apps/sim/lib/mcp/service-pool.test.ts | 56 ++++---- apps/sim/lib/mcp/service.test.ts | 79 +++++------ packages/testing/src/mocks/schema.mock.ts | 2 + 27 files changed, 517 insertions(+), 1025 deletions(-) diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts index ea9c68b2b58..7e8b86e5a87 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts +++ b/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, headers: Record 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 diff --git a/apps/sim/app/api/copilot/chat/update-messages/route.test.ts b/apps/sim/app/api/copilot/chat/update-messages/route.test.ts index b6ec71a7a9a..ab842b87480 100644 --- a/apps/sim/app/api/copilot/chat/update-messages/route.test.ts +++ b/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): NextRequest { @@ -62,21 +34,15 @@ function createMockRequest(method: string, body: Record): 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 = [ { diff --git a/apps/sim/app/api/copilot/chats/route.test.ts b/apps/sim/app/api/copilot/chats/route.test.ts index 2ce59e2a41d..b0d55f718a6 100644 --- a/apps/sim/app/api/copilot/chats/route.test.ts +++ b/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 () => { diff --git a/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts b/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts index 9ce957da1ef..4918057f055 100644 --- a/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts +++ b/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts @@ -3,25 +3,19 @@ * * @vitest-environment node */ -import { authMockFns, workflowAuthzMockFns, workflowsUtilsMock } from '@sim/testing' +import { + authMockFns, + dbChainMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, + workflowAuthzMockFns, + workflowsUtilsMock, +} from '@sim/testing' import { NextRequest } from 'next/server' -import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' - -const { - mockSelect, - mockFrom, - mockWhere, - mockThen, - mockDelete, - mockDeleteWhere, - mockGetAccessibleCopilotChat, -} = vi.hoisted(() => ({ - mockSelect: vi.fn(), - mockFrom: vi.fn(), - mockWhere: vi.fn(), - mockThen: vi.fn(), - mockDelete: vi.fn(), - mockDeleteWhere: vi.fn(), +import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockGetAccessibleCopilotChat } = vi.hoisted(() => ({ mockGetAccessibleCopilotChat: vi.fn(), })) @@ -39,28 +33,12 @@ vi.mock('@/lib/copilot/chat/lifecycle', () => ({ getAccessibleCopilotChatAuth: mockGetAccessibleCopilotChat, })) -vi.mock('@sim/db', () => ({ - db: { - select: mockSelect, - delete: mockDelete, - }, -})) - -vi.mock('drizzle-orm', () => ({ - and: vi.fn((...conditions: unknown[]) => ({ conditions, type: 'and' })), - eq: vi.fn((field: unknown, value: unknown) => ({ field, value, type: 'eq' })), -})) - import { POST } from '@/app/api/copilot/checkpoints/revert/route' describe('Copilot Checkpoints Revert API Route', () => { - /** Queued results for successive `.then()` calls in the db select chain */ - let thenResults: unknown[] - beforeEach(() => { vi.clearAllMocks() - - thenResults = [] + resetDbChainMock() authMockFns.mockGetSession.mockResolvedValue(null) @@ -69,24 +47,6 @@ describe('Copilot Checkpoints Revert API Route', () => { status: 200, }) - mockSelect.mockReturnValue({ from: mockFrom }) - mockFrom.mockReturnValue({ where: mockWhere }) - mockWhere.mockReturnValue({ then: mockThen }) - - // Drizzle's .then() is a thenable: it receives a callback like (rows) => rows[0]. - // We invoke the callback with our mock rows array so the route gets the expected value. - mockThen.mockImplementation((callback: (rows: unknown[]) => unknown) => { - const result = thenResults.shift() - if (result instanceof Error) { - return Promise.reject(result) - } - const rows = result === undefined ? [] : [result] - return Promise.resolve(callback(rows)) - }) - - // Mock delete chain - mockDelete.mockReturnValue({ where: mockDeleteWhere }) - mockDeleteWhere.mockResolvedValue(undefined) mockGetAccessibleCopilotChat.mockResolvedValue({ id: 'chat-123', userId: 'user-123' }) global.fetch = vi.fn() @@ -114,10 +74,13 @@ describe('Copilot Checkpoints Revert API Route', () => { }) afterEach(() => { - vi.clearAllMocks() vi.restoreAllMocks() }) + afterAll(() => { + resetDbChainMock() + }) + /** Helper to set authenticated state */ function setAuthenticated(user = { id: 'user-123', email: 'test@example.com' }) { authMockFns.mockGetSession.mockResolvedValue({ user }) @@ -180,8 +143,7 @@ describe('Copilot Checkpoints Revert API Route', () => { it('should return 404 when checkpoint is not found', async () => { setAuthenticated() - // Mock checkpoint not found - thenResults.push(undefined) + queueTableRows(schemaMock.workflowCheckpoints, []) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { method: 'POST', @@ -199,8 +161,7 @@ describe('Copilot Checkpoints Revert API Route', () => { it('should return 404 when checkpoint belongs to different user', async () => { setAuthenticated() - // Mock checkpoint not found (due to user mismatch in query) - thenResults.push(undefined) + queueTableRows(schemaMock.workflowCheckpoints, []) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { method: 'POST', @@ -225,8 +186,8 @@ describe('Copilot Checkpoints Revert API Route', () => { workflowState: { blocks: {}, edges: [] }, } - thenResults.push(mockCheckpoint) // Checkpoint found - thenResults.push(undefined) // Workflow not found + queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) + queueTableRows(schemaMock.workflow, []) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { method: 'POST', @@ -256,8 +217,8 @@ describe('Copilot Checkpoints Revert API Route', () => { userId: 'different-user', } - thenResults.push(mockCheckpoint) // Checkpoint found - thenResults.push(mockWorkflow) // Workflow found but different user + queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) + queueTableRows(schemaMock.workflow, [mockWorkflow]) workflowAuthzMockFns.mockAuthorizeWorkflowByWorkspacePermission.mockResolvedValueOnce({ allowed: false, @@ -298,8 +259,8 @@ describe('Copilot Checkpoints Revert API Route', () => { userId: 'user-123', } - thenResults.push(mockCheckpoint) // Checkpoint found - thenResults.push(mockWorkflow) // Workflow found + queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) + queueTableRows(schemaMock.workflow, [mockWorkflow]) ;(global.fetch as any).mockResolvedValue({ ok: true, @@ -380,8 +341,8 @@ describe('Copilot Checkpoints Revert API Route', () => { userId: 'user-123', } - thenResults.push(mockCheckpoint) - thenResults.push(mockWorkflow) + queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) + queueTableRows(schemaMock.workflow, [mockWorkflow]) ;(global.fetch as any).mockResolvedValue({ ok: true, @@ -422,8 +383,8 @@ describe('Copilot Checkpoints Revert API Route', () => { userId: 'user-123', } - thenResults.push(mockCheckpoint) - thenResults.push(mockWorkflow) + queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) + queueTableRows(schemaMock.workflow, [mockWorkflow]) ;(global.fetch as any).mockResolvedValue({ ok: true, @@ -464,8 +425,8 @@ describe('Copilot Checkpoints Revert API Route', () => { userId: 'user-123', } - thenResults.push(mockCheckpoint) - thenResults.push(mockWorkflow) + queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) + queueTableRows(schemaMock.workflow, [mockWorkflow]) ;(global.fetch as any).mockResolvedValue({ ok: true, @@ -509,8 +470,8 @@ describe('Copilot Checkpoints Revert API Route', () => { userId: 'user-123', } - thenResults.push(mockCheckpoint) - thenResults.push(mockWorkflow) + queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) + queueTableRows(schemaMock.workflow, [mockWorkflow]) ;(global.fetch as any).mockResolvedValue({ ok: false, @@ -533,8 +494,9 @@ describe('Copilot Checkpoints Revert API Route', () => { it('should handle database errors during checkpoint lookup', async () => { setAuthenticated() - // Mock database error - thenResults.push(new Error('Database connection failed')) + dbChainMockFns.where.mockReturnValueOnce( + Promise.reject(new Error('Database connection failed')) + ) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { method: 'POST', @@ -559,8 +521,10 @@ describe('Copilot Checkpoints Revert API Route', () => { workflowState: { blocks: {}, edges: [] }, } - thenResults.push(mockCheckpoint) // Checkpoint found - thenResults.push(new Error('Database error during workflow lookup')) // Workflow lookup fails + dbChainMockFns.where.mockReturnValueOnce(Promise.resolve([mockCheckpoint])) + dbChainMockFns.where.mockReturnValueOnce( + Promise.reject(new Error('Database error during workflow lookup')) + ) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { method: 'POST', @@ -590,8 +554,8 @@ describe('Copilot Checkpoints Revert API Route', () => { userId: 'user-123', } - thenResults.push(mockCheckpoint) - thenResults.push(mockWorkflow) + queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) + queueTableRows(schemaMock.workflow, [mockWorkflow]) ;(global.fetch as any).mockRejectedValue(new Error('Network error')) @@ -641,8 +605,8 @@ describe('Copilot Checkpoints Revert API Route', () => { userId: 'user-123', } - thenResults.push(mockCheckpoint) - thenResults.push(mockWorkflow) + queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) + queueTableRows(schemaMock.workflow, [mockWorkflow]) ;(global.fetch as any).mockResolvedValue({ ok: true, @@ -690,8 +654,8 @@ describe('Copilot Checkpoints Revert API Route', () => { userId: 'user-123', } - thenResults.push(mockCheckpoint) - thenResults.push(mockWorkflow) + queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) + queueTableRows(schemaMock.workflow, [mockWorkflow]) ;(global.fetch as any).mockResolvedValue({ ok: true, @@ -758,8 +722,8 @@ describe('Copilot Checkpoints Revert API Route', () => { userId: 'user-123', } - thenResults.push(mockCheckpoint) - thenResults.push(mockWorkflow) + queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) + queueTableRows(schemaMock.workflow, [mockWorkflow]) ;(global.fetch as any).mockResolvedValue({ ok: true, diff --git a/apps/sim/app/api/copilot/checkpoints/route.test.ts b/apps/sim/app/api/copilot/checkpoints/route.test.ts index f7dc748786b..8ee9df0bc5d 100644 --- a/apps/sim/app/api/copilot/checkpoints/route.test.ts +++ b/apps/sim/app/api/copilot/checkpoints/route.test.ts @@ -3,43 +3,20 @@ * * @vitest-environment node */ -import { authMockFns, workflowAuthzMockFns, workflowsUtilsMock } from '@sim/testing' +import { + authMockFns, + dbChainMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, + workflowAuthzMockFns, + workflowsUtilsMock, +} from '@sim/testing' import { NextRequest } from 'next/server' -import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' - -const { - mockSelect, - mockFrom, - mockWhere, - mockLimit, - mockOrderBy, - mockInsert, - mockValues, - mockReturning, - mockGetAccessibleCopilotChat, -} = vi.hoisted(() => ({ - mockSelect: vi.fn(), - mockFrom: vi.fn(), - mockWhere: vi.fn(), - mockLimit: vi.fn(), - mockOrderBy: vi.fn(), - mockInsert: vi.fn(), - mockValues: vi.fn(), - mockReturning: vi.fn(), - mockGetAccessibleCopilotChat: vi.fn(), -})) - -vi.mock('@sim/db', () => ({ - db: { - select: mockSelect, - insert: mockInsert, - }, -})) +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' -vi.mock('drizzle-orm', () => ({ - and: vi.fn((...conditions: unknown[]) => ({ conditions, type: 'and' })), - eq: vi.fn((field: unknown, value: unknown) => ({ field, value, type: 'eq' })), - desc: vi.fn((field: unknown) => ({ field, type: 'desc' })), +const { mockGetAccessibleCopilotChat } = vi.hoisted(() => ({ + mockGetAccessibleCopilotChat: vi.fn(), })) vi.mock('@/lib/copilot/chat/lifecycle', () => ({ @@ -62,19 +39,10 @@ function createMockRequest(method: string, body: Record): NextR describe('Copilot Checkpoints API Route', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() authMockFns.mockGetSession.mockResolvedValue(null) - mockSelect.mockReturnValue({ from: mockFrom }) - mockFrom.mockReturnValue({ where: mockWhere }) - mockWhere.mockReturnValue({ - orderBy: mockOrderBy, - limit: mockLimit, - }) - mockOrderBy.mockResolvedValue([]) - mockLimit.mockResolvedValue([]) - mockInsert.mockReturnValue({ values: mockValues }) - mockValues.mockReturnValue({ returning: mockReturning }) mockGetAccessibleCopilotChat.mockResolvedValue({ id: 'chat-123', userId: 'user-123', @@ -85,8 +53,8 @@ describe('Copilot Checkpoints API Route', () => { }) }) - afterEach(() => { - vi.restoreAllMocks() + afterAll(() => { + resetDbChainMock() }) describe('POST', () => { @@ -165,7 +133,7 @@ describe('Copilot Checkpoints API Route', () => { createdAt: new Date('2024-01-01'), updatedAt: new Date('2024-01-01'), } - mockReturning.mockResolvedValue([checkpoint]) + dbChainMockFns.returning.mockResolvedValueOnce([checkpoint]) const workflowState = { blocks: [], connections: [] } const req = createMockRequest('POST', { @@ -192,8 +160,8 @@ describe('Copilot Checkpoints API Route', () => { }, }) - expect(mockInsert).toHaveBeenCalled() - expect(mockValues).toHaveBeenCalledWith({ + expect(dbChainMockFns.insert).toHaveBeenCalled() + expect(dbChainMockFns.values).toHaveBeenCalledWith({ userId: 'user-123', workflowId: 'workflow-123', chatId: 'chat-123', @@ -214,7 +182,7 @@ describe('Copilot Checkpoints API Route', () => { createdAt: new Date('2024-01-01'), updatedAt: new Date('2024-01-01'), } - mockReturning.mockResolvedValue([checkpoint]) + dbChainMockFns.returning.mockResolvedValueOnce([checkpoint]) const workflowState = { blocks: [] } const req = createMockRequest('POST', { @@ -234,7 +202,7 @@ describe('Copilot Checkpoints API Route', () => { it('should handle database errors during checkpoint creation', async () => { authMockFns.mockGetSession.mockResolvedValue({ user: { id: 'user-123' } }) - mockReturning.mockRejectedValue(new Error('Database insert failed')) + dbChainMockFns.returning.mockRejectedValueOnce(new Error('Database insert failed')) const req = createMockRequest('POST', { workflowId: 'workflow-123', @@ -317,7 +285,7 @@ describe('Copilot Checkpoints API Route', () => { }, ] - mockOrderBy.mockResolvedValue(mockCheckpoints) + queueTableRows(schemaMock.workflowCheckpoints, mockCheckpoints) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints?chatId=chat-123') @@ -349,15 +317,15 @@ describe('Copilot Checkpoints API Route', () => { ], }) - expect(mockSelect).toHaveBeenCalled() - expect(mockWhere).toHaveBeenCalled() - expect(mockOrderBy).toHaveBeenCalled() + expect(dbChainMockFns.select).toHaveBeenCalled() + expect(dbChainMockFns.where).toHaveBeenCalled() + expect(dbChainMockFns.orderBy).toHaveBeenCalled() }) it('should handle database errors when fetching checkpoints', async () => { authMockFns.mockGetSession.mockResolvedValue({ user: { id: 'user-123' } }) - mockOrderBy.mockRejectedValue(new Error('Database query failed')) + dbChainMockFns.orderBy.mockRejectedValueOnce(new Error('Database query failed')) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints?chatId=chat-123') @@ -371,7 +339,7 @@ describe('Copilot Checkpoints API Route', () => { it('should return empty array when no checkpoints found', async () => { authMockFns.mockGetSession.mockResolvedValue({ user: { id: 'user-123' } }) - mockOrderBy.mockResolvedValue([]) + queueTableRows(schemaMock.workflowCheckpoints, []) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints?chatId=chat-123') diff --git a/apps/sim/app/api/knowledge/[id]/connectors/[connectorId]/documents/route.test.ts b/apps/sim/app/api/knowledge/[id]/connectors/[connectorId]/documents/route.test.ts index 4f39294d055..bcbf905bb27 100644 --- a/apps/sim/app/api/knowledge/[id]/connectors/[connectorId]/documents/route.test.ts +++ b/apps/sim/app/api/knowledge/[id]/connectors/[connectorId]/documents/route.test.ts @@ -4,31 +4,20 @@ import { auditMock, createMockRequest, + dbChainMock, + dbChainMockFns, hybridAuthMockFns, knowledgeApiUtilsMock, knowledgeApiUtilsMockFns, requestUtilsMockFns, + resetDbChainMock, } from '@sim/testing' -import { 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(), - orderBy: vi.fn().mockResolvedValue([]), - limit: vi.fn().mockResolvedValue([]), - update: vi.fn().mockReturnThis(), - set: vi.fn().mockReturnThis(), - returning: vi.fn().mockResolvedValue([]), - } - return { mockDbChain: chain } -}) +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const mockCheckAccess = knowledgeApiUtilsMockFns.mockCheckKnowledgeBaseAccess const mockCheckWriteAccess = knowledgeApiUtilsMockFns.mockCheckKnowledgeBaseWriteAccess -vi.mock('@sim/db', () => ({ db: mockDbChain })) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/app/api/knowledge/utils', () => knowledgeApiUtilsMock) vi.mock('@sim/audit', () => auditMock) @@ -39,15 +28,12 @@ describe('Connector Documents API Route', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() requestUtilsMockFns.mockGenerateRequestId.mockReturnValue('test-req-id') - mockDbChain.select.mockReturnThis() - mockDbChain.from.mockReturnThis() - mockDbChain.where.mockReturnThis() - mockDbChain.orderBy.mockResolvedValue([]) - mockDbChain.limit.mockResolvedValue([]) - mockDbChain.update.mockReturnThis() - mockDbChain.set.mockReturnThis() - mockDbChain.returning.mockResolvedValue([]) + }) + + afterAll(() => { + resetDbChainMock() }) describe('GET', () => { @@ -69,7 +55,7 @@ describe('Connector Documents API Route', () => { userId: 'user-1', }) mockCheckAccess.mockResolvedValue({ hasAccess: true }) - mockDbChain.limit.mockResolvedValueOnce([]) + dbChainMockFns.limit.mockResolvedValueOnce([]) const req = createMockRequest('GET') const response = await GET(req as never, { params: mockParams }) @@ -85,8 +71,8 @@ describe('Connector Documents API Route', () => { mockCheckAccess.mockResolvedValue({ hasAccess: true }) const doc = { id: 'doc-1', filename: 'test.txt', userExcluded: false } - mockDbChain.limit.mockResolvedValueOnce([{ id: 'conn-456' }]) - mockDbChain.orderBy.mockResolvedValueOnce([doc]) + dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'conn-456' }]) + dbChainMockFns.orderBy.mockResolvedValueOnce([doc]) const url = 'http://localhost/api/knowledge/kb-123/connectors/conn-456/documents' const req = createMockRequest('GET', undefined, undefined, url) @@ -106,8 +92,8 @@ describe('Connector Documents API Route', () => { }) mockCheckAccess.mockResolvedValue({ hasAccess: true }) - mockDbChain.limit.mockResolvedValueOnce([{ id: 'conn-456' }]) - mockDbChain.orderBy + dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'conn-456' }]) + dbChainMockFns.orderBy .mockResolvedValueOnce([{ id: 'doc-1', userExcluded: false }]) .mockResolvedValueOnce([{ id: 'doc-2', userExcluded: true }]) @@ -143,7 +129,7 @@ describe('Connector Documents API Route', () => { userId: 'user-1', }) mockCheckWriteAccess.mockResolvedValue({ hasAccess: true }) - mockDbChain.limit.mockResolvedValueOnce([{ id: 'conn-456' }]) + dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'conn-456' }]) const req = createMockRequest('PATCH', { documentIds: [] }) const response = await PATCH(req as never, { params: mockParams }) @@ -157,7 +143,7 @@ describe('Connector Documents API Route', () => { userId: 'user-1', }) mockCheckWriteAccess.mockResolvedValue({ hasAccess: true }) - mockDbChain.limit.mockResolvedValueOnce([]) + dbChainMockFns.limit.mockResolvedValueOnce([]) const req = createMockRequest('PATCH', { operation: 'restore', documentIds: ['doc-1'] }) const response = await PATCH(req as never, { params: mockParams }) @@ -176,8 +162,8 @@ describe('Connector Documents API Route', () => { hasAccess: true, knowledgeBase: { workspaceId: 'ws-1', name: 'Test KB' }, }) - mockDbChain.limit.mockResolvedValueOnce([{ id: 'conn-456' }]) - mockDbChain.returning.mockResolvedValueOnce([{ id: 'doc-1' }]) + dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'conn-456' }]) + dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'doc-1' }]) const req = createMockRequest('PATCH', { operation: 'restore', documentIds: ['doc-1'] }) const response = await PATCH(req as never, { params: mockParams }) @@ -198,8 +184,8 @@ describe('Connector Documents API Route', () => { hasAccess: true, knowledgeBase: { workspaceId: 'ws-1', name: 'Test KB' }, }) - mockDbChain.limit.mockResolvedValueOnce([{ id: 'conn-456' }]) - mockDbChain.returning.mockResolvedValueOnce([{ id: 'doc-2' }, { id: 'doc-3' }]) + dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'conn-456' }]) + dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'doc-2' }, { id: 'doc-3' }]) const req = createMockRequest('PATCH', { operation: 'exclude', diff --git a/apps/sim/app/api/knowledge/[id]/connectors/[connectorId]/route.test.ts b/apps/sim/app/api/knowledge/[id]/connectors/[connectorId]/route.test.ts index 4f8f7e7fca4..4784c21943a 100644 --- a/apps/sim/app/api/knowledge/[id]/connectors/[connectorId]/route.test.ts +++ b/apps/sim/app/api/knowledge/[id]/connectors/[connectorId]/route.test.ts @@ -5,39 +5,26 @@ import { auditMock, authOAuthUtilsMock, createMockRequest, + dbChainMock, + dbChainMockFns, hybridAuthMockFns, knowledgeApiUtilsMock, knowledgeApiUtilsMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const { mockDbChain, mockHasWorkspaceLiveSyncAccess, mockValidateConfig } = vi.hoisted(() => { - const chain = { - select: vi.fn().mockReturnThis(), - from: vi.fn().mockReturnThis(), - where: vi.fn().mockReturnThis(), - orderBy: vi.fn().mockReturnThis(), - limit: vi.fn().mockResolvedValue([]), - execute: vi.fn().mockResolvedValue(undefined), - transaction: vi.fn(), - insert: vi.fn().mockReturnThis(), - values: vi.fn().mockResolvedValue(undefined), - update: vi.fn().mockReturnThis(), - delete: vi.fn().mockReturnThis(), - set: vi.fn().mockReturnThis(), - returning: vi.fn().mockResolvedValue([]), - } - return { - mockDbChain: chain, - mockHasWorkspaceLiveSyncAccess: vi.fn(), - mockValidateConfig: vi.fn(), - } -}) +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockHasWorkspaceLiveSyncAccess, mockValidateConfig } = vi.hoisted(() => ({ + mockHasWorkspaceLiveSyncAccess: vi.fn(), + mockValidateConfig: vi.fn(), +})) const mockCheckAccess = knowledgeApiUtilsMockFns.mockCheckKnowledgeBaseAccess const mockCheckWriteAccess = knowledgeApiUtilsMockFns.mockCheckKnowledgeBaseWriteAccess -vi.mock('@sim/db', () => ({ db: mockDbChain })) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/app/api/knowledge/utils', () => knowledgeApiUtilsMock) vi.mock('@/app/api/auth/oauth/utils', () => authOAuthUtilsMock) vi.mock('@/connectors/registry.server', () => ({ @@ -63,19 +50,11 @@ describe('Knowledge Connector By ID API Route', () => { beforeEach(() => { vi.clearAllMocks() - mockDbChain.select.mockReturnThis() - mockDbChain.from.mockReturnThis() - mockDbChain.where.mockReturnThis() - mockDbChain.orderBy.mockReturnThis() - mockDbChain.limit.mockResolvedValue([]) - mockDbChain.execute.mockResolvedValue(undefined) - mockDbChain.transaction.mockImplementation( - async (callback: (tx: typeof mockDbChain) => unknown) => callback(mockDbChain) - ) - mockDbChain.update.mockReturnThis() - mockDbChain.delete.mockReturnThis() - mockDbChain.set.mockReturnThis() - mockDbChain.returning.mockResolvedValue([]) + resetDbChainMock() + }) + + afterAll(() => { + resetDbChainMock() }) describe('GET', () => { @@ -110,7 +89,7 @@ describe('Knowledge Connector By ID API Route', () => { userId: 'user-1', }) mockCheckAccess.mockResolvedValue({ hasAccess: true }) - mockDbChain.limit.mockResolvedValueOnce([]) + dbChainMockFns.limit.mockResolvedValueOnce([]) const req = createMockRequest('GET') const response = await GET(req, { params: mockParams }) @@ -128,7 +107,7 @@ describe('Knowledge Connector By ID API Route', () => { const mockConnector = { id: 'conn-456', connectorType: 'jira', status: 'active' } const mockLogs = [{ id: 'log-1', status: 'completed' }] - mockDbChain.limit.mockResolvedValueOnce([mockConnector]).mockResolvedValueOnce(mockLogs) + dbChainMockFns.limit.mockResolvedValueOnce([mockConnector]).mockResolvedValueOnce(mockLogs) const req = createMockRequest('GET') const response = await GET(req, { params: mockParams }) @@ -175,7 +154,7 @@ describe('Knowledge Connector By ID API Route', () => { userId: 'user-1', }) mockCheckWriteAccess.mockResolvedValue({ hasAccess: true }) - mockDbChain.limit.mockResolvedValueOnce([]) + dbChainMockFns.limit.mockResolvedValueOnce([]) const req = createMockRequest('PATCH', { sourceConfig: { project: 'NEW' } }) const response = await PATCH(req, { params: mockParams }) @@ -197,7 +176,7 @@ describe('Knowledge Connector By ID API Route', () => { mockHasWorkspaceLiveSyncAccess.mockResolvedValue(true) const updatedConnector = { id: 'conn-456', status: 'paused', syncIntervalMinutes: 5 } - mockDbChain.limit.mockResolvedValueOnce([updatedConnector]) + dbChainMockFns.limit.mockResolvedValueOnce([updatedConnector]) const req = createMockRequest('PATCH', { status: 'paused', syncIntervalMinutes: 5 }) const response = await PATCH(req, { params: mockParams }) @@ -252,12 +231,9 @@ describe('Knowledge Connector By ID API Route', () => { hasAccess: true, knowledgeBase: { workspaceId: 'ws-1', name: 'Test KB' }, }) - mockDbChain.where - .mockReturnValueOnce(mockDbChain) - .mockResolvedValueOnce([{ id: 'doc-1', fileUrl: '/api/uploads/test.txt' }]) - .mockReturnValueOnce(mockDbChain) - mockDbChain.limit.mockResolvedValueOnce([{ id: 'conn-456', connectorType: 'jira' }]) - mockDbChain.returning.mockResolvedValueOnce([{ id: 'conn-456' }]) + dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'conn-456', connectorType: 'jira' }]) + queueTableRows(schemaMock.document, [{ id: 'doc-1', fileUrl: '/api/uploads/test.txt' }]) + dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'conn-456' }]) const req = createMockRequest('DELETE') const response = await DELETE(req, { params: mockParams }) diff --git a/apps/sim/app/api/knowledge/[id]/connectors/[connectorId]/sync/route.test.ts b/apps/sim/app/api/knowledge/[id]/connectors/[connectorId]/sync/route.test.ts index e84f5cd4da6..54ef8f3fc4c 100644 --- a/apps/sim/app/api/knowledge/[id]/connectors/[connectorId]/sync/route.test.ts +++ b/apps/sim/app/api/knowledge/[id]/connectors/[connectorId]/sync/route.test.ts @@ -4,33 +4,24 @@ import { auditMock, createMockRequest, + dbChainMock, + dbChainMockFns, hybridAuthMockFns, knowledgeApiUtilsMock, knowledgeApiUtilsMockFns, requestUtilsMockFns, + resetDbChainMock, } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const { mockDispatchSync, mockDbChain, mockResolveBillingAttribution } = vi.hoisted(() => { - const chain = { - select: vi.fn().mockReturnThis(), - from: vi.fn().mockReturnThis(), - where: vi.fn().mockReturnThis(), - orderBy: vi.fn().mockResolvedValue([]), - limit: vi.fn().mockResolvedValue([]), - update: vi.fn().mockReturnThis(), - set: vi.fn().mockReturnThis(), - } - return { - mockDispatchSync: vi.fn().mockResolvedValue(undefined), - mockDbChain: chain, - mockResolveBillingAttribution: vi.fn(), - } -}) +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockDispatchSync, mockResolveBillingAttribution } = vi.hoisted(() => ({ + mockDispatchSync: vi.fn().mockResolvedValue(undefined), + mockResolveBillingAttribution: vi.fn(), +})) const mockCheckWriteAccess = knowledgeApiUtilsMockFns.mockCheckKnowledgeBaseWriteAccess -vi.mock('@sim/db', () => ({ db: mockDbChain })) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/app/api/knowledge/utils', () => knowledgeApiUtilsMock) vi.mock('@/lib/billing/core/billing-attribution', () => ({ requireBillingAttributionHeader: vi.fn(), @@ -48,14 +39,12 @@ describe('Connector Manual Sync API Route', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() requestUtilsMockFns.mockGenerateRequestId.mockReturnValue('test-req-id') - mockDbChain.select.mockReturnThis() - mockDbChain.from.mockReturnThis() - mockDbChain.where.mockReturnThis() - mockDbChain.orderBy.mockResolvedValue([]) - mockDbChain.limit.mockResolvedValue([]) - mockDbChain.update.mockReturnThis() - mockDbChain.set.mockReturnThis() + }) + + afterAll(() => { + resetDbChainMock() }) it('returns 401 when unauthenticated', async () => { @@ -76,7 +65,7 @@ describe('Connector Manual Sync API Route', () => { userId: 'user-1', }) mockCheckWriteAccess.mockResolvedValue({ hasAccess: true }) - mockDbChain.limit.mockResolvedValueOnce([]) + dbChainMockFns.limit.mockResolvedValueOnce([]) const req = createMockRequest('POST') const response = await POST(req as never, { params: mockParams }) @@ -90,7 +79,7 @@ describe('Connector Manual Sync API Route', () => { userId: 'user-1', }) mockCheckWriteAccess.mockResolvedValue({ hasAccess: true }) - mockDbChain.limit.mockResolvedValueOnce([{ id: 'conn-456', status: 'syncing' }]) + dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'conn-456', status: 'syncing' }]) const req = createMockRequest('POST') const response = await POST(req as never, { params: mockParams }) @@ -122,7 +111,7 @@ describe('Connector Manual Sync API Route', () => { hasAccess: true, knowledgeBase: { workspaceId: 'ws-1', name: 'Test KB' }, }) - mockDbChain.limit.mockResolvedValueOnce([{ id: 'conn-456', status: 'active' }]) + dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'conn-456', status: 'active' }]) mockResolveBillingAttribution.mockResolvedValue(billingAttribution) const req = createMockRequest('POST') @@ -166,7 +155,7 @@ describe('Connector Manual Sync API Route', () => { hasAccess: true, knowledgeBase: { workspaceId: 'ws-1', name: 'Test KB' }, }) - mockDbChain.limit.mockResolvedValueOnce([{ id: 'conn-456', status: 'active' }]) + dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'conn-456', status: 'active' }]) mockResolveBillingAttribution.mockResolvedValue(billingAttribution) const req = createMockRequest( diff --git a/apps/sim/app/api/knowledge/[id]/connectors/route.test.ts b/apps/sim/app/api/knowledge/[id]/connectors/route.test.ts index 2ac24c18df0..d1f9f852ae4 100644 --- a/apps/sim/app/api/knowledge/[id]/connectors/route.test.ts +++ b/apps/sim/app/api/knowledge/[id]/connectors/route.test.ts @@ -5,45 +5,34 @@ import { auditMock, authOAuthUtilsMock, createMockRequest, + dbChainMock, + dbChainMockFns, hybridAuthMockFns, knowledgeApiUtilsMock, knowledgeApiUtilsMockFns, + resetDbChainMock, } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { mockCaptureServerEvent, - mockDbChain, mockDispatchSync, mockEncryptApiKey, mockHasWorkspaceLiveSyncAccess, mockResolveBillingAttribution, mockValidateConfig, -} = vi.hoisted(() => { - const chain = { - select: vi.fn().mockReturnThis(), - from: vi.fn().mockReturnThis(), - where: vi.fn().mockReturnThis(), - limit: vi.fn(), - execute: vi.fn(), - transaction: vi.fn(), - insert: vi.fn().mockReturnThis(), - values: vi.fn(), - } - return { - mockCaptureServerEvent: vi.fn(), - mockDbChain: chain, - mockDispatchSync: vi.fn(), - mockEncryptApiKey: vi.fn(), - mockHasWorkspaceLiveSyncAccess: vi.fn(), - mockResolveBillingAttribution: vi.fn(), - mockValidateConfig: vi.fn(), - } -}) +} = vi.hoisted(() => ({ + mockCaptureServerEvent: vi.fn(), + mockDispatchSync: vi.fn(), + mockEncryptApiKey: vi.fn(), + mockHasWorkspaceLiveSyncAccess: vi.fn(), + mockResolveBillingAttribution: vi.fn(), + mockValidateConfig: vi.fn(), +})) const mockCheckWriteAccess = knowledgeApiUtilsMockFns.mockCheckKnowledgeBaseWriteAccess -vi.mock('@sim/db', () => ({ db: mockDbChain })) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@sim/audit', () => auditMock) vi.mock('@/app/api/knowledge/utils', () => knowledgeApiUtilsMock) vi.mock('@/app/api/auth/oauth/utils', () => authOAuthUtilsMock) @@ -103,20 +92,16 @@ describe('Knowledge Connectors API Route', () => { beforeEach(() => { vi.clearAllMocks() - mockDbChain.select.mockReturnThis() - mockDbChain.from.mockReturnThis() - mockDbChain.where.mockReturnThis() - mockDbChain.insert.mockReturnThis() - mockDbChain.execute.mockResolvedValue(undefined) - mockDbChain.values.mockResolvedValue(undefined) - mockDbChain.transaction.mockImplementation( - async (callback: (tx: typeof mockDbChain) => unknown) => callback(mockDbChain) - ) + resetDbChainMock() mockDispatchSync.mockResolvedValue(undefined) mockEncryptApiKey.mockResolvedValue({ encrypted: 'encrypted-api-key' }) mockValidateConfig.mockResolvedValue({ valid: true }) }) + afterAll(() => { + resetDbChainMock() + }) + it('queues the authenticated actor with the paid workspace payer', async () => { hybridAuthMockFns.mockCheckSessionOrInternalAuth.mockResolvedValue({ success: true, @@ -135,7 +120,7 @@ describe('Knowledge Connectors API Route', () => { }) mockHasWorkspaceLiveSyncAccess.mockResolvedValue(true) mockResolveBillingAttribution.mockResolvedValue(BILLING_ATTRIBUTION) - mockDbChain.limit.mockResolvedValueOnce([{ id: 'knowledge-base-1' }]).mockResolvedValueOnce([ + dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'knowledge-base-1' }]).mockResolvedValueOnce([ { id: 'connector-1', knowledgeBaseId: 'knowledge-base-1', diff --git a/apps/sim/app/api/mcp/servers/[id]/refresh/route.test.ts b/apps/sim/app/api/mcp/servers/[id]/refresh/route.test.ts index e3f217fbdde..a04dd7edf74 100644 --- a/apps/sim/app/api/mcp/servers/[id]/refresh/route.test.ts +++ b/apps/sim/app/api/mcp/servers/[id]/refresh/route.test.ts @@ -1,22 +1,16 @@ /** * @vitest-environment node */ +import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' import type { NextRequest } from 'next/server' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' -const { mockClearCache, mockDiscoverServerTools, mockSelect, mockUpdateSet } = vi.hoisted(() => ({ +const { mockClearCache, mockDiscoverServerTools } = vi.hoisted(() => ({ mockClearCache: vi.fn(), mockDiscoverServerTools: vi.fn(), - mockSelect: vi.fn(), - mockUpdateSet: vi.fn(), })) -vi.mock('@sim/db', () => ({ - db: { - select: mockSelect, - update: vi.fn().mockReturnValue({ set: mockUpdateSet }), - }, -})) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/core/utils/with-route-handler', () => ({ withRouteHandler: (handler: unknown) => handler, @@ -68,23 +62,16 @@ const persistedServer = { toolCount: 0, } -function selectRows(rows: unknown[]) { - return { - from: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - limit: vi.fn().mockResolvedValue(rows), - }), - }), - } -} - describe('MCP server refresh route', () => { beforeEach(() => { vi.clearAllMocks() - mockSelect.mockReturnValueOnce(selectRows([initialServer])) - mockUpdateSet.mockReturnValue({ - where: vi.fn().mockReturnValue({ returning: vi.fn().mockResolvedValue([persistedServer]) }), - }) + resetDbChainMock() + dbChainMockFns.limit.mockResolvedValueOnce([initialServer]) + dbChainMockFns.returning.mockResolvedValue([persistedServer]) + }) + + afterAll(() => { + resetDbChainMock() }) it('preserves the service-persisted OAuth pending status', async () => { @@ -102,7 +89,7 @@ describe('MCP server refresh route', () => { error: null, }) ) - expect(mockUpdateSet).not.toHaveBeenCalledWith( + expect(dbChainMockFns.set).not.toHaveBeenCalledWith( expect.objectContaining({ connectionStatus: expect.anything() }) ) }) @@ -112,11 +99,7 @@ describe('MCP server refresh route', () => { mockDiscoverServerTools.mockRejectedValueOnce( new Error(`Upstream reflected ${reflectedSecret}`) ) - mockUpdateSet.mockReturnValueOnce({ - where: vi.fn().mockReturnValue({ - returning: vi.fn().mockResolvedValue([initialServer]), - }), - }) + dbChainMockFns.returning.mockResolvedValueOnce([initialServer]) const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { method: 'POST', @@ -142,11 +125,7 @@ describe('MCP server refresh route', () => { lastConnected: new Date(Date.now() + 60_000), toolCount: 7, } - mockUpdateSet.mockReturnValueOnce({ - where: vi.fn().mockReturnValue({ - returning: vi.fn().mockResolvedValue([newerSuccessfulServer]), - }), - }) + dbChainMockFns.returning.mockResolvedValueOnce([newerSuccessfulServer]) const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { method: 'POST', @@ -175,12 +154,17 @@ describe('MCP server refresh route', () => { serverName: 'OAuth Server', }, ]) - // The route's server lookup consumes the first select (beforeEach). The sync's - // workflow select is left unmocked, so it throws — exercising the guard that - // keeps a secondary sync failure from turning a successful refresh into a 500. - mockUpdateSet.mockReturnValueOnce({ - where: vi.fn().mockReturnValue({ returning: vi.fn().mockResolvedValue([initialServer]) }), - }) + // The route's server lookup consumes the first from() (verbatim mini-builder); + // the sync's workflow select then throws — exercising the guard that keeps a + // secondary sync failure from turning a successful refresh into a 500. + dbChainMockFns.from + .mockImplementationOnce(() => ({ + where: () => ({ limit: () => Promise.resolve([initialServer]) }), + })) + .mockImplementationOnce(() => { + throw new Error('workflow select unavailable') + }) + dbChainMockFns.returning.mockResolvedValueOnce([initialServer]) const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', { method: 'POST', diff --git a/apps/sim/app/api/mcp/servers/route.test.ts b/apps/sim/app/api/mcp/servers/route.test.ts index e831c802e0a..2d397432507 100644 --- a/apps/sim/app/api/mcp/servers/route.test.ts +++ b/apps/sim/app/api/mcp/servers/route.test.ts @@ -1,18 +1,15 @@ /** * @vitest-environment node */ +import { dbChainMock, resetDbChainMock } from '@sim/testing' import type { NextRequest } from 'next/server' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { mockPerformDeleteMcpServer } = vi.hoisted(() => ({ mockPerformDeleteMcpServer: vi.fn(), })) -vi.mock('@sim/db', () => ({ - db: { - select: vi.fn(), - }, -})) +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/mcp/middleware', () => ({ getParsedBody: () => undefined, @@ -57,6 +54,11 @@ function createDeleteRequest(serverId = 'server-1') { describe('MCP servers DELETE route', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() + }) + + afterAll(() => { + resetDbChainMock() }) it('returns 404 when orchestration reports a missing server', async () => { diff --git a/apps/sim/app/api/mothership/chats/[chatId]/fork/route.test.ts b/apps/sim/app/api/mothership/chats/[chatId]/fork/route.test.ts index a3c42f0f07e..e7ab8202a6e 100644 --- a/apps/sim/app/api/mothership/chats/[chatId]/fork/route.test.ts +++ b/apps/sim/app/api/mothership/chats/[chatId]/fork/route.test.ts @@ -1,13 +1,11 @@ /** * @vitest-environment node */ -import { copilotHttpMock, copilotHttpMockFns } from '@sim/testing' +import { copilotHttpMock, copilotHttpMockFns, dbChainMockFns, resetDbChainMock } from '@sim/testing' import { NextRequest } from 'next/server' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { - mockTransaction, - mockSelectRows, mockFilterForkableChatFiles, mockListForkableChatFiles, mockPlanChatFileCopies, @@ -18,11 +16,8 @@ const { mockFetchGo, mockPublishStatusChanged, mockCaptureServerEvent, - mockDeleteWhere, mockRemoveChatResources, } = vi.hoisted(() => ({ - mockTransaction: vi.fn(), - mockSelectRows: vi.fn(), // Real (pure) cut semantics so tests drive selection through row.messageId: // rows with a NULL/undefined messageId are kept in every fork. mockFilterForkableChatFiles: vi.fn( @@ -38,52 +33,9 @@ const { mockFetchGo: vi.fn(), mockPublishStatusChanged: vi.fn(), mockCaptureServerEvent: vi.fn(), - mockDeleteWhere: vi.fn(), mockRemoveChatResources: vi.fn(), })) -vi.mock('@sim/db', () => ({ - db: { - select: () => ({ - from: () => ({ - where: () => ({ - limit: () => mockSelectRows(), - }), - }), - }), - delete: () => ({ - where: mockDeleteWhere, - }), - transaction: mockTransaction, - }, -})) - -vi.mock('@sim/db/schema', () => ({ - copilotChats: { - id: 'copilotChats.id', - userId: 'copilotChats.userId', - type: 'copilotChats.type', - workspaceId: 'copilotChats.workspaceId', - title: 'copilotChats.title', - model: 'copilotChats.model', - resources: 'copilotChats.resources', - previewYaml: 'copilotChats.previewYaml', - planArtifact: 'copilotChats.planArtifact', - config: 'copilotChats.config', - deletedAt: 'copilotChats.deletedAt', - }, - workspaceFiles: { - id: 'workspaceFiles.id', - }, -})) - -vi.mock('drizzle-orm', () => ({ - and: vi.fn((...conditions: unknown[]) => ({ type: 'and', conditions })), - eq: vi.fn((field: unknown, value: unknown) => ({ type: 'eq', field, value })), - inArray: vi.fn((field: unknown, values: unknown) => ({ type: 'inArray', field, values })), - isNull: vi.fn((field: unknown) => ({ type: 'isNull', field })), -})) - vi.mock('@/lib/copilot/resources/persistence', () => ({ removeChatResources: mockRemoveChatResources, })) @@ -168,30 +120,6 @@ const threeMessages = [ }, ] -/** Chat rows inserted through the mock transaction, captured for title assertions. */ -let insertedChatRows: Array> = [] -/** tx.update(...).set(...) payloads, captured for resource-rewrite assertions. */ -let updatedChatRows: Array> = [] - -function makeTx() { - return { - insert: () => ({ - values: (v: Record) => { - insertedChatRows.push(v) - return { - returning: async () => [{ id: 'row-id', workspaceId: 'ws-1' }], - } - }, - }), - update: () => ({ - set: (v: Record) => { - updatedChatRows.push(v) - return { where: async () => undefined } - }, - }), - } -} - function createRequest(chatId: string, body?: unknown) { return new NextRequest(`http://localhost:3000/api/mothership/chats/${chatId}/fork`, { method: 'POST', @@ -207,13 +135,13 @@ function makeContext(chatId: string) { describe('POST /api/mothership/chats/[chatId]/fork', () => { beforeEach(() => { vi.clearAllMocks() - insertedChatRows = [] - updatedChatRows = [] + resetDbChainMock() copilotHttpMockFns.mockAuthenticateCopilotRequestSessionOnly.mockResolvedValue({ userId: 'user-1', isAuthenticated: true, }) - mockSelectRows.mockResolvedValue([parentRow]) + dbChainMockFns.limit.mockResolvedValue([parentRow]) + dbChainMockFns.returning.mockResolvedValue([{ id: 'row-id', workspaceId: 'ws-1' }]) mockListForkableChatFiles.mockResolvedValue([]) mockLoadCopilotChatMessages.mockResolvedValue(threeMessages) mockPlanChatFileCopies.mockResolvedValue({ @@ -223,13 +151,13 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => { }) mockExecuteChatFileBlobCopies.mockResolvedValue({ copied: 0, failed: 0, failedCopyIds: [] }) mockAppendCopilotChatMessages.mockResolvedValue(undefined) - mockDeleteWhere.mockResolvedValue(undefined) mockRemoveChatResources.mockResolvedValue(undefined) mockAssertActiveWorkspaceAccess.mockResolvedValue(undefined) mockFetchGo.mockResolvedValue({ ok: true }) - mockTransaction.mockImplementation(async (cb: (tx: unknown) => Promise) => - cb(makeTx()) - ) + }) + + afterAll(() => { + resetDbChainMock() }) it('rejects unauthenticated callers', async () => { @@ -242,14 +170,14 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => { }) it('404s when the chat belongs to another user', async () => { - mockSelectRows.mockResolvedValue([{ ...parentRow, userId: 'someone-else' }]) + dbChainMockFns.limit.mockResolvedValue([{ ...parentRow, userId: 'someone-else' }]) const res = await POST(createRequest('chat-1'), makeContext('chat-1')) expect(res.status).toBe(404) - expect(mockTransaction).not.toHaveBeenCalled() + expect(dbChainMockFns.transaction).not.toHaveBeenCalled() }) it('404s for non-mothership chats', async () => { - mockSelectRows.mockResolvedValue([{ ...parentRow, type: 'copilot' }]) + dbChainMockFns.limit.mockResolvedValue([{ ...parentRow, type: 'copilot' }]) const res = await POST(createRequest('chat-1'), makeContext('chat-1')) expect(res.status).toBe(404) }) @@ -257,13 +185,13 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => { it('400s when upToMessageId is missing', async () => { const res = await POST(createRequest('chat-1', {}), makeContext('chat-1')) expect(res.status).toBe(400) - expect(mockTransaction).not.toHaveBeenCalled() + expect(dbChainMockFns.transaction).not.toHaveBeenCalled() }) it('400s when upToMessageId is an empty string', async () => { const res = await POST(createRequest('chat-1', { upToMessageId: '' }), makeContext('chat-1')) expect(res.status).toBe(400) - expect(mockTransaction).not.toHaveBeenCalled() + expect(dbChainMockFns.transaction).not.toHaveBeenCalled() }) it('400s when the message is not in the chat', async () => { @@ -272,7 +200,7 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => { makeContext('chat-1') ) expect(res.status).toBe(400) - expect(mockTransaction).not.toHaveBeenCalled() + expect(dbChainMockFns.transaction).not.toHaveBeenCalled() }) it('applies the timeline cut: kept message ids drive the file selection', async () => { @@ -370,7 +298,7 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => { ) // Forks are titled "Fork | ". - expect(insertedChatRows[0].title).toBe('Fork | Generate Logs') + expect(dbChainMockFns.values.mock.calls[0][0].title).toBe('Fork | Generate Logs') }) it('still succeeds when the copilot-service clone fails (best-effort)', async () => { @@ -393,9 +321,9 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => { // The dead rows (committed, but no bytes behind them) are hard-deleted so // they vanish from VFS listings and name resolution… - expect(mockDeleteWhere).toHaveBeenCalledWith({ + expect(dbChainMockFns.where).toHaveBeenCalledWith({ type: 'inArray', - field: 'workspaceFiles.id', + column: 'id', values: ['wf_dead1', 'wf_dead2'], }) // …and their resource chips are dropped from the new chat. @@ -411,7 +339,7 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => { const cleanRes = await POST(createRequest('chat-1'), makeContext('chat-1')) expect('failedFileCopies' in (await cleanRes.json())).toBe(false) - expect(mockDeleteWhere).not.toHaveBeenCalled() + expect(dbChainMockFns.delete).not.toHaveBeenCalled() expect(mockRemoveChatResources).not.toHaveBeenCalled() }) @@ -421,7 +349,7 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => { // copies the kept upload AND the pre-cut apple; only the post-cut banana // stays behind, so only its resource is dropped — not left pointing at // the source chat. - mockSelectRows.mockResolvedValue([ + dbChainMockFns.limit.mockResolvedValue([ { ...parentRow, resources: [ @@ -457,8 +385,8 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => { OLD_FILE_ID, 'wf_apple', ]) - expect(updatedChatRows).toHaveLength(1) - expect(updatedChatRows[0].resources).toEqual([ + expect(dbChainMockFns.set).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.set.mock.calls[0][0].resources).toEqual([ { type: 'file', id: NEW_FILE_ID, title: 'cat.png' }, { type: 'file', id: 'wf_apple_copy', title: 'apple.png' }, { type: 'file', id: 'wf_shared', title: 'shared.pdf' }, @@ -469,7 +397,7 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => { it('drops ghosts even when the fork copies no files at all', async () => { // Fork cut before the chat's only upload arrived: a guard that skips the // resources update when idMap is empty would leave the ghost in place. - mockSelectRows.mockResolvedValue([ + dbChainMockFns.limit.mockResolvedValue([ { ...parentRow, resources: [{ type: 'file', id: 'wf_banana', title: 'banana.png' }], @@ -482,7 +410,7 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => { const res = await POST(createRequest('chat-1'), makeContext('chat-1')) expect(res.status).toBe(200) - expect(updatedChatRows).toHaveLength(1) - expect(updatedChatRows[0].resources).toEqual([]) + expect(dbChainMockFns.set).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.set.mock.calls[0][0].resources).toEqual([]) }) }) diff --git a/apps/sim/app/api/mothership/chats/[chatId]/restore/route.test.ts b/apps/sim/app/api/mothership/chats/[chatId]/restore/route.test.ts index 31d49d8a653..ee9438038b8 100644 --- a/apps/sim/app/api/mothership/chats/[chatId]/restore/route.test.ts +++ b/apps/sim/app/api/mothership/chats/[chatId]/restore/route.test.ts @@ -1,59 +1,15 @@ /** * @vitest-environment node */ -import { copilotHttpMock, copilotHttpMockFns } from '@sim/testing' +import { copilotHttpMock, copilotHttpMockFns, dbChainMockFns, resetDbChainMock } from '@sim/testing' import { NextRequest } from 'next/server' -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const { - mockDbSelect, - mockSelectFrom, - mockSelectWhere, - mockSelectLimit, - mockDbUpdate, - mockDbSet, - mockUpdateWhere, - mockDbReturning, - mockAssertActiveWorkspaceAccess, - mockPublishStatusChanged, -} = vi.hoisted(() => ({ - mockDbSelect: vi.fn(), - mockSelectFrom: vi.fn(), - mockSelectWhere: vi.fn(), - mockSelectLimit: vi.fn(), - mockDbUpdate: vi.fn(), - mockDbSet: vi.fn(), - mockUpdateWhere: vi.fn(), - mockDbReturning: vi.fn(), +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockAssertActiveWorkspaceAccess, mockPublishStatusChanged } = vi.hoisted(() => ({ mockAssertActiveWorkspaceAccess: vi.fn(), mockPublishStatusChanged: vi.fn(), })) -vi.mock('@sim/db', () => ({ - db: { - select: mockDbSelect, - update: mockDbUpdate, - }, -})) - -vi.mock('@sim/db/schema', () => ({ - copilotChats: { - id: 'copilotChats.id', - userId: 'copilotChats.userId', - type: 'copilotChats.type', - workspaceId: 'copilotChats.workspaceId', - updatedAt: 'copilotChats.updatedAt', - lastSeenAt: 'copilotChats.lastSeenAt', - deletedAt: 'copilotChats.deletedAt', - }, -})) - -vi.mock('drizzle-orm', () => ({ - and: vi.fn((...conditions: unknown[]) => ({ type: 'and', conditions })), - eq: vi.fn((field: unknown, value: unknown) => ({ type: 'eq', field, value })), - isNotNull: vi.fn((field: unknown) => ({ type: 'isNotNull', field })), -})) - vi.mock('@/lib/copilot/request/http', () => ({ ...copilotHttpMock, createForbiddenResponse: vi.fn((message: string) => ({ @@ -92,21 +48,20 @@ function makeContext(chatId: string) { describe('POST /api/mothership/chats/[chatId]/restore', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() copilotHttpMockFns.mockAuthenticateCopilotRequestSessionOnly.mockResolvedValue({ userId: 'user-1', isAuthenticated: true, }) - mockSelectLimit.mockResolvedValue([{ workspaceId: 'workspace-1' }]) - mockSelectWhere.mockReturnValue({ limit: mockSelectLimit }) - mockSelectFrom.mockReturnValue({ where: mockSelectWhere }) - mockDbSelect.mockReturnValue({ from: mockSelectFrom }) - mockDbReturning.mockResolvedValue([{ workspaceId: 'workspace-1' }]) - mockUpdateWhere.mockReturnValue({ returning: mockDbReturning }) - mockDbSet.mockReturnValue({ where: mockUpdateWhere }) - mockDbUpdate.mockReturnValue({ set: mockDbSet }) + dbChainMockFns.limit.mockResolvedValue([{ workspaceId: 'workspace-1' }]) + dbChainMockFns.returning.mockResolvedValue([{ workspaceId: 'workspace-1' }]) mockAssertActiveWorkspaceAccess.mockResolvedValue(undefined) }) + afterAll(() => { + resetDbChainMock() + }) + it('returns 401 when unauthenticated', async () => { copilotHttpMockFns.mockAuthenticateCopilotRequestSessionOnly.mockResolvedValueOnce({ userId: null, @@ -115,16 +70,16 @@ describe('POST /api/mothership/chats/[chatId]/restore', () => { const response = await POST(makeRequest('chat-1'), makeContext('chat-1')) expect(response.status).toBe(401) - expect(mockDbUpdate).not.toHaveBeenCalled() + expect(dbChainMockFns.update).not.toHaveBeenCalled() }) it('returns 404 when no soft-deleted chat is owned by the caller', async () => { - mockSelectLimit.mockResolvedValueOnce([]) + dbChainMockFns.limit.mockResolvedValueOnce([]) const response = await POST(makeRequest('chat-missing'), makeContext('chat-missing')) expect(response.status).toBe(404) expect(mockAssertActiveWorkspaceAccess).not.toHaveBeenCalled() - expect(mockDbUpdate).not.toHaveBeenCalled() + expect(dbChainMockFns.update).not.toHaveBeenCalled() }) it('returns 403 when the caller lost access to the workspace', async () => { @@ -132,7 +87,7 @@ describe('POST /api/mothership/chats/[chatId]/restore', () => { const response = await POST(makeRequest('chat-1'), makeContext('chat-1')) expect(response.status).toBe(403) - expect(mockDbUpdate).not.toHaveBeenCalled() + expect(dbChainMockFns.update).not.toHaveBeenCalled() }) it('restores the chat, bumping updatedAt and lastSeenAt, and publishes the event', async () => { @@ -141,7 +96,7 @@ describe('POST /api/mothership/chats/[chatId]/restore', () => { expect(response.status).toBe(200) expect(await response.json()).toEqual({ success: true }) expect(mockAssertActiveWorkspaceAccess).toHaveBeenCalledWith('workspace-1', 'user-1') - expect(mockDbSet).toHaveBeenCalledWith({ + expect(dbChainMockFns.set).toHaveBeenCalledWith({ deletedAt: null, updatedAt: expect.any(Date), lastSeenAt: expect.any(Date), @@ -154,7 +109,7 @@ describe('POST /api/mothership/chats/[chatId]/restore', () => { }) it('returns 404 when the chat is restored concurrently before the update lands', async () => { - mockDbReturning.mockResolvedValueOnce([]) + dbChainMockFns.returning.mockResolvedValueOnce([]) const response = await POST(makeRequest('chat-1'), makeContext('chat-1')) expect(response.status).toBe(404) diff --git a/apps/sim/app/api/mothership/chats/[chatId]/route.test.ts b/apps/sim/app/api/mothership/chats/[chatId]/route.test.ts index 095c5a11cc5..d0f2a64c00e 100644 --- a/apps/sim/app/api/mothership/chats/[chatId]/route.test.ts +++ b/apps/sim/app/api/mothership/chats/[chatId]/route.test.ts @@ -1,15 +1,11 @@ /** * @vitest-environment node */ -import { copilotHttpMock, copilotHttpMockFns } from '@sim/testing' +import { copilotHttpMock, copilotHttpMockFns, dbChainMockFns, resetDbChainMock } from '@sim/testing' import { NextRequest } from 'next/server' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { - mockDbUpdate, - mockDbSet, - mockDbReturning, - mockDbWhere, mockDecrementStorageUsageForBillingContext, mockDecrementStorageUsageForBillingContextInTx, mockGetAccessibleCopilotChat, @@ -18,10 +14,6 @@ const { mockReadFilePreviewSessions, mockGetLatestRunForStream, } = vi.hoisted(() => ({ - mockDbUpdate: vi.fn(), - mockDbSet: vi.fn(), - mockDbReturning: vi.fn(), - mockDbWhere: vi.fn(), mockDecrementStorageUsageForBillingContext: vi.fn(), mockDecrementStorageUsageForBillingContextInTx: vi.fn(), mockGetAccessibleCopilotChat: vi.fn(), @@ -31,38 +23,6 @@ const { mockGetLatestRunForStream: vi.fn(), })) -vi.mock('@sim/db', () => ({ - db: { - update: mockDbUpdate, - }, -})) - -vi.mock('@sim/db/schema', () => ({ - copilotChats: { - id: 'copilotChats.id', - userId: 'copilotChats.userId', - type: 'copilotChats.type', - updatedAt: 'copilotChats.updatedAt', - lastSeenAt: 'copilotChats.lastSeenAt', - workspaceId: 'copilotChats.workspaceId', - deletedAt: 'copilotChats.deletedAt', - }, -})) - -vi.mock('drizzle-orm', () => ({ - and: vi.fn((...conditions: unknown[]) => ({ type: 'and', conditions })), - eq: vi.fn((field: unknown, value: unknown) => ({ type: 'eq', field, value })), - isNull: vi.fn((field: unknown) => ({ type: 'isNull', field })), - sql: Object.assign( - vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ - type: 'sql', - strings, - values, - })), - { raw: vi.fn() } - ), -})) - vi.mock('@/lib/copilot/request/http', () => copilotHttpMock) vi.mock('@/lib/copilot/chat/lifecycle', () => ({ @@ -123,9 +83,14 @@ function createRequest(chatId: string) { }) } +afterAll(() => { + resetDbChainMock() +}) + describe('GET /api/mothership/chats/[chatId]', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() copilotHttpMockFns.mockAuthenticateCopilotRequestSessionOnly.mockResolvedValue({ userId: 'user-1', isAuthenticated: true, @@ -314,6 +279,7 @@ describe('GET /api/mothership/chats/[chatId]', () => { describe('DELETE /api/mothership/chats/[chatId]', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() copilotHttpMockFns.mockAuthenticateCopilotRequestSessionOnly.mockResolvedValue({ userId: 'user-1', isAuthenticated: true, @@ -323,10 +289,7 @@ describe('DELETE /api/mothership/chats/[chatId]', () => { type: 'mothership', workspaceId: 'workspace-1', }) - mockDbUpdate.mockReturnValue({ set: mockDbSet }) - mockDbSet.mockReturnValue({ where: mockDbWhere }) - mockDbWhere.mockReturnValue({ returning: mockDbReturning }) - mockDbReturning.mockResolvedValue([{ workspaceId: 'workspace-1' }]) + dbChainMockFns.returning.mockResolvedValue([{ workspaceId: 'workspace-1' }]) }) it('soft-deletes an unbilled chat without decrementing workspace or payer storage', async () => { @@ -338,8 +301,8 @@ describe('DELETE /api/mothership/chats/[chatId]', () => { ) expect(response.status).toBe(200) - expect(mockDbUpdate).toHaveBeenCalled() - expect(mockDbSet).toHaveBeenCalledWith({ deletedAt: expect.any(Date) }) + expect(dbChainMockFns.update).toHaveBeenCalled() + expect(dbChainMockFns.set).toHaveBeenCalledWith({ deletedAt: expect.any(Date) }) expect(mockDecrementStorageUsageForBillingContext).not.toHaveBeenCalled() expect(mockDecrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled() }) diff --git a/apps/sim/app/api/mothership/chats/read/route.test.ts b/apps/sim/app/api/mothership/chats/read/route.test.ts index ba2270510ab..f8e4a1b40ea 100644 --- a/apps/sim/app/api/mothership/chats/read/route.test.ts +++ b/apps/sim/app/api/mothership/chats/read/route.test.ts @@ -1,39 +1,14 @@ /** * @vitest-environment node */ -import { copilotHttpMock, copilotHttpMockFns } from '@sim/testing' +import { copilotHttpMock, copilotHttpMockFns, dbChainMockFns, resetDbChainMock } from '@sim/testing' import { NextRequest } from 'next/server' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' -const { mockUpdate, mockSet, mockWhere, mockParseRequest } = vi.hoisted(() => ({ - mockUpdate: vi.fn(), - mockSet: vi.fn(), - mockWhere: vi.fn(), +const { mockParseRequest } = vi.hoisted(() => ({ mockParseRequest: vi.fn(), })) -vi.mock('@sim/db', () => ({ - db: { update: mockUpdate }, -})) - -vi.mock('@sim/db/schema', () => ({ - copilotChats: { - id: 'copilotChats.id', - userId: 'copilotChats.userId', - updatedAt: 'copilotChats.updatedAt', - lastSeenAt: 'copilotChats.lastSeenAt', - }, -})) - -vi.mock('drizzle-orm', () => ({ - and: vi.fn((...conditions: unknown[]) => ({ type: 'and', conditions })), - eq: vi.fn((field: unknown, value: unknown) => ({ type: 'eq', field, value })), - or: vi.fn((...conditions: unknown[]) => ({ type: 'or', conditions })), - isNull: vi.fn((field: unknown) => ({ type: 'isNull', field })), - lt: vi.fn((field: unknown, value: unknown) => ({ type: 'lt', field, value })), - sql: vi.fn(() => ({ type: 'sql' })), -})) - vi.mock('@/lib/copilot/request/http', () => copilotHttpMock) vi.mock('@/lib/api/server', () => ({ parseRequest: mockParseRequest })) vi.mock('@/lib/api/contracts/mothership-chats', () => ({ markMothershipChatReadContract: {} })) @@ -50,22 +25,24 @@ function createRequest() { describe('POST /api/mothership/chats/read', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() copilotHttpMockFns.mockAuthenticateCopilotRequestSessionOnly.mockResolvedValue({ userId: 'user-1', isAuthenticated: true, }) mockParseRequest.mockResolvedValue({ success: true, data: { body: { chatId: 'chat-1' } } }) - mockWhere.mockResolvedValue(undefined) - mockSet.mockReturnValue({ where: mockWhere }) - mockUpdate.mockReturnValue({ set: mockSet }) + }) + + afterAll(() => { + resetDbChainMock() }) it('guards the lastSeenAt write with the unread predicate (only writes when unread)', async () => { const res = await POST(createRequest()) expect(res.status).toBe(200) - expect(mockUpdate).toHaveBeenCalledTimes(1) - const whereArg = mockWhere.mock.calls[0][0] as { + expect(dbChainMockFns.update).toHaveBeenCalledTimes(1) + const whereArg = dbChainMockFns.where.mock.calls[0][0] as { type: string conditions: Array<{ type: string; conditions?: unknown[] }> } @@ -75,8 +52,8 @@ describe('POST /api/mothership/chats/read', () => { expect(orClause).toBeDefined() expect(orClause?.conditions).toEqual( expect.arrayContaining([ - { type: 'isNull', field: 'copilotChats.lastSeenAt' }, - { type: 'lt', field: 'copilotChats.lastSeenAt', value: 'copilotChats.updatedAt' }, + { type: 'isNull', column: 'lastSeenAt' }, + { type: 'lt', left: 'lastSeenAt', right: 'updatedAt' }, ]) ) }) @@ -88,6 +65,6 @@ describe('POST /api/mothership/chats/read', () => { }) const res = await POST(createRequest()) expect(res.status).toBe(401) - expect(mockUpdate).not.toHaveBeenCalled() + expect(dbChainMockFns.update).not.toHaveBeenCalled() }) }) diff --git a/apps/sim/app/api/mothership/chats/route.test.ts b/apps/sim/app/api/mothership/chats/route.test.ts index 890e228a0b5..84296f0dd5f 100644 --- a/apps/sim/app/api/mothership/chats/route.test.ts +++ b/apps/sim/app/api/mothership/chats/route.test.ts @@ -1,47 +1,20 @@ /** * @vitest-environment node */ -import { copilotHttpMock, copilotHttpMockFns, permissionsMock } from '@sim/testing' +import { + copilotHttpMock, + copilotHttpMockFns, + dbChainMockFns, + permissionsMock, + queueTableRows, + resetDbChainMock, + schemaMock, +} from '@sim/testing' import { NextRequest } from 'next/server' -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const { mockSelect, mockFrom, mockWhere, mockOrderBy, mockReconcileChatStreamMarkers } = vi.hoisted( - () => ({ - mockSelect: vi.fn(), - mockFrom: vi.fn(), - mockWhere: vi.fn(), - mockOrderBy: vi.fn(), - mockReconcileChatStreamMarkers: vi.fn(), - }) -) - -vi.mock('@sim/db', () => ({ - db: { - select: mockSelect, - }, -})) - -vi.mock('@sim/db/schema', () => ({ - copilotChats: { - id: 'copilotChats.id', - title: 'copilotChats.title', - userId: 'copilotChats.userId', - workspaceId: 'copilotChats.workspaceId', - type: 'copilotChats.type', - updatedAt: 'copilotChats.updatedAt', - conversationId: 'copilotChats.conversationId', - lastSeenAt: 'copilotChats.lastSeenAt', - pinned: 'copilotChats.pinned', - deletedAt: 'copilotChats.deletedAt', - }, -})) +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' -vi.mock('drizzle-orm', () => ({ - and: vi.fn((...conditions: unknown[]) => ({ type: 'and', conditions })), - desc: vi.fn((field: unknown) => ({ type: 'desc', field })), - eq: vi.fn((field: unknown, value: unknown) => ({ type: 'eq', field, value })), - isNull: vi.fn((field: unknown) => ({ type: 'isNull', field })), - isNotNull: vi.fn((field: unknown) => ({ type: 'isNotNull', field })), +const { mockReconcileChatStreamMarkers } = vi.hoisted(() => ({ + mockReconcileChatStreamMarkers: vi.fn(), })) vi.mock('@/lib/copilot/request/http', () => copilotHttpMock) @@ -70,17 +43,13 @@ function createRequest(workspaceId: string) { describe('GET /api/mothership/chats', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() copilotHttpMockFns.mockAuthenticateCopilotRequestSessionOnly.mockResolvedValue({ userId: 'user-1', isAuthenticated: true, }) - mockOrderBy.mockResolvedValue([]) - mockWhere.mockReturnValue({ orderBy: mockOrderBy }) - mockFrom.mockReturnValue({ where: mockWhere }) - mockSelect.mockReturnValue({ from: mockFrom }) - mockReconcileChatStreamMarkers.mockImplementation( async (candidates: Array<{ chatId: string; streamId: string | null }>) => new Map( @@ -98,7 +67,7 @@ describe('GET /api/mothership/chats', () => { it('clears activeStreamId on chats whose redis lock has expired (stuck-yellow bug)', async () => { const now = new Date('2026-05-11T12:00:00Z') - mockOrderBy.mockResolvedValueOnce([ + queueTableRows(schemaMock.copilotChats, [ { id: 'chat-stuck', title: 'Stuck chat', @@ -151,7 +120,7 @@ describe('GET /api/mothership/chats', () => { it('preserves chats when no chat has a stream marker set', async () => { const now = new Date('2026-05-11T12:00:00Z') - mockOrderBy.mockResolvedValueOnce([ + queueTableRows(schemaMock.copilotChats, [ { id: 'chat-1', title: null, updatedAt: now, activeStreamId: null, lastSeenAt: null }, { id: 'chat-2', title: null, updatedAt: now, activeStreamId: null, lastSeenAt: null }, ]) @@ -175,7 +144,7 @@ describe('GET /api/mothership/chats', () => { it('leaves activeStreamId untouched when redis confirms every lock is live', async () => { const now = new Date('2026-05-11T12:00:00Z') - mockOrderBy.mockResolvedValueOnce([ + queueTableRows(schemaMock.copilotChats, [ { id: 'chat-a', title: null, updatedAt: now, activeStreamId: 'stream-a', lastSeenAt: null }, { id: 'chat-b', title: null, updatedAt: now, activeStreamId: 'stream-b', lastSeenAt: null }, ]) @@ -191,7 +160,7 @@ describe('GET /api/mothership/chats', () => { it('uses Redis lock owner when it differs from a stale activeStreamId', async () => { const now = new Date('2026-05-11T12:00:00Z') - mockOrderBy.mockResolvedValueOnce([ + queueTableRows(schemaMock.copilotChats, [ { id: 'chat-mismatch', title: null, @@ -223,7 +192,11 @@ describe('GET /api/mothership/chats', () => { const response = await GET(createRequest('ws-1')) expect(response.status).toBe(401) - expect(mockSelect).not.toHaveBeenCalled() + expect(dbChainMockFns.select).not.toHaveBeenCalled() expect(mockReconcileChatStreamMarkers).not.toHaveBeenCalled() }) + + afterAll(() => { + resetDbChainMock() + }) }) diff --git a/apps/sim/app/api/workspaces/[id]/api-keys/route.test.ts b/apps/sim/app/api/workspaces/[id]/api-keys/route.test.ts index 67b354a5268..d5d3b874c78 100644 --- a/apps/sim/app/api/workspaces/[id]/api-keys/route.test.ts +++ b/apps/sim/app/api/workspaces/[id]/api-keys/route.test.ts @@ -1,33 +1,19 @@ /** * @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 { mockGetApiKeyDisplayFormat, mockGetSession, mockGetUserEntityPermissions, mockGetWorkspaceById, - mockOrderBy, } = vi.hoisted(() => ({ mockGetApiKeyDisplayFormat: vi.fn(), mockGetSession: vi.fn(), mockGetUserEntityPermissions: vi.fn(), mockGetWorkspaceById: vi.fn(), - mockOrderBy: vi.fn(), -})) - -vi.mock('@sim/db', () => ({ - db: { - select: vi.fn(() => ({ - from: vi.fn(() => ({ - where: vi.fn(() => ({ - orderBy: mockOrderBy, - })), - })), - })), - }, })) vi.mock('@/lib/api-key/auth', () => ({ @@ -53,11 +39,12 @@ import { GET } from '@/app/api/workspaces/[id]/api-keys/route' describe('GET /api/workspaces/[id]/api-keys', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockGetSession.mockResolvedValue({ user: { id: 'reader-1' } }) mockGetWorkspaceById.mockResolvedValue({ id: 'workspace-1' }) mockGetUserEntityPermissions.mockResolvedValue('read') mockGetApiKeyDisplayFormat.mockResolvedValue('sim_••••legacy') - mockOrderBy.mockResolvedValue([ + queueTableRows(schemaMock.apiKey, [ { id: 'key-1', name: 'Legacy key', @@ -70,6 +57,10 @@ describe('GET /api/workspaces/[id]/api-keys', () => { ]) }) + afterAll(() => { + resetDbChainMock() + }) + it('returns metadata without exposing the stored key value', async () => { const response = await GET(createMockRequest('GET'), { params: Promise.resolve({ id: 'workspace-1' }), diff --git a/apps/sim/app/api/workspaces/[id]/byok-keys/route.test.ts b/apps/sim/app/api/workspaces/[id]/byok-keys/route.test.ts index 186839287d7..89b7f60b7b0 100644 --- a/apps/sim/app/api/workspaces/[id]/byok-keys/route.test.ts +++ b/apps/sim/app/api/workspaces/[id]/byok-keys/route.test.ts @@ -1,63 +1,29 @@ /** * @vitest-environment node */ -import { auditMock, createMockRequest } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + auditMock, + createMockRequest, + dbChainMockFns, + queueTableRows, + resetDbChainMock, + schemaMock, +} from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { - mockDbState, mockGetSession, mockGetUserEntityPermissions, mockGetWorkspaceById, mockEncryptSecret, mockDecryptSecret, - mockUpdateSet, - mockInsertValues, - mockDeleteWhere, -} = vi.hoisted(() => { - const state = { - selectResults: [] as unknown[][], - insertReturning: [] as unknown[], - deleteReturning: [] as unknown[], - } - return { - mockDbState: state, - mockGetSession: vi.fn(), - mockGetUserEntityPermissions: vi.fn(), - mockGetWorkspaceById: vi.fn(), - mockEncryptSecret: vi.fn(), - mockDecryptSecret: vi.fn(), - mockUpdateSet: vi.fn(() => ({ where: vi.fn(() => Promise.resolve()) })), - mockInsertValues: vi.fn(() => ({ - returning: vi.fn(() => Promise.resolve(state.insertReturning)), - })), - mockDeleteWhere: vi.fn(() => ({ - returning: vi.fn(() => Promise.resolve(state.deleteReturning)), - })), - } -}) - -vi.mock('@sim/db', () => { - const dbMock: Record = { - select: vi.fn(() => { - const chain: Record = {} - chain.from = vi.fn().mockReturnValue(chain) - chain.where = vi.fn().mockImplementation(() => { - const result: any = Promise.resolve(mockDbState.selectResults.shift() ?? []) - result.limit = vi.fn(() => result) - result.orderBy = vi.fn(() => result) - return result - }) - return chain - }), - update: vi.fn(() => ({ set: mockUpdateSet })), - insert: vi.fn(() => ({ values: mockInsertValues })), - delete: vi.fn(() => ({ where: mockDeleteWhere })), - execute: vi.fn(() => Promise.resolve([])), - } - dbMock.transaction = vi.fn(async (callback: (tx: unknown) => unknown) => callback(dbMock)) - return { db: dbMock } -}) +} = vi.hoisted(() => ({ + mockGetSession: vi.fn(), + mockGetUserEntityPermissions: vi.fn(), + mockGetWorkspaceById: vi.fn(), + mockEncryptSecret: vi.fn(), + mockDecryptSecret: vi.fn(), +})) vi.mock('@sim/audit', () => auditMock) @@ -97,9 +63,7 @@ const storedKeyRow = (id: string, name: string | null = null) => ({ describe('workspace BYOK keys route', () => { beforeEach(() => { vi.clearAllMocks() - mockDbState.selectResults = [] - mockDbState.insertReturning = [] - mockDbState.deleteReturning = [] + resetDbChainMock() mockGetSession.mockResolvedValue({ user: { id: 'user-1' } }) mockGetUserEntityPermissions.mockResolvedValue('admin') @@ -110,9 +74,16 @@ describe('workspace BYOK keys route', () => { })) }) + afterAll(() => { + resetDbChainMock() + }) + describe('GET', () => { it('lists every stored key with name and masked value', async () => { - mockDbState.selectResults = [[storedKeyRow('key-1', 'Production'), storedKeyRow('key-2')]] + queueTableRows(schemaMock.workspaceBYOKKeys, [ + storedKeyRow('key-1', 'Production'), + storedKeyRow('key-2'), + ]) const res = await GET(createMockRequest('GET'), routeContext) @@ -143,14 +114,14 @@ describe('workspace BYOK keys route', () => { ) expect(res.status).toBe(403) - expect(mockInsertValues).not.toHaveBeenCalled() + expect(dbChainMockFns.values).not.toHaveBeenCalled() }) it('adds a new key even when the provider already has keys', async () => { - mockDbState.selectResults = [[{ keyCount: 2 }]] - mockDbState.insertReturning = [ + queueTableRows(schemaMock.workspaceBYOKKeys, [{ keyCount: 2 }]) + dbChainMockFns.returning.mockResolvedValueOnce([ { id: 'key-3', providerId: 'openai', name: 'Backup', createdAt: new Date() }, - ] + ]) const res = await POST( createMockRequest('POST', { providerId: 'openai', apiKey: 'sk-new-key', name: 'Backup' }), @@ -161,7 +132,7 @@ describe('workspace BYOK keys route', () => { const body = await res.json() expect(body.success).toBe(true) expect(body.key).toMatchObject({ id: 'key-3', name: 'Backup' }) - expect(mockInsertValues).toHaveBeenCalledWith( + expect(dbChainMockFns.values).toHaveBeenCalledWith( expect.objectContaining({ workspaceId: WORKSPACE_ID, providerId: 'openai', @@ -172,10 +143,10 @@ describe('workspace BYOK keys route', () => { }) it('stores a null name when none is provided', async () => { - mockDbState.selectResults = [[{ keyCount: 0 }]] - mockDbState.insertReturning = [ + queueTableRows(schemaMock.workspaceBYOKKeys, [{ keyCount: 0 }]) + dbChainMockFns.returning.mockResolvedValueOnce([ { id: 'key-1', providerId: 'openai', name: null, createdAt: new Date() }, - ] + ]) const res = await POST( createMockRequest('POST', { providerId: 'openai', apiKey: 'sk-new-key' }), @@ -183,11 +154,11 @@ describe('workspace BYOK keys route', () => { ) expect(res.status).toBe(200) - expect(mockInsertValues).toHaveBeenCalledWith(expect.objectContaining({ name: null })) + expect(dbChainMockFns.values).toHaveBeenCalledWith(expect.objectContaining({ name: null })) }) it('rejects adding a key beyond the per-provider cap', async () => { - mockDbState.selectResults = [[{ keyCount: 10 }]] + queueTableRows(schemaMock.workspaceBYOKKeys, [{ keyCount: 10 }]) const res = await POST( createMockRequest('POST', { providerId: 'openai', apiKey: 'sk-new-key' }), @@ -197,12 +168,12 @@ describe('workspace BYOK keys route', () => { expect(res.status).toBe(400) const body = await res.json() expect(body.error).toContain('at most 10 keys') - expect(mockInsertValues).not.toHaveBeenCalled() + expect(dbChainMockFns.values).not.toHaveBeenCalled() expect(mockEncryptSecret).not.toHaveBeenCalled() }) it('updates the targeted key in place when keyId is provided', async () => { - mockDbState.selectResults = [[{ id: 'key-2', name: 'Old name' }]] + queueTableRows(schemaMock.workspaceBYOKKeys, [{ id: 'key-2', name: 'Old name' }]) const res = await POST( createMockRequest('POST', { providerId: 'openai', apiKey: 'sk-rotated', keyId: 'key-2' }), @@ -212,14 +183,14 @@ describe('workspace BYOK keys route', () => { expect(res.status).toBe(200) const body = await res.json() expect(body.key).toMatchObject({ id: 'key-2', name: 'Old name' }) - expect(mockUpdateSet).toHaveBeenCalledWith( + expect(dbChainMockFns.set).toHaveBeenCalledWith( expect.objectContaining({ encryptedApiKey: 'encrypted-value', name: 'Old name' }) ) - expect(mockInsertValues).not.toHaveBeenCalled() + expect(dbChainMockFns.values).not.toHaveBeenCalled() }) it('clears the name when updating with an empty name', async () => { - mockDbState.selectResults = [[{ id: 'key-2', name: 'Old name' }]] + queueTableRows(schemaMock.workspaceBYOKKeys, [{ id: 'key-2', name: 'Old name' }]) const res = await POST( createMockRequest('POST', { @@ -232,11 +203,11 @@ describe('workspace BYOK keys route', () => { ) expect(res.status).toBe(200) - expect(mockUpdateSet).toHaveBeenCalledWith(expect.objectContaining({ name: null })) + expect(dbChainMockFns.set).toHaveBeenCalledWith(expect.objectContaining({ name: null })) }) it('returns 404 when the keyId does not exist in the workspace', async () => { - mockDbState.selectResults = [[]] + queueTableRows(schemaMock.workspaceBYOKKeys, []) const res = await POST( createMockRequest('POST', { providerId: 'openai', apiKey: 'sk-rotated', keyId: 'missing' }), @@ -244,7 +215,7 @@ describe('workspace BYOK keys route', () => { ) expect(res.status).toBe(404) - expect(mockUpdateSet).not.toHaveBeenCalled() + expect(dbChainMockFns.set).not.toHaveBeenCalled() expect(mockEncryptSecret).not.toHaveBeenCalled() }) @@ -260,7 +231,7 @@ describe('workspace BYOK keys route', () => { describe('DELETE', () => { it('deletes a single key when keyId is provided', async () => { - mockDbState.deleteReturning = [{ id: 'key-2' }] + dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'key-2' }]) const res = await DELETE( createMockRequest('DELETE', { providerId: 'openai', keyId: 'key-2' }), @@ -272,7 +243,7 @@ describe('workspace BYOK keys route', () => { }) it('returns 404 when keyId is provided but no key matches', async () => { - mockDbState.deleteReturning = [] + dbChainMockFns.returning.mockResolvedValueOnce([]) const res = await DELETE( createMockRequest('DELETE', { providerId: 'openai', keyId: 'missing' }), @@ -283,7 +254,7 @@ describe('workspace BYOK keys route', () => { }) it('deletes all provider keys when keyId is omitted', async () => { - mockDbState.deleteReturning = [{ id: 'key-1' }, { id: 'key-2' }] + dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'key-1' }, { id: 'key-2' }]) const res = await DELETE(createMockRequest('DELETE', { providerId: 'openai' }), routeContext) @@ -292,7 +263,7 @@ describe('workspace BYOK keys route', () => { }) it('succeeds when keyId is omitted and the provider has no keys', async () => { - mockDbState.deleteReturning = [] + dbChainMockFns.returning.mockResolvedValueOnce([]) const res = await DELETE(createMockRequest('DELETE', { providerId: 'openai' }), routeContext) @@ -306,7 +277,7 @@ describe('workspace BYOK keys route', () => { const res = await DELETE(createMockRequest('DELETE', { providerId: 'openai' }), routeContext) expect(res.status).toBe(403) - expect(mockDeleteWhere).not.toHaveBeenCalled() + expect(dbChainMockFns.delete).not.toHaveBeenCalled() }) }) }) diff --git a/apps/sim/app/api/workspaces/[id]/fork/excluded-workflows/route.test.ts b/apps/sim/app/api/workspaces/[id]/fork/excluded-workflows/route.test.ts index f882cb9bf54..a623def8937 100644 --- a/apps/sim/app/api/workspaces/[id]/fork/excluded-workflows/route.test.ts +++ b/apps/sim/app/api/workspaces/[id]/fork/excluded-workflows/route.test.ts @@ -1,16 +1,22 @@ /** * @vitest-environment node */ -import { auditMock, auditMockFns, createMockRequest } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' - -const { mockGetSession, mockAssertWorkspaceAdminAccess, mockDbUpdate, mockCaptureServerEvent } = - vi.hoisted(() => ({ +import { + auditMock, + auditMockFns, + createMockRequest, + dbChainMockFns, + resetDbChainMock, +} from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockGetSession, mockAssertWorkspaceAdminAccess, mockCaptureServerEvent } = vi.hoisted( + () => ({ mockGetSession: vi.fn(), mockAssertWorkspaceAdminAccess: vi.fn(), - mockDbUpdate: vi.fn(), mockCaptureServerEvent: vi.fn(), - })) + }) +) vi.mock('@/lib/auth', () => ({ auth: { api: { getSession: vi.fn() } }, @@ -27,10 +33,6 @@ vi.mock('@/lib/posthog/server', () => ({ captureServerEvent: mockCaptureServerEvent, })) -vi.mock('@sim/db', () => ({ - db: { update: () => mockDbUpdate() }, -})) - import { PUT } from '@/app/api/workspaces/[id]/fork/excluded-workflows/route' const WORKSPACE_ID = 'workspace-1' @@ -38,23 +40,22 @@ const ADMIN_ID = 'user-1' const routeContext = { params: Promise.resolve({ id: WORKSPACE_ID }) } function mockUpdateReturning(rows: Array<{ id: string; name: string }>) { - mockDbUpdate.mockReturnValue({ - set: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ - returning: vi.fn().mockResolvedValue(rows), - }), - }), - }) + dbChainMockFns.returning.mockResolvedValue(rows) } describe('fork excluded-workflows route', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockGetSession.mockResolvedValue({ user: { id: ADMIN_ID } }) mockAssertWorkspaceAdminAccess.mockResolvedValue({ id: WORKSPACE_ID, name: 'My Workspace' }) mockUpdateReturning([]) }) + afterAll(() => { + resetDbChainMock() + }) + it('returns 401 when there is no session', async () => { mockGetSession.mockResolvedValue(null) @@ -74,7 +75,7 @@ describe('fork excluded-workflows route', () => { ) expect(res.status).toBe(400) - expect(mockDbUpdate).not.toHaveBeenCalled() + expect(dbChainMockFns.update).not.toHaveBeenCalled() }) it('requires workspace admin (and the fork entitlement gate) before writing', async () => { diff --git a/apps/sim/app/api/workspaces/invitations/route.test.ts b/apps/sim/app/api/workspaces/invitations/route.test.ts index ddfe8a59213..d4eefc75af4 100644 --- a/apps/sim/app/api/workspaces/invitations/route.test.ts +++ b/apps/sim/app/api/workspaces/invitations/route.test.ts @@ -9,9 +9,11 @@ import { permissionsMock, permissionsMockFns, posthogServerMock, + queueTableRows, + resetDbChainMock, schemaMock, } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { mockGetWorkspaceInvitePolicy, @@ -22,7 +24,6 @@ const { mockSendInvitationEmail, mockCancelPendingInvitation, mockFindPendingGrantForWorkspaceEmail, - mockDbResults, } = vi.hoisted(() => ({ mockGetWorkspaceInvitePolicy: vi.fn(), mockValidateInvitationsAllowed: vi.fn().mockResolvedValue(undefined), @@ -32,30 +33,8 @@ const { mockSendInvitationEmail: vi.fn(), mockCancelPendingInvitation: vi.fn(), mockFindPendingGrantForWorkspaceEmail: vi.fn(), - mockDbResults: { value: [] as any[] }, })) -vi.mock('@sim/db', () => ({ - db: { - select: vi.fn().mockImplementation(() => { - const chain: any = {} - chain.from = vi.fn().mockReturnValue(chain) - chain.innerJoin = vi.fn().mockReturnValue(chain) - chain.where = vi.fn().mockReturnValue(chain) - chain.limit = vi - .fn() - .mockImplementation(() => Promise.resolve(mockDbResults.value.shift() || [])) - chain.then = vi.fn().mockImplementation((callback: (rows: any[]) => unknown) => { - const result = mockDbResults.value.shift() || [] - return Promise.resolve(callback ? callback(result) : result) - }) - return chain - }), - }, -})) - -vi.mock('@sim/db/schema', () => schemaMock) - vi.mock('@/lib/auth', () => authMock) vi.mock('@/lib/workspaces/permissions/utils', () => permissionsMock) @@ -112,7 +91,7 @@ import { POST } from '@/app/api/workspaces/invitations/batch/route' describe('POST /api/workspaces/invitations/batch', () => { beforeEach(() => { vi.clearAllMocks() - mockDbResults.value = [] + resetDbChainMock() mockGetSession.mockResolvedValue({ user: { id: 'user-1', email: 'owner@test.com', name: 'Owner User' }, }) @@ -149,6 +128,10 @@ describe('POST /api/workspaces/invitations/batch', () => { mockFindPendingGrantForWorkspaceEmail.mockResolvedValue(null) }) + afterAll(() => { + resetDbChainMock() + }) + it('blocks invites for personal workspaces with an upgrade prompt', async () => { mockGetWorkspaceWithOwner.mockResolvedValueOnce({ id: 'workspace-1', @@ -165,7 +148,6 @@ describe('POST /api/workspaces/invitations/batch', () => { organizationId: null, upgradeRequired: true, }) - mockDbResults.value = [] const request = createMockRequest('POST', { workspaceId: 'workspace-1', @@ -196,7 +178,6 @@ describe('POST /api/workspaces/invitations/batch', () => { organizationId: null, upgradeRequired: true, }) - mockDbResults.value = [] const request = createMockRequest('POST', { workspaceId: 'workspace-1', @@ -234,7 +215,6 @@ describe('POST /api/workspaces/invitations/batch', () => { maxSeats: 5, availableSeats: 0, }) - mockDbResults.value = [[]] const request = createMockRequest('POST', { workspaceId: 'workspace-1', @@ -277,7 +257,7 @@ describe('POST /api/workspaces/invitations/batch', () => { role: 'member', memberId: 'member-1', }) - mockDbResults.value = [[{ id: 'existing-user', email: 'new@example.com' }]] + queueTableRows(schemaMock.user, [{ id: 'existing-user', email: 'new@example.com' }]) const request = createMockRequest('POST', { workspaceId: 'workspace-1', @@ -311,7 +291,6 @@ describe('POST /api/workspaces/invitations/batch', () => { workspaceMode: 'grandfathered_shared', billedAccountUserId: 'user-1', }) - mockDbResults.value = [[]] const request = createMockRequest('POST', { workspaceId: 'workspace-1', @@ -336,7 +315,6 @@ describe('POST /api/workspaces/invitations/batch', () => { }) it('creates multiple workspace invitations in one batch request', async () => { - mockDbResults.value = [[], []] mockCreatePendingInvitation .mockResolvedValueOnce({ invitationId: 'inv-1', @@ -382,7 +360,6 @@ describe('POST /api/workspaces/invitations/batch', () => { success: false, error: 'mailer unavailable', }) - mockDbResults.value = [[]] const request = createMockRequest('POST', { workspaceId: 'workspace-1', diff --git a/apps/sim/lib/copilot/tools/handlers/deployment/manage.test.ts b/apps/sim/lib/copilot/tools/handlers/deployment/manage.test.ts index 66ed0a0d8fc..eb2cc11467b 100644 --- a/apps/sim/lib/copilot/tools/handlers/deployment/manage.test.ts +++ b/apps/sim/lib/copilot/tools/handlers/deployment/manage.test.ts @@ -2,8 +2,15 @@ * @vitest-environment node */ -import { auditMock, workflowsOrchestrationMock, workflowsOrchestrationMockFns } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + auditMock, + queueTableRows, + resetDbChainMock, + schemaMock, + workflowsOrchestrationMock, + workflowsOrchestrationMockFns, +} from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' import type { ExecutionContext } from '@/lib/copilot/request/types' const { ensureWorkflowAccessMock, checkNeedsRedeploymentMock } = vi.hoisted(() => ({ @@ -23,20 +30,6 @@ const { resolveWorkflowStateRefMock, generateWorkflowDiffSummaryMock, listWorkfl listWorkflowVersionsMock: vi.fn(), })) -vi.mock('@sim/db', () => ({ - db: { - select: vi.fn(), - insert: vi.fn(), - update: vi.fn(), - delete: vi.fn(), - }, - chat: {}, - workflow: {}, - workflowDeploymentVersion: {}, - workflowMcpServer: {}, - workflowMcpTool: {}, -})) - vi.mock('@sim/audit', () => auditMock) vi.mock('@/lib/mcp/pubsub', () => ({ @@ -81,7 +74,6 @@ vi.mock('@/lib/workflows/persistence/utils', () => ({ updateDeploymentVersionMetadata: vi.fn(), })) -import { db } from '@sim/db' import { executeCheckDeploymentStatus, executeDiffWorkflows, @@ -90,16 +82,9 @@ import { executePromoteToLive, } from './manage' -function selectChain(result: unknown[], resolveOnWhere = false) { - const chain = { - from: vi.fn(() => chain), - innerJoin: vi.fn(() => chain), - where: vi.fn(() => (resolveOnWhere ? Promise.resolve(result) : chain)), - orderBy: vi.fn(() => Promise.resolve(result)), - limit: vi.fn(() => Promise.resolve(result)), - } - return chain -} +afterAll(() => { + resetDbChainMock() +}) describe('executeLoadDeployment', () => { beforeEach(() => { @@ -332,6 +317,7 @@ describe('executeDiffWorkflows', () => { describe('executeCheckDeploymentStatus', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() ensureWorkflowAccessMock.mockResolvedValue({ workflow: { id: 'wf-1', workspaceId: 'ws-1', name: 'Test Workflow' }, }) @@ -353,10 +339,7 @@ describe('executeCheckDeploymentStatus', () => { latestDeploymentAttempt: null, warnings: [], }) - vi.mocked(db.select) - .mockReturnValueOnce(selectChain([{ deployedAt: new Date('2026-05-28') }]) as never) - .mockReturnValueOnce(selectChain([]) as never) - .mockReturnValueOnce(selectChain([], true) as never) + queueTableRows(schemaMock.workflow, [{ deployedAt: new Date('2026-05-28') }]) checkNeedsRedeploymentMock.mockResolvedValueOnce(true) const result = await executeCheckDeploymentStatus({ workflowId: 'wf-1' }, { @@ -376,10 +359,7 @@ describe('executeCheckDeploymentStatus', () => { }) it('does not check redeployment freshness for undeployed APIs', async () => { - vi.mocked(db.select) - .mockReturnValueOnce(selectChain([{ deployedAt: null }]) as never) - .mockReturnValueOnce(selectChain([]) as never) - .mockReturnValueOnce(selectChain([], true) as never) + queueTableRows(schemaMock.workflow, [{ deployedAt: null }]) const result = await executeCheckDeploymentStatus({ workflowId: 'wf-1' }, { userId: 'user-1', diff --git a/apps/sim/lib/copilot/tools/handlers/resources.test.ts b/apps/sim/lib/copilot/tools/handlers/resources.test.ts index af6abf456fc..c27fbd737cd 100644 --- a/apps/sim/lib/copilot/tools/handlers/resources.test.ts +++ b/apps/sim/lib/copilot/tools/handlers/resources.test.ts @@ -9,12 +9,6 @@ const { getWorkspaceFileMock, resolveWorkspaceFileReferenceMock } = vi.hoisted(( resolveWorkspaceFileReferenceMock: vi.fn(), })) -vi.mock('@sim/db', () => ({ - db: {}, -})) - -vi.mock('@sim/db/schema', () => ({})) - vi.mock('@/lib/uploads/contexts/workspace/workspace-file-manager', () => ({ getWorkspaceFile: getWorkspaceFileMock, resolveWorkspaceFileReference: resolveWorkspaceFileReferenceMock, diff --git a/apps/sim/lib/copilot/tools/handlers/vfs-mutate.test.ts b/apps/sim/lib/copilot/tools/handlers/vfs-mutate.test.ts index 30847fdcc07..1b7860da1df 100644 --- a/apps/sim/lib/copilot/tools/handlers/vfs-mutate.test.ts +++ b/apps/sim/lib/copilot/tools/handlers/vfs-mutate.test.ts @@ -1,7 +1,8 @@ /** * @vitest-environment node */ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { dbChainMock, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const mocks = vi.hoisted(() => ({ ensureWorkspaceAccess: vi.fn(), @@ -25,19 +26,9 @@ const mocks = vi.hoisted(() => ({ getKnowledgeBases: vi.fn(), updateKnowledgeBase: vi.fn(), checkKnowledgeBaseWriteAccess: vi.fn(), - workflowRows: vi.fn(), })) -vi.mock('@sim/db', () => ({ - db: { - select: () => ({ - from: () => ({ - where: () => Promise.resolve(mocks.workflowRows()), - }), - }), - }, - workflow: { id: 'id', name: 'name', folderId: 'folderId', workspaceId: 'workspaceId' }, -})) +vi.mock('@sim/db', () => ({ ...dbChainMock, ...schemaMock })) vi.mock('@sim/platform-authz/workflow', () => ({ assertFolderMutable: mocks.assertFolderMutable, @@ -102,18 +93,22 @@ const context = { userId: 'user-1', workspaceId: 'ws-1' } as ExecutionContext describe('vfs mv/cp', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mocks.ensureWorkspaceAccess.mockResolvedValue(undefined) mocks.ensureWorkflowAccess.mockResolvedValue({ workspaceId: 'ws-1', workflow: {} }) mocks.assertFolderMutable.mockResolvedValue(undefined) mocks.assertWorkflowMutable.mockResolvedValue(undefined) mocks.verifyFolderWorkspace.mockResolvedValue(true) mocks.listFolders.mockResolvedValue([]) - mocks.workflowRows.mockReturnValue([]) mocks.getWorkspaceFileByName.mockResolvedValue(null) mocks.findWorkspaceFileFolderIdByPath.mockResolvedValue(null) mocks.ensureWorkspaceFileFolderPath.mockResolvedValue('ensured-folder') }) + afterAll(() => { + resetDbChainMock() + }) + describe('category rules', () => { it('rejects cross-category moves', async () => { const result = await executeVfsMv( @@ -286,7 +281,7 @@ describe('vfs mv/cp', () => { describe('workflows', () => { it('renames a workflow at root', async () => { - mocks.workflowRows.mockReturnValue([{ id: 'wf-1', name: 'Old Name', folderId: null }]) + queueTableRows(schemaMock.workflow, [{ id: 'wf-1', name: 'Old Name', folderId: null }]) mocks.performUpdateWorkflow.mockResolvedValue({ success: true }) const result = await executeVfsMv( @@ -306,7 +301,7 @@ describe('vfs mv/cp', () => { mocks.listFolders.mockResolvedValue([ { folderId: 'fold-1', folderName: 'Archive', parentId: null }, ]) - mocks.workflowRows.mockReturnValue([{ id: 'wf-1', name: 'My Workflow', folderId: null }]) + queueTableRows(schemaMock.workflow, [{ id: 'wf-1', name: 'My Workflow', folderId: null }]) mocks.performUpdateWorkflow.mockResolvedValue({ success: true }) const result = await executeVfsMv( @@ -322,7 +317,7 @@ describe('vfs mv/cp', () => { }) it('surfaces locked-workflow rejections per item', async () => { - mocks.workflowRows.mockReturnValue([{ id: 'wf-1', name: 'Locked One', folderId: null }]) + queueTableRows(schemaMock.workflow, [{ id: 'wf-1', name: 'Locked One', folderId: null }]) mocks.assertWorkflowMutable.mockRejectedValue(new Error('Workflow is locked')) const result = await executeVfsMv( @@ -335,7 +330,7 @@ describe('vfs mv/cp', () => { }) it('duplicates a workflow with cp (locked source allowed)', async () => { - mocks.workflowRows.mockReturnValue([{ id: 'wf-1', name: 'Template', folderId: null }]) + queueTableRows(schemaMock.workflow, [{ id: 'wf-1', name: 'Template', folderId: null }]) mocks.duplicateWorkflow.mockResolvedValue({ id: 'wf-2', name: 'My Copy' }) const result = await executeVfsCp( diff --git a/apps/sim/lib/copilot/tools/handlers/workflow/mutations.test.ts b/apps/sim/lib/copilot/tools/handlers/workflow/mutations.test.ts index 2ec9cb72264..41e3c6c922e 100644 --- a/apps/sim/lib/copilot/tools/handlers/workflow/mutations.test.ts +++ b/apps/sim/lib/copilot/tools/handlers/workflow/mutations.test.ts @@ -1,7 +1,7 @@ /** * @vitest-environment node */ -import { createEnvMock, workflowAuthzMockFns } from '@sim/testing' +import { createEnvMock, dbChainMock, schemaMock, workflowAuthzMockFns } from '@sim/testing' import { beforeEach, describe, expect, it, vi } from 'vitest' import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution' import type { ExecutionContext } from '@/lib/copilot/request/types' @@ -44,10 +44,7 @@ vi.mock('@sim/audit', () => ({ recordAudit: recordAuditMock, })) -vi.mock('@sim/db', () => ({ - db: {}, - workflow: {}, -})) +vi.mock('@sim/db', () => ({ ...dbChainMock, ...schemaMock })) vi.mock('@/lib/api-key/orchestration', () => ({ performCreateWorkspaceApiKey: vi.fn(), diff --git a/apps/sim/lib/mcp/service-pool.test.ts b/apps/sim/lib/mcp/service-pool.test.ts index be7fb6de847..406471bcced 100644 --- a/apps/sim/lib/mcp/service-pool.test.ts +++ b/apps/sim/lib/mcp/service-pool.test.ts @@ -9,8 +9,8 @@ */ import { UnauthorizedError } from '@modelcontextprotocol/sdk/client/auth.js' import { StreamableHTTPError } from '@modelcontextprotocol/sdk/client/streamableHttp.js' -import { loggerMock } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { dbChainMock, dbChainMockFns, loggerMock, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { MockMcpClient, @@ -71,32 +71,24 @@ vi.mock('@/lib/mcp/client', () => ({ McpClient: MockMcpClient })) const WORKSPACE_ID = 'ws-1' const USER_ID = 'user-1' -vi.mock('@sim/db', () => { - const where = () => { - const rows = Promise.resolve([ - { - id: 'server-1', - name: 'Server 1', - description: null, - transport: 'streamable-http', - url: 'https://server-1.example.com/mcp', - authType: 'headers', - workspaceId: WORKSPACE_ID, - headers: {}, - timeout: 30000, - retries: 3, - enabled: true, - deletedAt: null, - createdAt: new Date('2026-01-01T00:00:00Z'), - updatedAt: new Date('2026-01-01T00:00:00Z'), - }, - ]) - return Object.assign(rows, { limit: (n: number) => rows.then((r) => r.slice(0, n)) }) - } - return { - db: { select: vi.fn().mockReturnValue({ from: vi.fn().mockReturnValue({ where }) }) }, - } -}) +const SERVER_ROW = { + id: 'server-1', + name: 'Server 1', + description: null, + transport: 'streamable-http', + url: 'https://server-1.example.com/mcp', + authType: 'headers', + workspaceId: WORKSPACE_ID, + headers: {}, + timeout: 30000, + retries: 3, + enabled: true, + deletedAt: null, + createdAt: new Date('2026-01-01T00:00:00Z'), + updatedAt: new Date('2026-01-01T00:00:00Z'), +} + +vi.mock('@sim/db', () => dbChainMock) vi.mock('@/lib/mcp/domain-check', () => ({ isMcpDomainAllowed: () => true, validateMcpDomain: () => {}, @@ -121,11 +113,19 @@ import { mcpService } from '@/lib/mcp/service' describe('McpService connection reuse wiring', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() + // Every select here is getServerConfig's `.where(...).limit(1)`; the + // persistent override keeps the row available across retries. + dbChainMockFns.limit.mockResolvedValue([SERVER_ROW]) mockResolveEnvVars.mockImplementation(async (config: unknown) => ({ config })) mockAcquire.mockResolvedValue({ client: poolClient, release: mockRelease }) mockCallTool.mockResolvedValue({ content: [] }) }) + afterAll(() => { + resetDbChainMock() + }) + it('leases from the pool (keyed by server+workspace+user) and never disconnects on a hit', async () => { await mcpService.executeTool(USER_ID, 'server-1', { name: 'do', arguments: {} }, WORKSPACE_ID) diff --git a/apps/sim/lib/mcp/service.test.ts b/apps/sim/lib/mcp/service.test.ts index 1fc6e1e4b8f..04f33b2b4ff 100644 --- a/apps/sim/lib/mcp/service.test.ts +++ b/apps/sim/lib/mcp/service.test.ts @@ -3,8 +3,8 @@ */ import { UnauthorizedError } from '@modelcontextprotocol/sdk/client/auth.js' -import { loggerMock } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { dbChainMock, dbChainMockFns, loggerMock, resetDbChainMock } from '@sim/testing' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { MockMcpClient, @@ -17,13 +17,10 @@ const { mockValidateSsrf, mockIsDomainAllowed, mockCacheAdapter, - mockUpdateSet, - mockUpdateReturning, } = vi.hoisted(() => { const mockListTools = vi.fn() const mockConnect = vi.fn() const mockDisconnect = vi.fn() - const mockUpdateReturning = vi.fn().mockResolvedValue([{ id: 'server-1' }]) // In-memory cache adapter so the service never touches the real Redis the // local .env points at (unreachable in CI/sandbox → hangs). Honors TTL via // an expiry timestamp so negative-cache assertions behave like production. @@ -73,34 +70,27 @@ const { mockValidateDomain: vi.fn(), mockValidateSsrf: vi.fn(), mockIsDomainAllowed: vi.fn(() => true), - mockUpdateReturning, - mockUpdateSet: vi.fn().mockReturnValue({ - where: vi.fn().mockReturnValue({ returning: mockUpdateReturning }), - }), } }) -vi.mock('@sim/db', () => { - // `where(...)` resolves to the workspace's rows AND exposes `.limit()` for - // chains like `getServerConfig` that do `select().from().where().limit(1)`. - const where = (...args: unknown[]) => { - const rowsPromise = Promise.resolve(mockGetWorkspaceServersRows(...args)) - const thenable = Object.assign(rowsPromise, { - limit: (n: number) => rowsPromise.then((rows) => rows.slice(0, n)), - }) - return thenable - } - return { - db: { - select: vi.fn().mockReturnValue({ - from: vi.fn().mockReturnValue({ where }), - }), - update: vi.fn().mockReturnValue({ set: mockUpdateSet }), - insert: vi.fn(), - delete: vi.fn(), - }, - } -}) +vi.mock('@sim/db', () => dbChainMock) + +/** + * Routes every select chain to `mockGetWorkspaceServersRows`: `where(...)` + * resolves the workspace's rows AND exposes `.limit()` for chains like + * `getServerConfig` that do `select().from().where().limit(1)`. + */ +function wireSelectsToWorkspaceRows() { + dbChainMockFns.from.mockImplementation(() => { + const rows = Promise.resolve(mockGetWorkspaceServersRows()) + return { + where: () => + Object.assign(rows, { + limit: (n: number) => rows.then((r: unknown[]) => r.slice(0, n)), + }), + } + }) +} vi.mock('@/lib/mcp/client', () => ({ McpClient: MockMcpClient, @@ -173,6 +163,9 @@ function tool(name: string, serverId: string) { describe('McpService.discoverTools per-server caching', () => { beforeEach(async () => { vi.clearAllMocks() + resetDbChainMock() + wireSelectsToWorkspaceRows() + dbChainMockFns.returning.mockResolvedValue([{ id: 'server-1' }]) // `clearAllMocks` does not drain `.mockResolvedValueOnce` queues; reset // listTools so a previous test's unconsumed mock doesn't leak into the next. mockListTools.mockReset() @@ -184,12 +177,14 @@ describe('McpService.discoverTools per-server caching', () => { ) mockConnect.mockResolvedValue(undefined) mockDisconnect.mockResolvedValue(undefined) - mockUpdateReturning.mockReset() - mockUpdateReturning.mockResolvedValue([{ id: 'server-1' }]) // The McpService singleton holds cache state across imports. await mcpService.clearCache() }) + afterAll(() => { + resetDbChainMock() + }) + it('caches each server independently after first discovery', async () => { mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A'), dbRow('mcp-b', 'B')]) mockListTools @@ -349,7 +344,7 @@ describe('McpService.discoverTools per-server caching', () => { expect(first).toEqual([]) await vi.waitFor(() => { - expect(mockUpdateSet).toHaveBeenCalledWith( + expect(dbChainMockFns.set).toHaveBeenCalledWith( expect.objectContaining({ connectionStatus: 'disconnected', lastError: 'Authentication failed', @@ -362,7 +357,7 @@ describe('McpService.discoverTools per-server caching', () => { expect.any(Number) ) }) - expect(JSON.stringify(mockUpdateSet.mock.calls)).not.toContain(reflectedCredential) + expect(JSON.stringify(dbChainMockFns.set.mock.calls)).not.toContain(reflectedCredential) expect(JSON.stringify(mockCacheAdapter.set.mock.calls)).not.toContain(reflectedCredential) expect(JSON.stringify(mockLogger?.warn.mock.calls)).not.toContain(reflectedCredential) @@ -380,7 +375,7 @@ describe('McpService.discoverTools per-server caching', () => { expect(first).toEqual([]) await vi.waitFor(() => { - expect(mockUpdateSet).toHaveBeenCalledWith( + expect(dbChainMockFns.set).toHaveBeenCalledWith( expect.objectContaining({ connectionStatus: 'disconnected', lastError: null, @@ -462,7 +457,7 @@ describe('McpService.discoverTools per-server caching', () => { 'Request timed out' ) - expect(mockUpdateSet).toHaveBeenCalledWith( + expect(dbChainMockFns.set).toHaveBeenCalledWith( expect.objectContaining({ connectionStatus: 'disconnected', // Raw SDK timeout text is mapped to a user-facing message before persisting. @@ -499,14 +494,14 @@ describe('McpService.discoverTools per-server caching', () => { reflectedCredential ) - expect(mockUpdateSet).toHaveBeenCalledWith( + expect(dbChainMockFns.set).toHaveBeenCalledWith( expect.objectContaining({ connectionStatus: 'disconnected', lastError: 'Authentication failed', statusConfig: { consecutiveFailures: 1, lastSuccessfulDiscovery: null }, }) ) - expect(JSON.stringify(mockUpdateSet.mock.calls)).not.toContain(reflectedCredential) + expect(JSON.stringify(dbChainMockFns.set.mock.calls)).not.toContain(reflectedCredential) expect(JSON.stringify(mockCacheAdapter.set.mock.calls)).not.toContain(reflectedCredential) expect(JSON.stringify(mockLogger?.warn.mock.calls)).not.toContain(reflectedCredential) @@ -538,7 +533,7 @@ describe('McpService.discoverTools per-server caching', () => { 'OAuth token rejected' ) - expect(mockUpdateSet).toHaveBeenCalledWith( + expect(dbChainMockFns.set).toHaveBeenCalledWith( expect.objectContaining({ connectionStatus: 'disconnected', lastError: null, @@ -569,7 +564,7 @@ describe('McpService.discoverTools per-server caching', () => { 'Connection refused' ) - expect(mockUpdateSet).toHaveBeenCalledWith( + expect(dbChainMockFns.set).toHaveBeenCalledWith( expect.objectContaining({ connectionStatus: 'error', statusConfig: { consecutiveFailures: 3, lastSuccessfulDiscovery: null }, @@ -585,7 +580,7 @@ describe('McpService.discoverTools per-server caching', () => { 'OAuth authorization required' ) - expect(mockUpdateSet).toHaveBeenCalledWith( + expect(dbChainMockFns.set).toHaveBeenCalledWith( expect.objectContaining({ connectionStatus: 'disconnected', lastError: null, @@ -596,7 +591,7 @@ describe('McpService.discoverTools per-server caching', () => { it('does not negative-cache a failure older than a successful discovery', async () => { mockGetWorkspaceServersRows.mockResolvedValue([dbRow('mcp-a', 'A')]) mockListTools.mockRejectedValueOnce(new Error('Older request failed')) - mockUpdateReturning.mockResolvedValueOnce([]) + dbChainMockFns.returning.mockResolvedValueOnce([]) await expect(mcpService.discoverServerTools(USER_ID, 'mcp-a', WORKSPACE_ID)).rejects.toThrow( 'Older request failed' diff --git a/packages/testing/src/mocks/schema.mock.ts b/packages/testing/src/mocks/schema.mock.ts index 8bd29692bd7..ca251cba0ee 100644 --- a/packages/testing/src/mocks/schema.mock.ts +++ b/packages/testing/src/mocks/schema.mock.ts @@ -725,6 +725,8 @@ export const schemaMock = { config: 'config', resources: 'resources', lastSeenAt: 'lastSeenAt', + pinned: 'pinned', + deletedAt: 'deletedAt', createdAt: 'createdAt', updatedAt: 'updatedAt', }, From 3f189f539fd62938ddc93aa5740b5ef0cbec551e Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 22 Jul 2026 15:13:21 -0700 Subject: [PATCH 2/2] fix(testing): drain unconsumed ...Once overrides in resetDbChainMock MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../testing/src/mocks/database.mock.test.ts | 6 ++ packages/testing/src/mocks/database.mock.ts | 55 +++++-------------- 2 files changed, 21 insertions(+), 40 deletions(-) diff --git a/packages/testing/src/mocks/database.mock.test.ts b/packages/testing/src/mocks/database.mock.test.ts index 6e5f58fa178..94145bd3f0d 100644 --- a/packages/testing/src/mocks/database.mock.test.ts +++ b/packages/testing/src/mocks/database.mock.test.ts @@ -146,6 +146,12 @@ describe('database mock', () => { await expect(db.select().from(workflowTable).where({})).resolves.toEqual([]) }) + it('drains unconsumed ...Once overrides on resetDbChainMock', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'leftover' }]) + resetDbChainMock() + await expect(db.select().from(workflowTable).where({}).limit(1)).resolves.toEqual([]) + }) + it('clears queues and rewires defaults on resetDbChainMock', async () => { queueTableRows(workflowTable, [{ id: 'stale' }]) dbChainMockFns.where.mockReturnValue('broken' as never) diff --git a/packages/testing/src/mocks/database.mock.ts b/packages/testing/src/mocks/database.mock.ts index 897d2c5f0d1..f9699300af2 100644 --- a/packages/testing/src/mocks/database.mock.ts +++ b/packages/testing/src/mocks/database.mock.ts @@ -54,7 +54,9 @@ export function createMockSqlOperators() { exists: vi.fn((subquery) => ({ type: 'exists', subquery })), notExists: vi.fn((subquery) => ({ type: 'notExists', subquery })), like: vi.fn((column, pattern) => ({ type: 'like', column, pattern })), + notLike: vi.fn((column, pattern) => ({ type: 'notLike', column, pattern })), ilike: vi.fn((column, pattern) => ({ type: 'ilike', column, pattern })), + notIlike: vi.fn((column, pattern) => ({ type: 'notIlike', column, pattern })), desc: vi.fn((column) => ({ type: 'desc', column })), asc: vi.fn((column) => ({ type: 'asc', column })), } @@ -299,50 +301,23 @@ export const dbChainMockFns = { } /** - * Restores every `dbChainMockFns` entry to its default wiring and clears all - * table-routed row queues. Call this in `beforeEach` (after - * `vi.clearAllMocks()`) if any test uses `mockReturnValue` / - * `mockResolvedValue` (permanent overrides) or `queueTableRows` — this - * guarantees the next test starts with fresh defaults. - * - * Not needed if tests exclusively use the `...Once` variants, since those - * auto-expire after one call. + * Restores every `dbChainMockFns` entry to its default wiring, drains any + * unconsumed `...Once` overrides, and clears all table-routed row queues. + * Call this in `beforeEach` (after `vi.clearAllMocks()`) so each test starts + * from fresh defaults — a `...Once` override queued by a previous test but + * never consumed would otherwise leak into the next test (`vi.clearAllMocks` + * clears call history only, not once-queues). */ export function resetDbChainMock(): void { tableRowQueues.clear() - for (const spy of [ - select, - selectDistinct, - selectDistinctOn, - from, - where, - limit, - offset, - orderBy, - groupBy, - having, - forClause, - innerJoin, - leftJoin, - insert, - update, - set, - del, - ]) { - spy.mockImplementation(() => CHAIN_DEFAULT) + // mockReset restores the implementation passed to vi.fn() (the sentinel for + // structural spies, the real defaults for value terminals) AND drains any + // unconsumed ...Once overrides — covering the shared spies and the stable + // db-instance wrappers alike. + for (const spy of Object.values(dbChainMockFns)) { + ;(spy as ChainSpy).mockReset() } - returning.mockImplementation(() => Promise.resolve([] as unknown[])) - execute.mockImplementation(() => Promise.resolve([] as unknown[])) - query.mockImplementation(() => Promise.resolve([] as unknown[])) - onConflictDoUpdate.mockImplementation(() => ({ returning }) as unknown as Promise) - onConflictDoNothing.mockImplementation(() => ({ returning }) as unknown as Promise) - values.mockImplementation(() => ({ returning, onConflictDoUpdate, onConflictDoNothing })) - transaction.mockImplementation(async (cb: (tx: typeof dbChainMock.db) => unknown) => - cb(dbChainMock.db) - ) - // The stable db-instance entry points are wrappers around the spies above; a - // suite may have overridden them directly, so restore their original - // implementations too (mockReset restores the fn passed to vi.fn()). + query.mockReset() for (const key of [ 'select', 'selectDistinct',