diff --git a/apps/sim/lib/admin/dashboard-organizations.test.ts b/apps/sim/lib/admin/dashboard-organizations.test.ts index 00c87bc32cc..3bf092f9585 100644 --- a/apps/sim/lib/admin/dashboard-organizations.test.ts +++ b/apps/sim/lib/admin/dashboard-organizations.test.ts @@ -1,15 +1,12 @@ /** @vitest-environment node */ -import { databaseMock, dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' -import type { Mock } from 'vitest' +import { member, organization, permissions, subscription } from '@sim/db/schema' +import { dbChainMock, dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing' import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' vi.unmock('drizzle-orm') const mocks = vi.hoisted(() => ({ - queryRows: [] as unknown[][], - selectCalls: 0, - selectDistinctOnCalls: 0, provisionings: new Map(), })) @@ -65,78 +62,8 @@ vi.mock('@/lib/core/outbox/service', () => ({ enqueueOutboxEvent: vi.fn() })) import { listDashboardOrganizations, toDashboardConfigurationUpdate } from '@/lib/admin/dashboard' -/** - * `@sim/db` behavior is driven through the SHARED `dbChainMockFns` instances - * instead of a file-local factory object. This file mocks `@sim/db` with - * `dbChainMock` and installs the queued-rows select implementation in - * `beforeEach`; the setup-level `databaseMock` entry points are mirrored onto - * the same chain fns. Under `isolate: false` the module under test may have - * been loaded by an earlier suite in this worker with `@sim/db` bound to - * `databaseMock` — configuring both shared instances keeps either binding - * correct. - */ -function queryChain(rows: unknown[]) { - const chain: Record = {} - for (const method of ['from', 'innerJoin', 'leftJoin', 'where', 'orderBy', 'limit']) { - chain[method] = () => chain - } - chain.offset = () => Promise.resolve(rows) - chain.groupBy = () => Promise.resolve(rows) - chain.then = (resolve: (value: unknown[]) => unknown, reject: (error: unknown) => unknown) => - Promise.resolve(rows).then(resolve, reject) - return chain -} - -function queuedSelect() { - const rows = mocks.queryRows[mocks.selectCalls] ?? [] - mocks.selectCalls += 1 - return queryChain(rows) -} - -function queuedSelectDistinctOn() { - const rows = mocks.queryRows[mocks.selectCalls] ?? [] - mocks.selectCalls += 1 - mocks.selectDistinctOnCalls += 1 - return queryChain(rows) -} - -const GLOBAL_DB_KEYS = [ - 'select', - 'selectDistinct', - 'selectDistinctOn', - 'insert', - 'update', - 'delete', - 'transaction', -] as const - -const globalDb = databaseMock.db as unknown as Record<(typeof GLOBAL_DB_KEYS)[number], Mock> -const savedGlobalDbImpls = new Map< - (typeof GLOBAL_DB_KEYS)[number], - ((...args: unknown[]) => unknown) | undefined ->() - -/** Mirrors the setup-level databaseMock entry points onto the shared chain fns. */ -function delegateGlobalDbToChainMocks(): void { - for (const key of GLOBAL_DB_KEYS) { - const fn = globalDb[key] - if (typeof fn?.mockImplementation !== 'function') continue - if (!savedGlobalDbImpls.has(key)) savedGlobalDbImpls.set(key, fn.getMockImplementation()) - fn.mockImplementation((...args: unknown[]) => (dbChainMockFns[key] as Mock)(...args)) - } -} - -/** Restores the databaseMock entry points captured before this suite ran. */ -function restoreGlobalDb(): void { - for (const [key, impl] of savedGlobalDbImpls) { - if (impl) globalDb[key].mockImplementation(impl) - else globalDb[key].mockReset() - } -} - afterAll(() => { resetDbChainMock() - restoreGlobalDb() }) describe('toDashboardConfigurationUpdate', () => { @@ -173,48 +100,41 @@ describe('listDashboardOrganizations', () => { beforeEach(() => { vi.clearAllMocks() resetDbChainMock() - dbChainMockFns.select.mockImplementation(queuedSelect) - dbChainMockFns.selectDistinctOn.mockImplementation(queuedSelectDistinctOn) - delegateGlobalDbToChainMocks() - mocks.selectCalls = 0 - mocks.selectDistinctOnCalls = 0 mocks.provisionings = new Map() }) it('loads a page with a fixed batch of queries instead of querying once per organization', async () => { - mocks.queryRows = [ - [{ total: 2 }], - [ - { id: 'org-1', name: 'One', orgUsageLimit: '10', creditBalance: '1' }, - { id: 'org-2', name: 'Two', orgUsageLimit: '20', creditBalance: '2' }, - ], - [ - { - organizationId: 'org-1', - memberCount: 2, - ownerId: 'owner-1', - ownerName: 'Owner One', - ownerEmail: 'one@example.com', - }, - { - organizationId: 'org-2', - memberCount: 1, - ownerId: 'owner-2', - ownerName: 'Owner Two', - ownerEmail: 'two@example.com', - }, - ], - [{ organizationId: 'org-1', externalCollaboratorCount: 3 }], - [ - { - id: 'sub-1', - referenceId: 'org-1', - plan: 'team_6000', - status: 'active', - metadata: null, - }, - ], - ] + queueTableRows(organization, [{ total: 2 }]) + queueTableRows(organization, [ + { id: 'org-1', name: 'One', orgUsageLimit: '10', creditBalance: '1' }, + { id: 'org-2', name: 'Two', orgUsageLimit: '20', creditBalance: '2' }, + ]) + queueTableRows(member, [ + { + organizationId: 'org-1', + memberCount: 2, + ownerId: 'owner-1', + ownerName: 'Owner One', + ownerEmail: 'one@example.com', + }, + { + organizationId: 'org-2', + memberCount: 1, + ownerId: 'owner-2', + ownerName: 'Owner Two', + ownerEmail: 'two@example.com', + }, + ]) + queueTableRows(permissions, [{ organizationId: 'org-1', externalCollaboratorCount: 3 }]) + queueTableRows(subscription, [ + { + id: 'sub-1', + referenceId: 'org-1', + plan: 'team_6000', + status: 'active', + metadata: null, + }, + ]) const result = await listDashboardOrganizations({ search: '', limit: 50, offset: 0 }) @@ -231,7 +151,7 @@ describe('listDashboardOrganizations', () => { externalCollaboratorCount: 0, planLabel: 'No plan', }) - expect(mocks.selectCalls).toBe(5) - expect(mocks.selectDistinctOnCalls).toBe(1) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(4) + expect(dbChainMockFns.selectDistinctOn).toHaveBeenCalledTimes(1) }) }) diff --git a/apps/sim/lib/billing/core/plan.test.ts b/apps/sim/lib/billing/core/plan.test.ts index a9e855f4861..9e776770eb9 100644 --- a/apps/sim/lib/billing/core/plan.test.ts +++ b/apps/sim/lib/billing/core/plan.test.ts @@ -2,8 +2,7 @@ * @vitest-environment node */ import { member, organization, subscription } from '@sim/db/schema' -import { databaseMock, dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' -import type { Mock } from 'vitest' +import { dbChainMock, dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing' import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' vi.mock('@sim/db', () => dbChainMock) @@ -25,88 +24,20 @@ vi.mock('@/lib/billing/subscriptions/utils', () => ({ import { getHighestPrioritySubscription } from '@/lib/billing/core/plan' /** - * Drizzle mock for `getHighestPrioritySubscription`. It issues up to four - * queries keyed by table: + * `getHighestPrioritySubscription` issues up to four queries keyed by table: * - `subscription` for the user's personal subs (parallelized with members) * - `member` for the user's org memberships (parallelized with subs) * - `organization` for the org-existence follow-up * - `subscription` again for the org-scoped subs follow-up * - * The mock routes results by the table object passed to `.from()`, serving the - * (twice-read) `subscription` table from a FIFO queue (first read = personal, - * second = org). It records which tables were queried so we can assert the - * parallelized pair both run and that follow-ups are skipped when appropriate. - * - * The routing implementation is installed in `beforeEach` onto the SHARED - * `dbChainMockFns.select` (this file mocks `@sim/db` with `dbChainMock`) AND - * mirrored onto the setup-level `databaseMock` entry points. Under - * `isolate: false` the module under test may have been loaded by an earlier - * suite in this worker with `@sim/db` bound to `databaseMock` — pointing both - * shared instances at the same routing keeps either binding correct. Table - * identity likewise relies on the setup-level `@sim/db/schema` mock (no local - * schema factory), which is stable across suites in the shared worker. + * Results are routed by the table object passed to `.from()` via + * `queueTableRows` (FIFO per table: first `subscription` read = personal, + * second = org). `dbChainMockFns.from` call args record which tables were + * queried so we can assert the parallelized pair both run and that follow-ups + * are skipped when appropriate. */ -type TableName = 'subscription' | 'member' | 'organization' - -const TABLE_NAMES = new Map([ - [subscription, 'subscription'], - [member, 'member'], - [organization, 'organization'], -]) - -const resultsByTable: Record = { - subscription: [], - member: [], - organization: [], -} -const fromCalls: string[] = [] - -function routedSelect() { - return { - from: (table: unknown) => { - const name = TABLE_NAMES.get(table) - if (name) fromCalls.push(name) - const where = () => { - const queue = name ? resultsByTable[name] : undefined - const next = queue && queue.length > 0 ? queue.shift() : [] - return Promise.resolve(next ?? []) - } - return { where } - }, - } -} - -const GLOBAL_DB_KEYS = [ - 'select', - 'selectDistinct', - 'insert', - 'update', - 'delete', - 'transaction', -] as const - -const globalDb = databaseMock.db as unknown as Record<(typeof GLOBAL_DB_KEYS)[number], Mock> -const savedGlobalDbImpls = new Map< - (typeof GLOBAL_DB_KEYS)[number], - ((...args: unknown[]) => unknown) | undefined ->() - -/** Mirrors the setup-level databaseMock entry points onto the shared chain fns. */ -function delegateGlobalDbToChainMocks(): void { - for (const key of GLOBAL_DB_KEYS) { - const fn = globalDb[key] - if (typeof fn?.mockImplementation !== 'function') continue - if (!savedGlobalDbImpls.has(key)) savedGlobalDbImpls.set(key, fn.getMockImplementation()) - fn.mockImplementation((...args: unknown[]) => (dbChainMockFns[key] as Mock)(...args)) - } -} - -/** Restores the databaseMock entry points captured before this suite ran. */ -function restoreGlobalDb(): void { - for (const [key, impl] of savedGlobalDbImpls) { - if (impl) globalDb[key].mockImplementation(impl) - else globalDb[key].mockReset() - } +function fromTables(): unknown[] { + return dbChainMockFns.from.mock.calls.map(([table]) => table) } interface SubRow { @@ -124,32 +55,21 @@ function orgEnterprise(orgId: string): SubRow { return { id: 'sub-org-enterprise', referenceId: orgId, plan: 'enterprise', status: 'active' } } -function queue(table: TableName, rows: unknown[]) { - resultsByTable[table].push(rows) -} - describe('getHighestPrioritySubscription', () => { beforeEach(() => { vi.clearAllMocks() resetDbChainMock() - resultsByTable.subscription = [] - resultsByTable.member = [] - resultsByTable.organization = [] - fromCalls.length = 0 - dbChainMockFns.select.mockImplementation(routedSelect) - delegateGlobalDbToChainMocks() }) afterAll(() => { resetDbChainMock() - restoreGlobalDb() }) it('picks the org Enterprise sub over a personal Pro sub (priority order)', async () => { - queue('subscription', [personalPro('user-1')]) // personalSubs query - queue('member', [{ organizationId: 'org-1' }]) // memberships query - queue('organization', [{ id: 'org-1' }]) // org-existence query - queue('subscription', [orgEnterprise('org-1')]) // org-subscriptions query + queueTableRows(subscription, [personalPro('user-1')]) // personalSubs query + queueTableRows(member, [{ organizationId: 'org-1' }]) // memberships query + queueTableRows(organization, [{ id: 'org-1' }]) // org-existence query + queueTableRows(subscription, [orgEnterprise('org-1')]) // org-subscriptions query const result = await getHighestPrioritySubscription('user-1') @@ -159,10 +79,10 @@ describe('getHighestPrioritySubscription', () => { }) it('selection is deterministic regardless of which parallelized query resolves first', async () => { - queue('subscription', [personalPro('user-1')]) - queue('member', [{ organizationId: 'org-1' }]) - queue('organization', [{ id: 'org-1' }]) - queue('subscription', [orgEnterprise('org-1')]) + queueTableRows(subscription, [personalPro('user-1')]) + queueTableRows(member, [{ organizationId: 'org-1' }]) + queueTableRows(organization, [{ id: 'org-1' }]) + queueTableRows(subscription, [orgEnterprise('org-1')]) const result = await getHighestPrioritySubscription('user-1') @@ -170,35 +90,38 @@ describe('getHighestPrioritySubscription', () => { }) it('issues BOTH the personal-subscriptions and memberships queries (parallelized pair)', async () => { - queue('subscription', [personalPro('user-1')]) - queue('member', [{ organizationId: 'org-1' }]) - queue('organization', [{ id: 'org-1' }]) - queue('subscription', [orgEnterprise('org-1')]) + queueTableRows(subscription, [personalPro('user-1')]) + queueTableRows(member, [{ organizationId: 'org-1' }]) + queueTableRows(organization, [{ id: 'org-1' }]) + queueTableRows(subscription, [orgEnterprise('org-1')]) await getHighestPrioritySubscription('user-1') - expect(fromCalls).toContain('subscription') - expect(fromCalls).toContain('member') + expect(fromTables()).toContain(subscription) + expect(fromTables()).toContain(member) // First two queries are exactly the parallelized pair (in either order). - expect(fromCalls.slice(0, 2).sort()).toEqual(['member', 'subscription']) + const firstTwo = fromTables().slice(0, 2) + expect(firstTwo).toHaveLength(2) + expect(firstTwo).toContain(subscription) + expect(firstTwo).toContain(member) }) it('returns the personal sub and skips org follow-ups when there are no memberships', async () => { - queue('subscription', [personalPro('user-1')]) - queue('member', []) + queueTableRows(subscription, [personalPro('user-1')]) + queueTableRows(member, []) const result = await getHighestPrioritySubscription('user-1') expect(result?.id).toBe('sub-personal-pro') expect(result?.plan).toBe('pro') // org-existence + org-subscription follow-ups are NOT issued. - expect(fromCalls).not.toContain('organization') - expect(fromCalls.filter((t) => t === 'subscription')).toHaveLength(1) + expect(fromTables()).not.toContain(organization) + expect(fromTables().filter((t) => t === subscription)).toHaveLength(1) }) it('returns null when neither personal nor org subscriptions exist', async () => { - queue('subscription', []) - queue('member', []) + queueTableRows(subscription, []) + queueTableRows(member, []) const result = await getHighestPrioritySubscription('user-1') @@ -206,27 +129,27 @@ describe('getHighestPrioritySubscription', () => { }) it('excludes orphaned org memberships whose organization row no longer exists', async () => { - queue('subscription', []) - queue('member', [{ organizationId: 'ghost-org' }]) // membership points at a deleted org - queue('organization', []) + queueTableRows(subscription, []) + queueTableRows(member, [{ organizationId: 'ghost-org' }]) // membership points at a deleted org + queueTableRows(organization, []) const result = await getHighestPrioritySubscription('user-1') // Org subs are never fetched (no valid org ids) -> falls back to null. expect(result).toBeNull() - expect(fromCalls).toContain('organization') + expect(fromTables()).toContain(organization) // Only the initial personal-subs read on `subscription`; org-subs query skipped. - expect(fromCalls.filter((t) => t === 'subscription')).toHaveLength(1) + expect(fromTables().filter((t) => t === subscription)).toHaveLength(1) }) it('falls back to the personal sub when the only org is orphaned', async () => { - queue('subscription', [personalPro('user-1')]) - queue('member', [{ organizationId: 'ghost-org' }]) - queue('organization', []) + queueTableRows(subscription, [personalPro('user-1')]) + queueTableRows(member, [{ organizationId: 'ghost-org' }]) + queueTableRows(organization, []) const result = await getHighestPrioritySubscription('user-1') expect(result?.id).toBe('sub-personal-pro') - expect(fromCalls.filter((t) => t === 'subscription')).toHaveLength(1) + expect(fromTables().filter((t) => t === subscription)).toHaveLength(1) }) }) diff --git a/apps/sim/lib/billing/core/usage-log.test.ts b/apps/sim/lib/billing/core/usage-log.test.ts index 23206179b9f..6b1861ed1ac 100644 --- a/apps/sim/lib/billing/core/usage-log.test.ts +++ b/apps/sim/lib/billing/core/usage-log.test.ts @@ -2,8 +2,7 @@ * @vitest-environment node */ import { usageLog } from '@sim/db/schema' -import { databaseMock, dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' -import type { Mock } from 'vitest' +import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { @@ -49,59 +48,17 @@ import { } from '@/lib/billing/core/usage-log' /** - * `@sim/db` behavior is driven through the SHARED mock instances rather than a - * file-local factory object. This file mocks `@sim/db` with `dbChainMock` and - * wires its `insert` / `transaction` entry points to this file's `mockInsert` - * / `mockTransaction` in `beforeEach`; the setup-level `databaseMock` entry - * points are mirrored onto the same chain fns. Under `isolate: false` the - * module under test may have been loaded by an earlier suite in this worker - * with `@sim/db` bound to `databaseMock` — configuring both shared instances - * keeps either binding correct. + * Re-wires the shared db mocks (`dbChainMockFns`, backing the single shared + * `@sim/db` mock instance) to this file's insert/transaction chain. */ -const GLOBAL_DB_KEYS = [ - 'select', - 'selectDistinct', - 'insert', - 'update', - 'delete', - 'transaction', -] as const - -const globalDb = databaseMock.db as unknown as Record<(typeof GLOBAL_DB_KEYS)[number], Mock> -const savedGlobalDbImpls = new Map< - (typeof GLOBAL_DB_KEYS)[number], - ((...args: unknown[]) => unknown) | undefined ->() - -/** Mirrors the setup-level databaseMock entry points onto the shared chain fns. */ -function delegateGlobalDbToChainMocks(): void { - for (const key of GLOBAL_DB_KEYS) { - const fn = globalDb[key] - if (typeof fn?.mockImplementation !== 'function') continue - if (!savedGlobalDbImpls.has(key)) savedGlobalDbImpls.set(key, fn.getMockImplementation()) - fn.mockImplementation((...args: unknown[]) => (dbChainMockFns[key] as Mock)(...args)) - } -} - -/** Restores the databaseMock entry points captured before this suite ran. */ -function restoreGlobalDb(): void { - for (const [key, impl] of savedGlobalDbImpls) { - if (impl) globalDb[key].mockImplementation(impl) - else globalDb[key].mockReset() - } -} - -/** Re-wires the shared db mocks to this file's insert/transaction chain. */ function installSharedDbMocks(): void { resetDbChainMock() dbChainMockFns.insert.mockImplementation((...args: unknown[]) => mockInsert(...args)) dbChainMockFns.transaction.mockImplementation((...args: unknown[]) => mockTransaction(...args)) - delegateGlobalDbToChainMocks() } afterAll(() => { resetDbChainMock() - restoreGlobalDb() }) describe('recordUsage', () => { diff --git a/apps/sim/lib/billing/core/usage.test.ts b/apps/sim/lib/billing/core/usage.test.ts index 443cea29f7a..a46caa1517a 100644 --- a/apps/sim/lib/billing/core/usage.test.ts +++ b/apps/sim/lib/billing/core/usage.test.ts @@ -8,55 +8,13 @@ * * @vitest-environment node */ -import { databaseMock, dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' -import type { Mock } from 'vitest' +import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' vi.mock('@sim/db', () => dbChainMock) -/** - * Under `isolate: false` the module under test may have been loaded by an - * earlier suite in this shared worker with `@sim/db` bound to the setup-level - * `databaseMock` instead of this file's `dbChainMock`. Delegating the - * databaseMock entry points to the same shared chain fns keeps either binding - * correct. - */ -const GLOBAL_DB_KEYS = [ - 'select', - 'selectDistinct', - 'insert', - 'update', - 'delete', - 'transaction', -] as const - -const globalDb = databaseMock.db as unknown as Record<(typeof GLOBAL_DB_KEYS)[number], Mock> -const savedGlobalDbImpls = new Map< - (typeof GLOBAL_DB_KEYS)[number], - ((...args: unknown[]) => unknown) | undefined ->() - -/** Mirrors the setup-level databaseMock entry points onto the shared chain fns. */ -function delegateGlobalDbToChainMocks(): void { - for (const key of GLOBAL_DB_KEYS) { - const fn = globalDb[key] - if (typeof fn?.mockImplementation !== 'function') continue - if (!savedGlobalDbImpls.has(key)) savedGlobalDbImpls.set(key, fn.getMockImplementation()) - fn.mockImplementation((...args: unknown[]) => (dbChainMockFns[key] as Mock)(...args)) - } -} - -/** Restores the databaseMock entry points captured before this suite ran. */ -function restoreGlobalDb(): void { - for (const [key, impl] of savedGlobalDbImpls) { - if (impl) globalDb[key].mockImplementation(impl) - else globalDb[key].mockReset() - } -} - afterAll(() => { resetDbChainMock() - restoreGlobalDb() }) const { @@ -131,7 +89,6 @@ describe('getUserUsageLimit', () => { beforeEach(() => { vi.clearAllMocks() resetDbChainMock() - delegateGlobalDbToChainMocks() mockIsOrgScopedSubscription.mockReturnValue(false) mockGetHighestPrioritySubscription.mockResolvedValue(null) }) @@ -213,7 +170,6 @@ describe('syncUsageLimitsFromSubscription', () => { beforeEach(() => { vi.clearAllMocks() resetDbChainMock() - delegateGlobalDbToChainMocks() mockIsOrgScopedSubscription.mockReturnValue(false) }) diff --git a/apps/sim/lib/workspaces/utils.test.ts b/apps/sim/lib/workspaces/utils.test.ts index 9cbf74c2c3e..a1cd55d2ad7 100644 --- a/apps/sim/lib/workspaces/utils.test.ts +++ b/apps/sim/lib/workspaces/utils.test.ts @@ -1,8 +1,7 @@ /** * @vitest-environment node */ -import { databaseMock, dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' -import type { Mock } from 'vitest' +import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { mockChangeWorkspaceStoragePayerInTx } = vi.hoisted(() => ({ @@ -21,55 +20,8 @@ import { reassignWorkflowOwnershipForWorkspaceMemberRemovalTx, } from '@/lib/workspaces/utils' -/** - * `@sim/db` behavior is driven through the SHARED `dbChainMockFns` instances - * instead of a file-local factory object. Under `isolate: false` the module - * under test may have been loaded by an earlier suite in this shared worker - * with `@sim/db` bound to the setup-level `databaseMock` instead of this - * file's `dbChainMock`; delegating the databaseMock entry points to the same - * chain fns keeps either binding correct. - */ -const mockDb = { - select: dbChainMockFns.select, - transaction: dbChainMockFns.transaction, -} - -const GLOBAL_DB_KEYS = [ - 'select', - 'selectDistinct', - 'insert', - 'update', - 'delete', - 'transaction', -] as const - -const globalDb = databaseMock.db as unknown as Record<(typeof GLOBAL_DB_KEYS)[number], Mock> -const savedGlobalDbImpls = new Map< - (typeof GLOBAL_DB_KEYS)[number], - ((...args: unknown[]) => unknown) | undefined ->() - -/** Mirrors the setup-level databaseMock entry points onto the shared chain fns. */ -function delegateGlobalDbToChainMocks(): void { - for (const key of GLOBAL_DB_KEYS) { - const fn = globalDb[key] - if (typeof fn?.mockImplementation !== 'function') continue - if (!savedGlobalDbImpls.has(key)) savedGlobalDbImpls.set(key, fn.getMockImplementation()) - fn.mockImplementation((...args: unknown[]) => (dbChainMockFns[key] as Mock)(...args)) - } -} - -/** Restores the databaseMock entry points captured before this suite ran. */ -function restoreGlobalDb(): void { - for (const [key, impl] of savedGlobalDbImpls) { - if (impl) globalDb[key].mockImplementation(impl) - else globalDb[key].mockReset() - } -} - afterAll(() => { resetDbChainMock() - restoreGlobalDb() }) function createMockChain(finalResult: unknown) { @@ -114,7 +66,6 @@ describe('reassignBilledAccountForUser', () => { beforeEach(() => { vi.clearAllMocks() resetDbChainMock() - delegateGlobalDbToChainMocks() }) it('routes each resolved workspace through the payer helper in its own transaction', async () => { @@ -122,19 +73,19 @@ describe('reassignBilledAccountForUser', () => { const tx = { update: vi.fn().mockReturnValue(updateChain), } - mockDb.select.mockReturnValueOnce( + dbChainMockFns.select.mockReturnValueOnce( createMockChain([ { id: 'workspace-personal', ownerId: 'owner-1', organizationId: null }, { id: 'workspace-org', ownerId: 'owner-2', organizationId: 'org-1' }, ]) ) - mockDb.transaction.mockImplementation( + dbChainMockFns.transaction.mockImplementation( async (callback: (transaction: typeof tx) => Promise) => callback(tx) ) const result = await reassignBilledAccountForUser('departing-user') - expect(mockDb.transaction).toHaveBeenCalledTimes(2) + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(2) expect(mockChangeWorkspaceStoragePayerInTx).toHaveBeenNthCalledWith(1, tx, { workspaceId: 'workspace-personal', organizationId: null, @@ -169,7 +120,7 @@ describe('reassignBilledAccountForUser', () => { const tx = { update: vi.fn().mockReturnValue(updateChain), } - mockDb.select + dbChainMockFns.select .mockReturnValueOnce( createMockChain([ { id: 'workspace-admin', ownerId: 'departing-user', organizationId: null }, @@ -178,7 +129,7 @@ describe('reassignBilledAccountForUser', () => { ) .mockReturnValueOnce(createMockChain([{ userId: 'admin-1' }])) .mockReturnValueOnce(createMockChain([])) - mockDb.transaction.mockImplementation( + dbChainMockFns.transaction.mockImplementation( async (callback: (transaction: typeof tx) => Promise) => callback(tx) ) @@ -203,13 +154,13 @@ describe('reassignBilledAccountForUser', () => { const tx = { update: vi.fn().mockReturnValue(updateChain), } - mockDb.select.mockReturnValueOnce( + dbChainMockFns.select.mockReturnValueOnce( createMockChain([ { id: 'workspace-1', ownerId: 'owner-1', organizationId: null }, { id: 'workspace-2', ownerId: 'owner-2', organizationId: null }, ]) ) - mockDb.transaction.mockImplementation( + dbChainMockFns.transaction.mockImplementation( async (callback: (transaction: typeof tx) => Promise) => callback(tx) ) mockChangeWorkspaceStoragePayerInTx.mockRejectedValueOnce(new Error('payer transfer failed')) @@ -218,7 +169,7 @@ describe('reassignBilledAccountForUser', () => { 'payer transfer failed' ) - expect(mockDb.transaction).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(1) expect(mockChangeWorkspaceStoragePayerInTx).toHaveBeenCalledTimes(1) expect(tx.update).not.toHaveBeenCalled() }) @@ -228,7 +179,6 @@ describe('reassignWorkflowOwnershipForWorkspaceMemberRemovalTx', () => { beforeEach(() => { vi.clearAllMocks() resetDbChainMock() - delegateGlobalDbToChainMocks() }) it('reassigns departing member workflows to the workspace billed account', async () => { @@ -327,13 +277,12 @@ describe('listAccessibleWorkspaceRowsForUser', () => { beforeEach(() => { vi.clearAllMocks() resetDbChainMock() - delegateGlobalDbToChainMocks() }) it('elevates an org admin to admin on an org workspace where they hold a lower explicit grant', async () => { const orgWorkspace = { id: 'ws-1', name: 'Shared', ownerId: 'owner-x', organizationId: 'org-1' } - mockDb.select + dbChainMockFns.select .mockReturnValueOnce(createMockChain([{ workspace: orgWorkspace, permissionType: 'write' }])) .mockReturnValueOnce(createMockChain([{ organizationId: 'org-1', role: 'admin' }])) .mockReturnValueOnce(createMockChain([orgWorkspace])) @@ -352,7 +301,7 @@ describe('listAccessibleWorkspaceRowsForUser', () => { } const orgWorkspace = { id: 'ws-1', name: 'Shared', ownerId: 'owner-x', organizationId: 'org-1' } - mockDb.select + dbChainMockFns.select .mockReturnValueOnce( createMockChain([{ workspace: externalWorkspace, permissionType: 'write' }]) ) diff --git a/packages/testing/package.json b/packages/testing/package.json index 216129ae6c1..ee8f36d1bff 100644 --- a/packages/testing/package.json +++ b/packages/testing/package.json @@ -42,6 +42,7 @@ } }, "scripts": { + "test": "vitest run", "type-check": "tsc --noEmit", "lint": "biome check --write --unsafe .", "lint:check": "biome check .", diff --git a/packages/testing/src/mocks/database.mock.test.ts b/packages/testing/src/mocks/database.mock.test.ts new file mode 100644 index 00000000000..c32fc940156 --- /dev/null +++ b/packages/testing/src/mocks/database.mock.test.ts @@ -0,0 +1,156 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + databaseMock, + dbChainMock, + dbChainMockFns, + queueTableRows, + resetDbChainMock, +} from './database.mock' + +const workflowTable = { id: 'id', name: 'name' } +const memberTable = { id: 'id', userId: 'userId' } + +type MockDb = Record + +const db = dbChainMock.db as MockDb + +describe('database mock', () => { + beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + }) + + it('shares one db instance between dbChainMock and databaseMock', () => { + expect(databaseMock.db).toBe(dbChainMock.db) + expect(databaseMock.dbReplica).toBe(dbChainMock.db) + expect(dbChainMock.dbFor()).toBe(dbChainMock.db) + }) + + it('resolves empty arrays by default at every terminal', async () => { + await expect(db.select().from(workflowTable).where({})).resolves.toEqual([]) + await expect(db.select().from(workflowTable).where({}).limit(5)).resolves.toEqual([]) + await expect(db.select().from(workflowTable).where({}).orderBy('id')).resolves.toEqual([]) + await expect(db.insert(workflowTable).values({}).returning()).resolves.toEqual([]) + await expect(db.update(workflowTable).set({}).where({})).resolves.toEqual([]) + await expect(db.delete(workflowTable).where({})).resolves.toEqual([]) + }) + + it('routes queued rows to the chain reading that table', async () => { + queueTableRows(workflowTable, [{ id: 'w-1' }]) + queueTableRows(memberTable, [{ id: 'm-1' }, { id: 'm-2' }]) + + await expect(db.select().from(memberTable).where({})).resolves.toEqual([ + { id: 'm-1' }, + { id: 'm-2' }, + ]) + await expect(db.select().from(workflowTable).where({})).resolves.toEqual([{ id: 'w-1' }]) + }) + + it('consumes queued sets FIFO per table and falls back to empty', async () => { + queueTableRows(workflowTable, [{ id: 'first' }]) + queueTableRows(workflowTable, [{ id: 'second' }]) + + await expect(db.select().from(workflowTable).where({})).resolves.toEqual([{ id: 'first' }]) + await expect(db.select().from(workflowTable).where({})).resolves.toEqual([{ id: 'second' }]) + await expect(db.select().from(workflowTable).where({})).resolves.toEqual([]) + }) + + it('resolves queued rows through downstream terminals (limit/orderBy/joins)', async () => { + queueTableRows(workflowTable, [{ id: 'w-1' }]) + await expect(db.select().from(workflowTable).where({}).limit(1)).resolves.toEqual([ + { id: 'w-1' }, + ]) + + queueTableRows(workflowTable, [{ id: 'w-2' }]) + await expect( + db.select().from(workflowTable).innerJoin(memberTable, {}).where({}).orderBy('id') + ).resolves.toEqual([{ id: 'w-2' }]) + }) + + it('resolves queued rows when a from-chain is awaited directly (no where)', async () => { + queueTableRows(workflowTable, [{ id: 'direct' }]) + await expect(db.select().from(workflowTable)).resolves.toEqual([{ id: 'direct' }]) + await expect(db.select().from(workflowTable)).resolves.toEqual([]) + }) + + it('routes two direct-await builders constructed before either resolves', async () => { + queueTableRows(workflowTable, [{ id: 'w-first' }]) + queueTableRows(memberTable, [{ id: 'm-second' }]) + const first = db.select().from(workflowTable) + const second = db.select().from(memberTable) + await expect(first).resolves.toEqual([{ id: 'w-first' }]) + await expect(second).resolves.toEqual([{ id: 'm-second' }]) + }) + + it('routes rows queued for a table referenced only by a join', async () => { + queueTableRows(memberTable, [{ id: 'joined' }]) + await expect( + db.select().from(workflowTable).leftJoin(memberTable, {}).where({}) + ).resolves.toEqual([{ id: 'joined' }]) + }) + + it('prefers the from-table queue over a join-table queue', async () => { + queueTableRows(workflowTable, [{ id: 'from-row' }]) + queueTableRows(memberTable, [{ id: 'join-row' }]) + await expect( + db.select().from(workflowTable).innerJoin(memberTable, {}).where({}) + ).resolves.toEqual([{ id: 'from-row' }]) + }) + + it('never lets mutation chains consume select queues', async () => { + queueTableRows(workflowTable, [{ id: 'kept' }]) + await expect(db.update(workflowTable).set({}).where({})).resolves.toEqual([]) + await expect(db.delete(workflowTable).where({})).resolves.toEqual([]) + await expect(db.select().from(workflowTable).where({})).resolves.toEqual([{ id: 'kept' }]) + }) + + it('routes selectDistinctOn chains through the same table queues', async () => { + queueTableRows(memberTable, [{ id: 'm-1' }]) + await expect(db.selectDistinctOn(['id']).from(memberTable).where({})).resolves.toEqual([ + { id: 'm-1' }, + ]) + expect(dbChainMockFns.selectDistinctOn).toHaveBeenCalledTimes(1) + }) + + it('lets per-test ...Once overrides win over queued rows downstream', async () => { + queueTableRows(workflowTable, [{ id: 'queued' }]) + dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'override' }]) + await expect(db.select().from(workflowTable).where({}).limit(1)).resolves.toEqual([ + { id: 'override' }, + ]) + }) + + it('preserves a queued set when a terminal override resolves the chain', async () => { + queueTableRows(workflowTable, [{ id: 'queued' }]) + dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'override' }]) + await expect(db.select().from(workflowTable).where({}).limit(1)).resolves.toEqual([ + { id: 'override' }, + ]) + await expect(db.select().from(workflowTable).where({})).resolves.toEqual([{ id: 'queued' }]) + }) + + it('restores directly-overridden db entry points on resetDbChainMock', async () => { + ;(db.select as ReturnType).mockImplementation(() => { + throw new Error('broken') + }) + expect(() => db.select()).toThrow('broken') + resetDbChainMock() + await expect(db.select().from(workflowTable).where({})).resolves.toEqual([]) + }) + + it('clears queues and rewires defaults on resetDbChainMock', async () => { + queueTableRows(workflowTable, [{ id: 'stale' }]) + dbChainMockFns.where.mockReturnValue('broken' as never) + resetDbChainMock() + await expect(db.select().from(workflowTable).where({})).resolves.toEqual([]) + }) + + it('runs transactions against the same shared instance', async () => { + queueTableRows(workflowTable, [{ id: 'tx-row' }]) + const rows = await db.transaction(async (tx: MockDb) => { + expect(tx).toBe(db) + return tx.select().from(workflowTable).where({}) + }) + expect(rows).toEqual([{ id: 'tx-row' }]) + }) +}) diff --git a/packages/testing/src/mocks/database.mock.ts b/packages/testing/src/mocks/database.mock.ts index c70960d8eb5..a432bc43972 100644 --- a/packages/testing/src/mocks/database.mock.ts +++ b/packages/testing/src/mocks/database.mock.ts @@ -11,13 +11,11 @@ export function createMockSql() { toSQL: () => ({ sql: strings.join('?'), params: values }), }) - // Add sql.raw method used by some queries sqlFn.raw = (rawSql: string) => ({ rawSql, toSQL: () => ({ sql: rawSql, params: [] }), }) - // Add sql.join method used to combine multiple SQL fragments sqlFn.join = (fragments: any[], separator: any) => ({ fragments, separator, @@ -62,113 +60,209 @@ export function createMockSqlOperators() { } } +/** + * Table-routed result queues. + * + * `queueTableRows(schemaMock.member, [rowA, rowB])` enqueues one result set for + * the next select chain whose `.from()` (or a subsequent `innerJoin`/ + * `leftJoin`) references that table. Each chain consumes at most one queued + * set (FIFO per table, `.from()` table checked before join tables); chains + * against tables with no queued sets resolve the chain-fn defaults (empty + * array). Mutation chains (`update`/`delete`/`insert`) never consume select + * queues. Queues are cleared by `resetDbChainMock()`. + * + * The queue is keyed by table object identity, so pass the same schema-mock + * table object the code under test passes to `.from()` / the join. + */ +const tableRowQueues = new Map() + +/** + * Enqueues one result set for the next select chain reading `table`. + */ +export function queueTableRows(table: unknown, rows: unknown[]): void { + const queue = tableRowQueues.get(table) + if (queue) queue.push(rows) + else tableRowQueues.set(table, [rows]) +} + +/** Dequeues the first queued set among the given chain's tables (from before joins). */ +function dequeueChainRows(tables: unknown[]): unknown[] | null { + for (const table of tables) { + const queue = tableRowQueues.get(table) + if (queue && queue.length > 0) return queue.shift() ?? null + } + return null +} + /** * Pre-wired chain of vi.fn()s for drizzle-style DB queries. * - * Each builder step is a stable, module-level `vi.fn()` — safe to reference - * inside hoisted `vi.mock()` factories (same pattern as `authMockFns`). Chains - * are wired at module load time: + * Each chain step is recorded on a stable, module-level `vi.fn()` spy + * (`dbChainMockFns.*`) — safe to reference inside hoisted `vi.mock()` + * factories (same pattern as `authMockFns`): * * - `select().from().where()` → returns a builder with `.limit` / `.orderBy` / * `.returning` / `.groupBy` / `.for` terminals * - `select().from().innerJoin()|leftJoin()` → returns the same where-builder * - `insert().values().returning()` / `update().set().where()` / `delete().where()` * - * Terminals (`limit`, `orderBy`, `returning`, `groupBy`, `for`, `values`) - * default to resolving `[]` (or `undefined` for `values`). Override per-test - * with `dbChainMockFns.limit.mockResolvedValueOnce([...])`. `for` mirrors - * drizzle's `.for('update')` — it returns a Promise with `.limit` / `.orderBy` - * / `.returning` / `.groupBy` attached, so both `await .where().for('update')` - * (terminal) and `await .where().for('update').limit(1)` (chained) work. - * Override the terminal result with `dbChainMockFns.for.mockResolvedValueOnce( - * [...])`; override the chained result by mocking the downstream terminal - * (e.g. `dbChainMockFns.limit.mockResolvedValueOnce([...])`). + * Results resolve, in priority order: + * 1. a per-test override (`dbChainMockFns.limit.mockResolvedValueOnce([...])`) + * 2. rows queued for one of the chain's tables via `queueTableRows` + * 3. the default empty array + * + * Routing state lives in per-chain closures: each `select().from(t)` captures + * its own table list, so partially-built chains for different tables can be + * interleaved or awaited in any order without cross-talk. The shared spies + * carry only call history and per-test overrides — a spy's default + * implementation returns a sentinel that the chain replaces with the + * chain-local builder, while any `mock*` override on the spy wins verbatim. + * + * `for` mirrors drizzle's `.for('update')` — it returns a Promise with + * `.limit` / `.orderBy` / `.returning` / `.groupBy` attached, so both + * `await .where().for('update')` (terminal) and + * `await .where().for('update').limit(1)` (chained) work. * * `vi.clearAllMocks()` clears call history but preserves default wiring. Tests - * that replace a wiring with `mockReturnValue(...)` (not `...Once`) must re-wire - * in their own `beforeEach`. + * that replace a wiring with `mockReturnValue(...)` (not `...Once`) must + * re-wire via `resetDbChainMock()` in their own `beforeEach`. * * @example * ```ts - * import { dbChainMock, dbChainMockFns } from '@sim/testing' - * vi.mock('@sim/db', () => dbChainMock) + * import { dbChainMockFns, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing' + * + * beforeEach(() => { + * vi.clearAllMocks() + * resetDbChainMock() + * }) * * it('finds rows', async () => { - * dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'w-1' }]) - * // ... exercise code that hits db.select().from().where().limit() ... + * queueTableRows(schemaMock.workflow, [{ id: 'w-1' }]) + * // ... exercise code that hits db.select().from(workflow).where() ... * expect(dbChainMockFns.where).toHaveBeenCalled() * }) * ``` */ -const offset = vi.fn(() => Promise.resolve([] as unknown[])) -// `.limit()` returns a builder that is awaitable (default empty page) and also -// exposes `.offset()` for keyset/OFFSET paging (`.limit(n).offset(m)`). -const limitBuilder = () => { - const thenable: any = Promise.resolve([] as unknown[]) - thenable.offset = offset - return thenable -} -const limit = vi.fn(limitBuilder) +const CHAIN_DEFAULT = Symbol('db-chain-default') + +type ChainSpy = ReturnType any>> + +const chainSpy = (): ChainSpy => vi.fn((..._args: any[]) => CHAIN_DEFAULT as any) + +/** + * Records the call on the shared spy, honoring any per-test override; when the + * spy still has its default implementation, builds the chain-local default. + */ +const spyOrDefault = (spy: ChainSpy, buildDefault: (...args: any[]) => unknown) => + vi.fn((...args: any[]) => { + const result = spy(...args) + return result === CHAIN_DEFAULT ? buildDefault(...args) : result + }) + +// Shared spies: structural steps default to the sentinel (chain-local builders +// take over); value terminals keep real defaults. +const select = chainSpy() +const selectDistinct = chainSpy() +const selectDistinctOn = chainSpy() +const from = chainSpy() +const where = chainSpy() +const limit = chainSpy() +const offset = chainSpy() +const orderBy = chainSpy() +const groupBy = chainSpy() +const having = chainSpy() +const forClause = chainSpy() +const innerJoin = chainSpy() +const leftJoin = chainSpy() +const insert = chainSpy() +const update = chainSpy() +const set = chainSpy() +const del = chainSpy() const returning = vi.fn(() => Promise.resolve([] as unknown[])) const execute = vi.fn(() => Promise.resolve([] as unknown[])) +const query = vi.fn(() => Promise.resolve([] as unknown[])) +const onConflictDoUpdate = vi.fn(() => ({ returning }) as unknown as Promise) +const onConflictDoNothing = vi.fn(() => ({ returning }) as unknown as Promise) +const values = vi.fn(() => ({ returning, onConflictDoUpdate, onConflictDoNothing })) +const transaction: ReturnType = vi.fn( + async (cb: (tx: any) => unknown): Promise => cb(dbChainMock.db) +) -const terminalBuilder = () => { - const thenable: any = Promise.resolve([] as unknown[]) - thenable.limit = limit - thenable.orderBy = orderBy - thenable.returning = returning - thenable.groupBy = groupBy - thenable.for = forClause - return thenable +/** + * Lazy per-chain rows supplier: dequeues once, at the moment the FIRST default + * thenable actually resolves. A chain whose result comes from a per-test + * override never reaches a default resolution, so its queued set stays + * available for the next chain on that table. + */ +type RowsSupplier = () => unknown[] | null + +const chainRowsSupplier = (tables: unknown[]): RowsSupplier => { + let consumed = false + let rows: unknown[] | null = null + return () => { + if (!consumed) { + consumed = true + rows = dequeueChainRows(tables) + } + return rows + } } -const orderBy = vi.fn(terminalBuilder) -const having = vi.fn(terminalBuilder) -const groupBy = vi.fn(() => { - const builder = terminalBuilder() - builder.having = having - return builder +const noRows: RowsSupplier = () => null + +/** An awaitable chain step that resolves `getRows()` only when actually awaited. */ +const lazyRowsThenable = (getRows: RowsSupplier): any => ({ + then: (onFulfilled?: (rows: unknown[]) => unknown, onRejected?: (reason: unknown) => unknown) => + Promise.resolve((getRows() ?? []) as unknown[]).then(onFulfilled, onRejected), + catch: (onRejected?: (reason: unknown) => unknown) => + Promise.resolve((getRows() ?? []) as unknown[]).catch(onRejected), + finally: (onFinally?: () => void) => + Promise.resolve((getRows() ?? []) as unknown[]).finally(onFinally), }) -const forBuilder = terminalBuilder -const forClause = vi.fn(forBuilder) -const onConflictDoUpdate = vi.fn(() => ({ returning }) as unknown as Promise) -const onConflictDoNothing = vi.fn(() => ({ returning }) as unknown as Promise) +// `.limit()` returns a builder that is awaitable and also exposes `.offset()` +// for keyset/OFFSET paging (`.limit(n).offset(m)`). +const limitBuilder = (getRows: RowsSupplier) => { + const thenable = lazyRowsThenable(getRows) + thenable.offset = spyOrDefault(offset, () => lazyRowsThenable(getRows)) + return thenable +} -const whereBuilder = () => { - // Some call sites (e.g. `db.select().from(t).where(eq(...))` with no - // limit/orderBy) await the where directly. Make the builder a thenable so - // those calls resolve to the default empty array. - const thenable: any = Promise.resolve([] as unknown[]) - thenable.limit = limit - thenable.orderBy = orderBy +const terminalBuilder = (getRows: RowsSupplier): any => { + const thenable = lazyRowsThenable(getRows) + thenable.limit = spyOrDefault(limit, () => limitBuilder(getRows)) + thenable.orderBy = spyOrDefault(orderBy, () => terminalBuilder(getRows)) thenable.returning = returning - thenable.groupBy = groupBy - thenable.for = forClause + thenable.groupBy = spyOrDefault(groupBy, () => { + const builder = terminalBuilder(getRows) + builder.having = spyOrDefault(having, () => terminalBuilder(getRows)) + return builder + }) + thenable.for = spyOrDefault(forClause, () => terminalBuilder(getRows)) return thenable } -const where = vi.fn(whereBuilder) -const joinBuilder = (): { where: typeof where; innerJoin: any; leftJoin: any } => ({ - where, - innerJoin, - leftJoin, +// The from/join builder is itself a thenable so `await db.select().from(t)` +// (no where clause) also resolves table-routed rows; the chain's single lazy +// supplier means it never double-consumes no matter which step is awaited. +const joinBuilder = (tables: unknown[]): any => { + const getRows = chainRowsSupplier(tables) + const builder = lazyRowsThenable(getRows) + builder.where = spyOrDefault(where, () => terminalBuilder(getRows)) + builder.innerJoin = spyOrDefault(innerJoin, (table: unknown) => joinBuilder([...tables, table])) + builder.leftJoin = spyOrDefault(leftJoin, (table: unknown) => joinBuilder([...tables, table])) + return builder +} + +const selectBuilder = () => ({ + from: spyOrDefault(from, (table: unknown) => joinBuilder([table])), }) -const innerJoin: ReturnType = vi.fn(joinBuilder) -const leftJoin: ReturnType = vi.fn(joinBuilder) -const from = vi.fn(joinBuilder) -const select = vi.fn(() => ({ from })) -const selectDistinct = vi.fn(() => ({ from })) -const selectDistinctOn = vi.fn(() => ({ from })) -const values = vi.fn(() => ({ returning, onConflictDoUpdate, onConflictDoNothing })) -const insert = vi.fn(() => ({ values })) -const set = vi.fn(() => ({ where })) -const update = vi.fn(() => ({ set })) -const del = vi.fn(() => ({ where })) -const transaction: ReturnType = vi.fn( - async (cb: (tx: any) => unknown): Promise => cb(dbChainMock.db) -) +// Mutation chains route nothing: their where() resolves the plain default so a +// mutation can never consume rows queued for a select. +const mutationWhere = () => ({ + where: spyOrDefault(where, () => terminalBuilder(noRows)), +}) export const dbChainMockFns = { select, @@ -197,44 +291,78 @@ export const dbChainMockFns = { } /** - * Re-applies the default chain wiring to every `dbChainMockFns` entry. Call - * this in `beforeEach` (after `vi.clearAllMocks()`) if any test uses - * `mockReturnValue` / `mockResolvedValue` (permanent overrides) — this + * 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. */ export function resetDbChainMock(): void { - select.mockImplementation(() => ({ from })) - selectDistinct.mockImplementation(() => ({ from })) - selectDistinctOn.mockImplementation(() => ({ from })) - from.mockImplementation(joinBuilder) - innerJoin.mockImplementation(joinBuilder) - leftJoin.mockImplementation(joinBuilder) - where.mockImplementation(whereBuilder) - insert.mockImplementation(() => ({ values })) - values.mockImplementation(() => ({ returning, onConflictDoUpdate, onConflictDoNothing })) - onConflictDoUpdate.mockImplementation(() => ({ returning }) as unknown as Promise) - onConflictDoNothing.mockImplementation(() => ({ returning }) as unknown as Promise) - update.mockImplementation(() => ({ set })) - set.mockImplementation(() => ({ where })) - del.mockImplementation(() => ({ where })) - limit.mockImplementation(limitBuilder) - offset.mockImplementation(() => Promise.resolve([] as unknown[])) - orderBy.mockImplementation(terminalBuilder) + 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) + } returning.mockImplementation(() => Promise.resolve([] as unknown[])) - having.mockImplementation(terminalBuilder) - groupBy.mockImplementation(() => { - const builder = terminalBuilder() - builder.having = having - return builder - }) execute.mockImplementation(() => Promise.resolve([] as unknown[])) - forClause.mockImplementation(forBuilder) + 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()). + for (const key of [ + 'select', + 'selectDistinct', + 'selectDistinctOn', + 'insert', + 'update', + 'delete', + ] as const) { + ;(dbInstance[key] as ChainSpy).mockReset() + } +} + +/** + * The single shared `@sim/db` mock instance backing BOTH `dbChainMock` and + * `databaseMock`. Because every binding resolves to the same chain spies, a + * module bound to either export behaves identically — there is exactly one + * db-mock state to configure and reset. + */ +const dbInstance = { + select: spyOrDefault(select, selectBuilder), + selectDistinct: spyOrDefault(selectDistinct, selectBuilder), + selectDistinctOn: spyOrDefault(selectDistinctOn, selectBuilder), + insert: spyOrDefault(insert, () => ({ values })), + update: spyOrDefault(update, () => ({ set: spyOrDefault(set, mutationWhere) })), + delete: spyOrDefault(del, mutationWhere), + execute, + query, + transaction, } /** @@ -245,104 +373,30 @@ export function resetDbChainMock(): void { * vi.mock('@sim/db', () => dbChainMock) * ``` */ -const dbChainInstance = { - select, - selectDistinct, - selectDistinctOn, - insert, - update, - delete: del, - execute, - transaction, -} - export const dbChainMock = { - db: dbChainInstance, + db: dbInstance, /** Same instance as `db` so per-test chain overrides cover both clients. */ - dbReplica: dbChainInstance, + dbReplica: dbInstance, /** Sub-pool clients (`dbFor('cleanup' | 'exec')`) share the same instance too. */ - dbFor: () => dbChainInstance, + dbFor: () => dbInstance, runOutsideTransactionContext: (fn: () => T): T => fn(), instrumentPoolClient: (client: T): T => client, } /** - * Creates a mock database connection. - */ -export function createMockDb() { - // A `where(...)` result that is both awaitable (resolves to `[]`) and exposes - // `.limit`/`.orderBy`, so `select().from()[.leftJoin()].where()[.limit()]` - // works whether or not a terminal is chained. - const whereResult = () => { - const thenable: any = Promise.resolve([]) - thenable.limit = vi.fn(() => Promise.resolve([])) - thenable.orderBy = vi.fn(() => Promise.resolve([])) - return thenable - } - const fromBuilder = () => ({ - where: vi.fn(whereResult), - leftJoin: vi.fn(() => ({ where: vi.fn(whereResult) })), - innerJoin: vi.fn(() => ({ where: vi.fn(whereResult) })), - }) - - return { - select: vi.fn(() => ({ - from: vi.fn(fromBuilder), - })), - selectDistinct: vi.fn(() => ({ - from: vi.fn(fromBuilder), - })), - selectDistinctOn: vi.fn(() => ({ - from: vi.fn(fromBuilder), - })), - insert: vi.fn(() => ({ - values: vi.fn(() => ({ - returning: vi.fn(() => Promise.resolve([])), - onConflictDoUpdate: vi.fn(() => ({ - returning: vi.fn(() => Promise.resolve([])), - })), - onConflictDoNothing: vi.fn(() => ({ - returning: vi.fn(() => Promise.resolve([])), - })), - })), - })), - update: vi.fn(() => ({ - set: vi.fn(() => ({ - where: vi.fn(() => ({ - returning: vi.fn(() => Promise.resolve([])), - })), - })), - })), - delete: vi.fn(() => ({ - where: vi.fn(() => ({ - returning: vi.fn(() => Promise.resolve([])), - })), - })), - transaction: vi.fn(async (callback) => callback(createMockDb())), - query: vi.fn(() => Promise.resolve([])), - } -} - -/** - * Mock module for @sim/db. - * Use with vi.mock() to replace the real database. + * Mock module for `@sim/db` installed globally in vitest.setup.ts. Shares its + * `db` instance (and therefore all chain spies and table queues) with + * `dbChainMock`; additionally exposes the `sql` template tag and operator + * exports the real module provides. * * @example * ```ts * vi.mock('@sim/db', () => databaseMock) * ``` */ -const mockDbInstance = createMockDb() - export const databaseMock = { - db: mockDbInstance, - /** Same instance as `db` so per-test overrides cover both clients. */ - dbReplica: mockDbInstance, - /** Sub-pool clients (`dbFor('cleanup' | 'exec')`) share the same instance too. */ - dbFor: () => mockDbInstance, + ...dbChainMock, sql: createMockSql(), - runOutsideTransactionContext: (fn: () => T): T => fn(), - instrumentPoolClient: (client: T): T => client, ...createMockSqlOperators(), } diff --git a/packages/testing/src/mocks/index.ts b/packages/testing/src/mocks/index.ts index 5ddbd73aa95..e12a09a6b48 100644 --- a/packages/testing/src/mocks/index.ts +++ b/packages/testing/src/mocks/index.ts @@ -39,13 +39,13 @@ export { export { copilotHttpMock, copilotHttpMockFns } from './copilot-http.mock' // Database mocks export { - createMockDb, createMockSql, createMockSqlOperators, databaseMock, dbChainMock, dbChainMockFns, drizzleOrmMock, + queueTableRows, resetDbChainMock, } from './database.mock' // Encryption mocks