From b50cf7ad9c5062697bd47b8fcaca2fc3029fb79c Mon Sep 17 00:00:00 2001
From: Han <fenghan770@gmail.com>
Date: Mon, 18 Sep 2023 23:47:19 +0800
Subject: [PATCH] feat(config): add `diff` option (#4063)

Co-authored-by: Vladimir Sheremet <sleuths.slews0s@icloud.com>
---
 docs/config/index.md                          | 30 +++++++++++++++++++
 packages/browser/src/client/main.ts           |  4 +++
 packages/runner/src/run.ts                    | 13 ++++----
 packages/runner/src/types/runner.ts           |  2 ++
 packages/utils/src/error.ts                   |  6 ++--
 packages/vitest/src/browser.ts                |  2 +-
 packages/vitest/src/node/cli.ts               |  1 +
 packages/vitest/src/node/config.ts            |  7 +++++
 packages/vitest/src/runtime/entry.ts          |  1 +
 packages/vitest/src/runtime/runners/index.ts  |  3 ++
 packages/vitest/src/runtime/setup.common.ts   | 14 +++++++++
 packages/vitest/src/types/config.ts           |  5 ++++
 packages/vitest/src/types/index.ts            |  1 +
 test/browser/custom-diff-config.ts            |  4 +++
 test/browser/specs/runner.test.mjs            |  2 ++
 test/browser/vitest.config.mts                |  1 +
 .../fixtures/custom-diff-config.test.ts       |  5 ++++
 test/reporters/fixtures/custom-diff-config.ts |  4 +++
 .../reporters/fixtures/invalid-diff-config.ts |  1 +
 .../tests/custom-diff-config.spec.ts          | 23 ++++++++++++++
 20 files changed, 119 insertions(+), 10 deletions(-)
 create mode 100644 test/browser/custom-diff-config.ts
 create mode 100644 test/reporters/fixtures/custom-diff-config.test.ts
 create mode 100644 test/reporters/fixtures/custom-diff-config.ts
 create mode 100644 test/reporters/fixtures/invalid-diff-config.ts
 create mode 100644 test/reporters/tests/custom-diff-config.spec.ts

diff --git a/docs/config/index.md b/docs/config/index.md
index 035979eb28d6..0c7ab89bfb5e 100644
--- a/docs/config/index.md
+++ b/docs/config/index.md
@@ -1641,3 +1641,33 @@ export default defineConfig({
 })
 ```
 
+### diff
+
+- **Type:** `string`
+- **CLI:** `--diff=<value>`
+- **Version:** Since Vitest 0.34.5
+
+Path to a diff config that will be used to generate diff interface. Useful if you want to customize diff display.
+
+:::code-group
+```ts [vitest.diff.ts]
+import type { DiffOptions } from 'vitest'
+import c from 'picocolors'
+
+export default {
+  aIndicator: c.bold('--'),
+  bIndicator: c.bold('++'),
+  omitAnnotationLines: true,
+} satisfies DiffOptions
+```
+
+```ts [vitest.config.js]
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  test: {
+    diff: './vitest.diff.ts'
+  }
+})
+```
+:::
diff --git a/packages/browser/src/client/main.ts b/packages/browser/src/client/main.ts
index 47171f97c036..95abe1d22b3e 100644
--- a/packages/browser/src/client/main.ts
+++ b/packages/browser/src/client/main.ts
@@ -1,6 +1,7 @@
 import { createClient } from '@vitest/ws-client'
 import type { ResolvedConfig } from 'vitest'
 import type { CancelReason, VitestRunner } from '@vitest/runner'
+import type { VitestExecutor } from 'vitest/src/runtime/execute'
 import { createBrowserRunner } from './runner'
 import { importId } from './utils'
 import { setupConsoleLogSpy } from './logger'
@@ -101,6 +102,7 @@ async function runTests(paths: string[], config: ResolvedConfig) {
   const {
     startTests,
     setupCommonEnv,
+    loadDiffConfig,
     takeCoverageInsideWorker,
   } = await importId('vitest/browser') as typeof import('vitest/browser')
 
@@ -122,6 +124,8 @@ async function runTests(paths: string[], config: ResolvedConfig) {
     config.snapshotOptions.snapshotEnvironment = new BrowserSnapshotEnvironment()
 
   try {
+    runner.config.diffOptions = await loadDiffConfig(config, executor as VitestExecutor)
+
     await setupCommonEnv(config)
     const files = paths.map((path) => {
       return (`${config.root}/${path}`).replace(/\/+/g, '/')
diff --git a/packages/runner/src/run.ts b/packages/runner/src/run.ts
index 6765c9f95265..bd26172f11a8 100644
--- a/packages/runner/src/run.ts
+++ b/packages/runner/src/run.ts
@@ -1,6 +1,7 @@
 import limit from 'p-limit'
 import { getSafeTimers, shuffle } from '@vitest/utils'
 import { processError } from '@vitest/utils/error'
+import type { DiffOptions } from '@vitest/utils/diff'
 import type { VitestRunner } from './types/runner'
 import type { File, HookCleanupCallback, HookListener, SequenceHooks, Suite, SuiteHooks, Task, TaskMeta, TaskResult, TaskResultPack, TaskState, Test } from './types'
 import { partitionSuiteChildren } from './utils/suite'
@@ -173,7 +174,7 @@ export async function runTest(test: Test, runner: VitestRunner) {
         }
       }
       catch (e) {
-        failTask(test.result, e)
+        failTask(test.result, e, runner.config.diffOptions)
       }
 
       // skipped with new PendingError
@@ -189,7 +190,7 @@ export async function runTest(test: Test, runner: VitestRunner) {
         await callCleanupHooks(beforeEachCleanups)
       }
       catch (e) {
-        failTask(test.result, e)
+        failTask(test.result, e, runner.config.diffOptions)
       }
 
       if (test.result.state === 'pass')
@@ -233,7 +234,7 @@ export async function runTest(test: Test, runner: VitestRunner) {
   updateTask(test, runner)
 }
 
-function failTask(result: TaskResult, err: unknown) {
+function failTask(result: TaskResult, err: unknown, diffOptions?: DiffOptions) {
   if (err instanceof PendingError) {
     result.state = 'skip'
     return
@@ -244,7 +245,7 @@ function failTask(result: TaskResult, err: unknown) {
     ? err
     : [err]
   for (const e of errors) {
-    const error = processError(e)
+    const error = processError(e, diffOptions)
     result.error ??= error
     result.errors ??= []
     result.errors.push(error)
@@ -316,7 +317,7 @@ export async function runSuite(suite: Suite, runner: VitestRunner) {
       }
     }
     catch (e) {
-      failTask(suite.result, e)
+      failTask(suite.result, e, runner.config.diffOptions)
     }
 
     try {
@@ -324,7 +325,7 @@ export async function runSuite(suite: Suite, runner: VitestRunner) {
       await callCleanupHooks(beforeAllCleanups)
     }
     catch (e) {
-      failTask(suite.result, e)
+      failTask(suite.result, e, runner.config.diffOptions)
     }
 
     if (suite.mode === 'run') {
diff --git a/packages/runner/src/types/runner.ts b/packages/runner/src/types/runner.ts
index 25142152f9aa..f512731c959f 100644
--- a/packages/runner/src/types/runner.ts
+++ b/packages/runner/src/types/runner.ts
@@ -1,3 +1,4 @@
+import type { DiffOptions } from '@vitest/utils/diff'
 import type { File, SequenceHooks, SequenceSetupFiles, Suite, TaskResultPack, Test, TestContext } from './tasks'
 
 export interface VitestRunnerConfig {
@@ -21,6 +22,7 @@ export interface VitestRunnerConfig {
   testTimeout: number
   hookTimeout: number
   retry: number
+  diffOptions?: DiffOptions
 }
 
 export type VitestRunnerImportSource = 'collect' | 'setup'
diff --git a/packages/utils/src/error.ts b/packages/utils/src/error.ts
index c5403e54b26d..dc2aeb0d5aac 100644
--- a/packages/utils/src/error.ts
+++ b/packages/utils/src/error.ts
@@ -1,4 +1,4 @@
-import { diff } from './diff'
+import { type DiffOptions, diff } from './diff'
 import { format } from './display'
 import { deepClone, getOwnProperties, getType } from './helpers'
 import { stringify } from './stringify'
@@ -86,7 +86,7 @@ function normalizeErrorMessage(message: string) {
   return message.replace(/__vite_ssr_import_\d+__\./g, '')
 }
 
-export function processError(err: any) {
+export function processError(err: any, diffOptions?: DiffOptions) {
   if (!err || typeof err !== 'object')
     return { message: err }
   // stack is not serialized in worker communication
@@ -101,7 +101,7 @@ export function processError(err: any) {
     const clonedExpected = deepClone(err.expected, { forceWritable: true })
 
     const { replacedActual, replacedExpected } = replaceAsymmetricMatcher(clonedActual, clonedExpected)
-    err.diff = diff(replacedExpected, replacedActual)
+    err.diff = diff(replacedExpected, replacedActual, diffOptions)
   }
 
   if (typeof err.expected !== 'string')
diff --git a/packages/vitest/src/browser.ts b/packages/vitest/src/browser.ts
index 3171c38ac81b..83d98341f387 100644
--- a/packages/vitest/src/browser.ts
+++ b/packages/vitest/src/browser.ts
@@ -1,3 +1,3 @@
 export { startTests } from '@vitest/runner'
-export { setupCommonEnv } from './runtime/setup.common'
+export { setupCommonEnv, loadDiffConfig } from './runtime/setup.common'
 export { takeCoverageInsideWorker, stopCoverageInsideWorker, getCoverageProvider, startCoverageInsideWorker } from './integrations/coverage'
diff --git a/packages/vitest/src/node/cli.ts b/packages/vitest/src/node/cli.ts
index 25872408d6b7..287ec9a7d950 100644
--- a/packages/vitest/src/node/cli.ts
+++ b/packages/vitest/src/node/cli.ts
@@ -51,6 +51,7 @@ cli
   .option('--test-timeout <time>', 'Default timeout of a test in milliseconds (default: 5000)')
   .option('--bail <number>', 'Stop test execution when given number of tests have failed', { default: 0 })
   .option('--retry <times>', 'Retry the test specific number of times if it fails', { default: 0 })
+  .option('--diff <path>', 'Path to a diff config that will be used to generate diff interface')
   .help()
 
 cli
diff --git a/packages/vitest/src/node/config.ts b/packages/vitest/src/node/config.ts
index ea8219b86f83..c124d62ff682 100644
--- a/packages/vitest/src/node/config.ts
+++ b/packages/vitest/src/node/config.ts
@@ -289,6 +289,13 @@ export function resolveConfig(
     ...resolved.setupFiles,
   ]
 
+  if (resolved.diff) {
+    resolved.diff = normalize(
+      resolveModule(resolved.diff, { paths: [resolved.root] })
+        ?? resolve(resolved.root, resolved.diff))
+    resolved.forceRerunTriggers.push(resolved.diff)
+  }
+
   // the server has been created, we don't need to override vite.server options
   resolved.api = resolveApiServerConfig(options)
 
diff --git a/packages/vitest/src/runtime/entry.ts b/packages/vitest/src/runtime/entry.ts
index 0dffd3f2b648..bc53a708dc7d 100644
--- a/packages/vitest/src/runtime/entry.ts
+++ b/packages/vitest/src/runtime/entry.ts
@@ -20,6 +20,7 @@ export async function run(files: string[], config: ResolvedConfig, environment:
     setupChaiConfig(config.chaiConfig)
 
   const runner = await resolveTestRunner(config, executor)
+
   workerState.onCancel.then(reason => runner.onCancel?.(reason))
 
   workerState.durations.prepare = performance.now() - workerState.durations.prepare
diff --git a/packages/vitest/src/runtime/runners/index.ts b/packages/vitest/src/runtime/runners/index.ts
index 761c431147e7..fe22c048b443 100644
--- a/packages/vitest/src/runtime/runners/index.ts
+++ b/packages/vitest/src/runtime/runners/index.ts
@@ -6,6 +6,7 @@ import { distDir } from '../../paths'
 import { getWorkerState } from '../../utils/global'
 import { rpc } from '../rpc'
 import { takeCoverageInsideWorker } from '../../integrations/coverage'
+import { loadDiffConfig } from '../setup.common'
 
 const runnersFile = resolve(distDir, 'runners.js')
 
@@ -37,6 +38,8 @@ export async function resolveTestRunner(config: ResolvedConfig, executor: Vitest
   if (!testRunner.importFile)
     throw new Error('Runner must implement "importFile" method.')
 
+  testRunner.config.diffOptions = await loadDiffConfig(config, executor)
+
   // patch some methods, so custom runners don't need to call RPC
   const originalOnTaskUpdate = testRunner.onTaskUpdate
   testRunner.onTaskUpdate = async (task) => {
diff --git a/packages/vitest/src/runtime/setup.common.ts b/packages/vitest/src/runtime/setup.common.ts
index 32d8f5fc2096..672e830272e3 100644
--- a/packages/vitest/src/runtime/setup.common.ts
+++ b/packages/vitest/src/runtime/setup.common.ts
@@ -1,6 +1,8 @@
 import { setSafeTimers } from '@vitest/utils'
 import { resetRunOnceCounter } from '../integrations/run-once'
 import type { ResolvedConfig } from '../types'
+import type { DiffOptions } from '../types/matcher-utils'
+import type { VitestExecutor } from './execute'
 
 let globalSetup = false
 export async function setupCommonEnv(config: ResolvedConfig) {
@@ -21,3 +23,15 @@ function setupDefines(defines: Record<string, any>) {
   for (const key in defines)
     (globalThis as any)[key] = defines[key]
 }
+
+export async function loadDiffConfig(config: ResolvedConfig, executor: VitestExecutor) {
+  if (typeof config.diff !== 'string')
+    return
+
+  const diffModule = await executor.executeId(config.diff)
+
+  if (diffModule && typeof diffModule.default === 'object' && diffModule.default != null)
+    return diffModule.default as DiffOptions
+  else
+    throw new Error(`invalid diff config file ${config.diff}. Must have a default export with config object`)
+}
diff --git a/packages/vitest/src/types/config.ts b/packages/vitest/src/types/config.ts
index 6763cc938899..434e0542e873 100644
--- a/packages/vitest/src/types/config.ts
+++ b/packages/vitest/src/types/config.ts
@@ -539,6 +539,11 @@ export interface InlineConfig {
    */
   snapshotFormat?: PrettyFormatOptions
 
+  /**
+   * Path to a module which has a default export of diff config.
+   */
+  diff?: string
+
   /**
    * Resolve custom snapshot path
    */
diff --git a/packages/vitest/src/types/index.ts b/packages/vitest/src/types/index.ts
index e0d92a8ed0a3..a1cb1f5c5ec5 100644
--- a/packages/vitest/src/types/index.ts
+++ b/packages/vitest/src/types/index.ts
@@ -13,6 +13,7 @@ export type * from './worker'
 export type * from './general'
 export type * from './coverage'
 export type * from './benchmark'
+export type { DiffOptions } from '@vitest/utils/diff'
 export type {
   EnhancedSpy,
   MockedFunction,
diff --git a/test/browser/custom-diff-config.ts b/test/browser/custom-diff-config.ts
new file mode 100644
index 000000000000..358794c614fe
--- /dev/null
+++ b/test/browser/custom-diff-config.ts
@@ -0,0 +1,4 @@
+export default {
+  aAnnotation: 'Expected to be',
+  bAnnotation: 'But got',
+}
diff --git a/test/browser/specs/runner.test.mjs b/test/browser/specs/runner.test.mjs
index 713876a325c9..3c9619d31652 100644
--- a/test/browser/specs/runner.test.mjs
+++ b/test/browser/specs/runner.test.mjs
@@ -34,6 +34,8 @@ await test('tests are actually running', async () => {
 await test('correctly prints error', () => {
   assert.match(stderr, /expected 1 to be 2/, 'prints failing error')
   assert.match(stderr, /- 2\s+\+ 1/, 'prints failing diff')
+  assert.match(stderr, /Expected to be/, 'prints \`Expected to be\`')
+  assert.match(stderr, /But got/, 'prints \`But got\`')
 })
 
 await test('logs are redirected to stdout', async () => {
diff --git a/test/browser/vitest.config.mts b/test/browser/vitest.config.mts
index 7247ea6baed6..a39f44e4fa1a 100644
--- a/test/browser/vitest.config.mts
+++ b/test/browser/vitest.config.mts
@@ -20,6 +20,7 @@ export default defineConfig({
     },
     open: false,
     isolate: false,
+    diff: './custom-diff-config.ts',
     outputFile: './browser.json',
     reporters: ['json', {
       onInit: noop,
diff --git a/test/reporters/fixtures/custom-diff-config.test.ts b/test/reporters/fixtures/custom-diff-config.test.ts
new file mode 100644
index 000000000000..7f52fd58593e
--- /dev/null
+++ b/test/reporters/fixtures/custom-diff-config.test.ts
@@ -0,0 +1,5 @@
+import { expect, test } from 'vitest'
+
+test('', () => {
+  expect({ foo: 1 }).toMatchInlineSnapshot('xxx')
+})
diff --git a/test/reporters/fixtures/custom-diff-config.ts b/test/reporters/fixtures/custom-diff-config.ts
new file mode 100644
index 000000000000..358794c614fe
--- /dev/null
+++ b/test/reporters/fixtures/custom-diff-config.ts
@@ -0,0 +1,4 @@
+export default {
+  aAnnotation: 'Expected to be',
+  bAnnotation: 'But got',
+}
diff --git a/test/reporters/fixtures/invalid-diff-config.ts b/test/reporters/fixtures/invalid-diff-config.ts
new file mode 100644
index 000000000000..2065be389c54
--- /dev/null
+++ b/test/reporters/fixtures/invalid-diff-config.ts
@@ -0,0 +1 @@
+export const diffOptions = {}
diff --git a/test/reporters/tests/custom-diff-config.spec.ts b/test/reporters/tests/custom-diff-config.spec.ts
new file mode 100644
index 000000000000..3445913eb245
--- /dev/null
+++ b/test/reporters/tests/custom-diff-config.spec.ts
@@ -0,0 +1,23 @@
+import { expect, test } from 'vitest'
+import { resolve } from 'pathe'
+import { runVitest } from '../../test-utils'
+
+test('custom diff config', async () => {
+  const filename = resolve('./fixtures/custom-diff-config.test.ts')
+  const diff = resolve('./fixtures/custom-diff-config.ts')
+  const { stderr } = await runVitest({ root: './fixtures', diff }, [filename])
+
+  expect(stderr).toBeTruthy()
+  expect(stderr).toContain('Expected to be')
+  expect(stderr).toContain('But got')
+})
+
+test('invalid diff config file', async () => {
+  const filename = resolve('./fixtures/custom-diff-config.test.ts')
+  const diff = resolve('./fixtures/invalid-diff-config.ts')
+  const { stderr } = await runVitest({ root: './fixtures', diff }, [filename])
+
+  expect(stderr).toBeTruthy()
+  expect(stderr).toContain('invalid diff config file')
+  expect(stderr).toContain('Must have a default export with config object')
+})