feat: add assertion helper to hide internal stack traces

docs/api/vi.md

@@ -1335,3 +1335,42 @@ function resetConfig(): void
 ```
 
 If [`vi.setConfig`](#vi-setconfig) was called before, this will reset config to the original state.
+
+### vi.defineHelper <Version>4.1.0</Version> {#vi-defineHelper}
+
+```ts
+function defineHelper<F extends (...args: any) => any>(fn: F): F
+```
+
+Wraps a function to create an assertion helper. When an assertion fails inside the helper, the error stack trace will point to where the helper was called, not inside the helper itself. This makes it easier to identify the source of test failures when using custom assertion functions.
+
+Works with both synchronous and asynchronous functions, and supports `expect.soft()`.
+
+```ts
+import { expect, vi } from 'vitest'
+
+const assertPair = vi.defineHelper((a, b) => {
+  expect(a).toEqual(b)
+})
+
+test('example', () => {
+  assertPair('left', 'right') // Error points to this line
+})
+```
+
+Example output:
+
+<!-- eslint-skip -->
+```js
+FAIL  example.test.ts > example
+AssertionError: expected 'left' to deeply equal 'right'
+
+Expected: "right"
+Received: "left"
+
+ ❯ example.test.ts:8:3
+      7| test('example', () => {
+      8|   assertPair('left', 'right')
+       |   ^
+      9| })
+```

packages/utils/src/source-map.ts

@@ -228,10 +228,16 @@ export function parseStacktrace(
   options: StackTraceParserOptions = {},
 ): ParsedStack[] {
   const { ignoreStackEntries = stackIgnorePatterns } = options
-  const stacks = !CHROME_IE_STACK_REGEXP.test(stack)
+  let stacks = !CHROME_IE_STACK_REGEXP.test(stack)
     ? parseFFOrSafariStackTrace(stack)
     : parseV8Stacktrace(stack)
 
+  // remove assertion helper's internal stacks
+  const helperIndex = stacks.findLastIndex(s => s.method === '__VITEST_HELPER__' || s.method === 'async*__VITEST_HELPER__')
+  if (helperIndex >= 0) {
+    stacks = stacks.slice(helperIndex + 1)
+  }
+
   return stacks.map((stack) => {
     if (options.getUrlId) {
       stack.file = options.getUrlId(stack.file)

packages/vitest/src/integrations/vi.ts

@@ -168,6 +168,40 @@ export interface VitestUtils {
    */
   waitFor: typeof waitFor
 
+  /**
+   * Wraps a function to create an assertion helper. When an assertion fails inside the helper,
+   * the error stack trace will point to where the helper was called, not inside the helper itself.
+   * Works with both synchronous and asynchronous functions, and supports `expect.soft()`.
+   *
+   * @example
+   * ```ts
+   * const myEqual = vi.defineHelper((x, y) => {
+   *   expect(x).toEqual(y)
+   * })
+   *
+   * test('example', () => {
+   *   myEqual('left', 'right') // Error points to this line
+   * })
+   * ```
+   * Example output:
+   * ```
+   * FAIL  example.test.ts > example
+   * AssertionError: expected 'left' to deeply equal 'right'
+   *
+   * Expected: "right"
+   * Received: "left"
+   *
+   *  ❯ example.test.ts:6:3
+   *       4| test('example', () => {
+   *       5|   myEqual('left', 'right')
+   *        |   ^
+   *       6| })
+   * ```
+   * @param fn The assertion function to wrap
+   * @returns A wrapped function with the same signature
+   */
+  defineHelper: <F extends (...args: any) => any>(fn: F) => F
+
   /**
    * This is similar to [`vi.waitFor`](https://vitest.dev/api/vi#vi-waitfor), but if the callback throws any errors, execution is immediately interrupted and an error message is received.
    *
@@ -576,6 +610,17 @@ function createVitest(): VitestUtils {
     fn,
     waitFor,
     waitUntil,
+    defineHelper: (fn) => {
+      return function __VITEST_HELPER__(...args: any[]): any {
+        const result = fn(...args)
+        if (result && typeof result === 'object' && typeof result.then === 'function') {
+          return (async function __VITEST_HELPER__() {
+            return await result
+          })()
+        }
+        return result
+      } as any
+    },
     hoisted<T>(factory: () => T): T {
       assertTypes(factory, '"vi.hoisted" factory', ['function'])
       return factory()

test/browser/fixtures/assertion-helper/basic.test.ts

@@ -0,0 +1,35 @@
+import { expect, test, vi } from 'vitest'
+
+const myEqual = vi.defineHelper((a: any, b: any) => {
+  expect(a).toEqual(b)
+})
+
+const myEqualAsync = vi.defineHelper(async (a: any, b: any) => {
+  await new Promise(r => setTimeout(r, 1))
+  expect(a).toEqual(b)
+})
+
+const myEqualSoft = vi.defineHelper((a: any, b: any) => {
+  expect.soft(a).toEqual(b)
+})
+
+const myEqualSoftAsync = vi.defineHelper(async (a: any, b: any) => {
+  await new Promise(r => setTimeout(r, 1))
+  expect.soft(a).toEqual(b)
+})
+
+test('sync', () => {
+  myEqual('sync', 'x')
+})
+
+test('async', async () => {
+  await myEqualAsync('async', 'x')
+})
+
+test('soft', () => {
+  myEqualSoft('soft', 'x')
+})
+
+test('soft async', async () => {
+  await myEqualSoftAsync('soft async', 'x')
+})

test/browser/fixtures/assertion-helper/vitest.config.ts

@@ -0,0 +1,14 @@
+import { fileURLToPath } from 'node:url'
+import { defineConfig } from 'vitest/config'
+import { instances, provider } from '../../settings'
+
+export default defineConfig({
+  cacheDir: fileURLToPath(new URL('./node_modules/.vite', import.meta.url)),
+  test: {
+    browser: {
+      enabled: true,
+      provider,
+      instances,
+    },
+  },
+})

test/browser/specs/assertion-helper.test.ts

@@ -0,0 +1,208 @@
+import path from 'node:path'
+import { expect, test } from 'vitest'
+import { buildTestProjectTree } from '../../test-utils'
+import { instances, runBrowserTests } from './utils'
+
+test('vi.defineHelper hides internal stack traces', async () => {
+  const { results, ctx } = await runBrowserTests({
+    root: './fixtures/assertion-helper',
+  })
+
+  const projectTree = buildTestProjectTree(results, (testCase) => {
+    const result = testCase.result()
+    return result.errors.map((e) => {
+      const stacks = e.stacks.map(s => ({
+        ...s,
+        file: path.relative(ctx.config.root, s.file),
+      }))
+      return ({ message: e.message, stacks })
+    })
+  })
+  expect(Object.keys(projectTree).sort()).toEqual(instances.map(i => i.browser).sort())
+
+  for (const [name, tree] of Object.entries(projectTree)) {
+    if (name === 'firefox') {
+      expect.soft(tree).toMatchInlineSnapshot(`
+        {
+          "basic.test.ts": {
+            "async": [
+              {
+                "message": "expected 'async' to deeply equal 'x'",
+                "stacks": [
+                  {
+                    "column": 8,
+                    "file": "basic.test.ts",
+                    "line": 26,
+                    "method": "",
+                  },
+                ],
+              },
+            ],
+            "soft": [
+              {
+                "message": "expected 'soft' to deeply equal 'x'",
+                "stacks": [
+                  {
+                    "column": 14,
+                    "file": "basic.test.ts",
+                    "line": 30,
+                    "method": "",
+                  },
+                ],
+              },
+            ],
+            "soft async": [
+              {
+                "message": "expected 'soft async' to deeply equal 'x'",
+                "stacks": [
+                  {
+                    "column": 8,
+                    "file": "basic.test.ts",
+                    "line": 34,
+                    "method": "",
+                  },
+                ],
+              },
+            ],
+            "sync": [
+              {
+                "message": "expected 'sync' to deeply equal 'x'",
+                "stacks": [
+                  {
+                    "column": 10,
+                    "file": "basic.test.ts",
+                    "line": 22,
+                    "method": "",
+                  },
+                ],
+              },
+            ],
+          },
+        }
+      `)
+    }
+    else if (name === 'webkit') {
+      // async stack trace is incomplete on webkit
+      // waiting for https://github.com/WebKit/WebKit/pull/57832 to land on playwright
+      // bun has already landed https://github.com/oven-sh/bun/pull/22517
+      expect.soft(tree).toMatchInlineSnapshot(`
+        {
+          "basic.test.ts": {
+            "async": [
+              {
+                "message": "expected 'async' to deeply equal 'x'",
+                "stacks": [
+                  {
+                    "column": 20,
+                    "file": "basic.test.ts",
+                    "line": 9,
+                    "method": "",
+                  },
+                ],
+              },
+            ],
+            "soft": [
+              {
+                "message": "expected 'soft' to deeply equal 'x'",
+                "stacks": [
+                  {
+                    "column": 14,
+                    "file": "basic.test.ts",
+                    "line": 30,
+                    "method": "",
+                  },
+                ],
+              },
+            ],
+            "soft async": [
+              {
+                "message": "expected 'soft async' to deeply equal 'x'",
+                "stacks": [
+                  {
+                    "column": 25,
+                    "file": "basic.test.ts",
+                    "line": 18,
+                    "method": "",
+                  },
+                ],
+              },
+            ],
+            "sync": [
+              {
+                "message": "expected 'sync' to deeply equal 'x'",
+                "stacks": [
+                  {
+                    "column": 10,
+                    "file": "basic.test.ts",
+                    "line": 22,
+                    "method": "",
+                  },
+                ],
+              },
+            ],
+          },
+        }
+      `)
+    }
+    else {
+      expect.soft(tree).toMatchInlineSnapshot(`
+        {
+          "basic.test.ts": {
+            "async": [
+              {
+                "message": "expected 'async' to deeply equal 'x'",
+                "stacks": [
+                  {
+                    "column": 2,
+                    "file": "basic.test.ts",
+                    "line": 26,
+                    "method": "",
+                  },
+                ],
+              },
+            ],
+            "soft": [
+              {
+                "message": "expected 'soft' to deeply equal 'x'",
+                "stacks": [
+                  {
+                    "column": 2,
+                    "file": "basic.test.ts",
+                    "line": 30,
+                    "method": "",
+                  },
+                ],
+              },
+            ],
+            "soft async": [
+              {
+                "message": "expected 'soft async' to deeply equal 'x'",
+                "stacks": [
+                  {
+                    "column": 2,
+                    "file": "basic.test.ts",
+                    "line": 34,
+                    "method": "",
+                  },
+                ],
+              },
+            ],
+            "sync": [
+              {
+                "message": "expected 'sync' to deeply equal 'x'",
+                "stacks": [
+                  {
+                    "column": 2,
+                    "file": "basic.test.ts",
+                    "line": 22,
+                    "method": "",
+                  },
+                ],
+              },
+            ],
+          },
+        }
+      `)
+    }
+  }
+})