fix(mock): vi.mock with line breaks in arguments (#502)

packages/vitest/src/integrations/snapshot/port/inlineSnapshot.ts

@@ -3,6 +3,7 @@ import type MagicString from 'magic-string'
 import detectIndent from 'detect-indent'
 import { rpc } from '../../../runtime/rpc'
 import { getOriginalPos, posToNumber } from '../../../utils/source-map'
+import { getCallLastIndex } from '../../../utils'
 
 export interface InlineSnapshot {
   snapshot: string
@@ -37,43 +38,14 @@ export async function saveInlineSnapshots(
 
 const startObjectRegex = /(?:toMatchInlineSnapshot|toThrowErrorMatchingInlineSnapshot)\s*\(\s*({)/m
 
-function getEndIndex(code: string) {
-  let charIndex = -1
-  let inString: string | null = null //
-  let startedBracers = 0
-  let endedBracers = 0
-  let beforeChar: string | null = null
-  while (charIndex <= code.length) {
-    beforeChar = code[charIndex]
-    charIndex++
-    const char = code[charIndex]
-
-    const isCharString = char === '"' || char === '\'' || char === '`'
-
-    if (isCharString && beforeChar !== '\\')
-      inString = inString === char ? null : char
-
-    if (!inString) {
-      if (char === '(')
-        startedBracers++
-      if (char === ')')
-        endedBracers++
-    }
-
-    if (startedBracers && endedBracers && startedBracers === endedBracers)
-      return charIndex
-  }
-  return null
-}
-
 function replaceObjectSnap(code: string, s: MagicString, index: number, newSnap: string, indent = '') {
   code = code.slice(index)
   const startMatch = startObjectRegex.exec(code)
   if (!startMatch)
     return false
 
   code = code.slice(startMatch.index)
-  const charIndex = getEndIndex(code)
+  const charIndex = getCallLastIndex(code)
   if (charIndex === null)
     return false
 

packages/vitest/src/plugins/mock.ts

@@ -1,86 +1,16 @@
 import type { Plugin } from 'vite'
 import MagicString from 'magic-string'
+import { getCallLastIndex, getRangeStatus } from '../utils'
 
-const mockRegexp = /\b((?:vitest|vi)\s*.\s*mock\(["`'\s](.*[@\w_-]+)["`'\s])[),]{1}/
+const mockRegexp = /^ *\b((?:vitest|vi)\s*.\s*mock\(["`'\s]+(.*[@\w_-]+)["`'\s]+)[),]{1};?/gm
 const pathRegexp = /\b(?:vitest|vi)\s*.\s*(unmock|importActual|importMock)\(["`'\s](.*[@\w_-]+)["`'\s]\);?/mg
 const vitestRegexp = /import {[^}]*}.*(?=["'`]vitest["`']).*/gm
 
-const isComment = (line: string) => {
-  const commentStarts = ['//', '/*', '*']
-
-  line = line.trim()
-
-  return commentStarts.some(cmt => line.startsWith(cmt))
-}
-
-interface MockCodeblock {
-  code: string
-  declaraton: string
-  path: string
-}
-
-const parseMocks = (code: string) => {
-  const splitted = code.split('\n')
-
-  const mockCalls: Record<string, MockCodeblock> = {}
-  let mockCall = 0
-  let lineIndex = -1
-
-  while (lineIndex < splitted.length) {
-    lineIndex++
-
-    const line = splitted[lineIndex]
-
-    if (line === undefined) break
-
-    const mock = mockCalls[mockCall] || {
-      code: '',
-      declaraton: '',
-      path: '',
-    }
-
-    if (!mock.code) {
-      const started = mockRegexp.exec(line)
-
-      if (!started || isComment(line)) continue
-
-      mock.code += `${line}\n`
-      mock.declaraton = started[1]
-      mock.path = started[2]
-
-      mockCalls[mockCall] = mock
-
-      // end at the same line
-      // we parse code after vite, so it contains semicolons
-      if (line.includes(');')) {
-        mockCall++
-        continue
-      }
-
-      continue
-    }
-
-    mock.code += `${line}\n`
-
-    mockCalls[mockCall] = mock
-
-    const startNumber = (mock.code.match(/{/g) || []).length
-    const endNumber = (mock.code.match(/}/g) || []).length
-
-    // we parse code after vite, so it contains semicolons
-    if (line.includes(');')) {
-      /**
-       * Check if number of {} is equal or this:
-       * vi.mock('path', () =>
-       *  loadStore()
-       * );
-       */
-      if (startNumber === endNumber || (startNumber === 0 && endNumber === 0))
-        mockCall++
-    }
-  }
-
-  return Object.values(mockCalls)
+const getMockLastIndex = (code: string): number | null => {
+  const index = getCallLastIndex(code)
+  if (index === null)
+    return null
+  return code[index + 1] === ';' ? index + 2 : index + 1
 }
 
 const getMethodCall = (method: string, actualPath: string, importPath: string) => {
@@ -111,19 +41,36 @@ export const MocksPlugin = (): Plugin => {
         m.overwrite(start, end, overwrite)
       }
 
-      if (mockRegexp.exec(code)) {
+      const mocks = code.matchAll(mockRegexp)
+
+      let previousIndex = 0
+
+      for (const mockResult of mocks) {
         // we need to parse parsed string because factory may contain importActual
-        const mocks = parseMocks(m?.toString() || code)
+        const lastIndex = getMockLastIndex(code.slice(mockResult.index!))
+        const [, declaration, path] = mockResult
 
-        for (const mock of mocks) {
-          const filepath = await this.resolve(mock.path, id)
+        if (lastIndex === null) continue
 
-          m ??= new MagicString(code)
+        const startIndex = mockResult.index!
 
-          const overwrite = getMethodCall('mock', filepath?.id || mock.path, mock.path)
+        const { insideComment, insideString } = getRangeStatus(code, previousIndex, startIndex)
 
-          m.prepend(mock.code.replace(mock.declaraton, overwrite))
-        }
+        if (insideComment || insideString)
+          continue
+
+        previousIndex = startIndex
+        const endIndex = startIndex + lastIndex
+
+        const filepath = await this.resolve(path, id)
+
+        m ??= new MagicString(code)
+
+        const overwrite = getMethodCall('mock', filepath?.id || path, path)
+
+        m.overwrite(startIndex, startIndex + declaration.length, overwrite)
+        m.prepend(`${m.slice(startIndex, endIndex)}\n`)
+        m.remove(startIndex, endIndex)
       }
 
       if (m) {

packages/vitest/src/utils/index.ts

@@ -110,4 +110,87 @@ export function deepMerge(target: any, source: any): any {
   return target
 }
 
+/**
+ * If code starts with a function call, will return its last index, respecting arguments.
+ * This will return 25 - last ending character of toMatch ")"
+ * Also works with callbacks
+ * ```
+ * toMatch({ test: '123' });
+ * toBeAliased('123')
+ * ```
+ */
+export function getCallLastIndex(code: string) {
+  let charIndex = -1
+  let inString: string | null = null
+  let startedBracers = 0
+  let endedBracers = 0
+  let beforeChar: string | null = null
+  while (charIndex <= code.length) {
+    beforeChar = code[charIndex]
+    charIndex++
+    const char = code[charIndex]
+
+    const isCharString = char === '"' || char === '\'' || char === '`'
+
+    if (isCharString && beforeChar !== '\\') {
+      if (inString === char)
+        inString = null
+      else if (!inString)
+        inString = char
+    }
+
+    if (!inString) {
+      if (char === '(')
+        startedBracers++
+      if (char === ')')
+        endedBracers++
+    }
+
+    if (startedBracers && endedBracers && startedBracers === endedBracers)
+      return charIndex
+  }
+  return null
+}
+
+export const getRangeStatus = (code: string, from: number, to: number) => {
+  let index = 0
+  let started = false
+  let ended = true
+  let inString: string | null = null
+  let beforeChar: string | null = null
+
+  while (index <= to) {
+    const char = code[index]
+    const sub = code[index] + code[index + 1]
+
+    const isCharString = char === '"' || char === '\'' || char === '`'
+
+    if (isCharString && beforeChar !== '\\') {
+      if (inString === char)
+        inString = null
+      else if (!inString)
+        inString = char
+    }
+
+    if (!inString && index >= from) {
+      if (sub === '/*') {
+        started = true
+        ended = false
+      }
+      if (sub === '*/' && started) {
+        started = false
+        ended = true
+      }
+    }
+
+    beforeChar = code[index]
+    index++
+  }
+
+  return {
+    insideComment: !ended,
+    insideString: inString !== null,
+  }
+}
+
 export { resolve as resolvePath }

test/core/test/mocked.test.js

@@ -0,0 +1,35 @@
+import { assert, test, vi } from 'vitest'
+import { two } from '../src/submodule'
+import { timeout } from '../src/timeout'
+
+/*
+    vi.mock('../src/timeout', () => ({ timeout: 0 }))
+ /* */
+
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+const text = `
+  vi.mock('../src/timeout', () => ({ timeout: 0 }))
+`
+
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+const textComment = `
+  vi.mock('../src/timeout', () => ({ timeout: 0 }))
+  /**
+    vi.mock('../src/timeout', () => ({ timeout: 0 }))
+    vi.mock('../src/timeout', () => ({ timeout: 0 }))
+  */
+`
+
+vi.mock(
+  '../src/submodule',
+  () => ({
+    two: 55,
+  }),
+)
+
+// vi.mock('../src/submodule')
+
+test('vitest correctly passes multiline vi.mock syntax', () => {
+  assert.equal(55, two)
+  assert.equal(100, timeout)
+})