fix(testing): make dbChainMock .for('update') chainable with .limit()

Author: waleedlatif1

.claude/rules/sim-testing.md

@@ -224,8 +224,12 @@ it('reads a row', async () => {
 - `db.execute()` resolves `[]`
 - `db.transaction(cb)` calls cb with `dbChainMock.db`
 
-`.for('update')` (Postgres row-level locking) is supported as a terminal on
-`where` builders. Override with `dbChainMockFns.for.mockResolvedValueOnce([...])`.
+`.for('update')` (Postgres row-level locking) is supported on `where`
+builders. It returns a thenable 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([...])`; for the chained
+form, mock the downstream terminal (e.g. `dbChainMockFns.limit.mockResolvedValueOnce([...])`).
 
 All terminals default to `Promise.resolve([])`. Override per-test with `dbChainMockFns.<terminal>.mockResolvedValueOnce(...)`.
 

packages/testing/src/mocks/database.mock.ts

@@ -71,8 +71,12 @@ export function createMockSqlOperators() {
  * 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')` row-lock terminal — override with
- * `dbChainMockFns.for.mockResolvedValueOnce([...])`.
+ * 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([...])`).
  *
  * `vi.clearAllMocks()` clears call history but preserves default wiring. Tests
  * that replace a wiring with `mockReturnValue(...)` (not `...Once`) must re-wire
@@ -95,7 +99,16 @@ const orderBy = vi.fn(() => Promise.resolve([] as unknown[]))
 const returning = vi.fn(() => Promise.resolve([] as unknown[]))
 const groupBy = vi.fn(() => Promise.resolve([] as unknown[]))
 const execute = vi.fn(() => Promise.resolve([] as unknown[]))
-const forClause = vi.fn(() => Promise.resolve([] as unknown[]))
+
+const forBuilder = () => {
+  const thenable: any = Promise.resolve([] as unknown[])
+  thenable.limit = limit
+  thenable.orderBy = orderBy
+  thenable.returning = returning
+  thenable.groupBy = groupBy
+  return thenable
+}
+const forClause = vi.fn(forBuilder)
 
 const onConflictDoUpdate = vi.fn(() => ({ returning }) as unknown as Promise<void>)
 const onConflictDoNothing = vi.fn(() => ({ returning }) as unknown as Promise<void>)
@@ -177,7 +190,7 @@ export function resetDbChainMock(): void {
   returning.mockImplementation(() => Promise.resolve([] as unknown[]))
   groupBy.mockImplementation(() => Promise.resolve([] as unknown[]))
   execute.mockImplementation(() => Promise.resolve([] as unknown[]))
-  forClause.mockImplementation(() => Promise.resolve([] as unknown[]))
+  forClause.mockImplementation(forBuilder)
   transaction.mockImplementation(async (cb: (tx: typeof dbChainMock.db) => unknown) =>
     cb(dbChainMock.db)
   )