From 30f728bc1a6c4cf794b184ff3ed0ee27888b924c Mon Sep 17 00:00:00 2001
From: Vladimir <sleuths.slews0s@icloud.com>
Date: Fri, 3 May 2024 10:21:03 +0200
Subject: [PATCH] feat: custom "snapshotEnvironment" option (#5449)

---
 docs/config/index.md                          | 28 ++++++++++++
 packages/browser/src/client/runner.ts         |  9 ++--
 packages/browser/src/client/snapshot.ts       | 11 +++--
 packages/vitest/package.json                  |  4 ++
 packages/vitest/rollup.config.js              |  3 ++
 packages/vitest/snapshot.d.ts                 |  1 +
 packages/vitest/src/index.ts                  |  1 +
 .../snapshot/environments/node.ts             | 11 ++---
 .../resolveSnapshotEnvironment.ts             | 18 ++++++++
 packages/vitest/src/node/cli/cli-config.ts    |  1 +
 packages/vitest/src/node/config.ts            |  3 ++
 packages/vitest/src/runtime/runBaseTests.ts   |  2 +-
 packages/vitest/src/runtime/runVmTests.ts     | 11 +++--
 packages/vitest/src/runtime/setup-node.ts     |  7 +--
 packages/vitest/src/snapshot.ts               |  4 ++
 packages/vitest/src/types/config.ts           |  5 +++
 test/browser/custom-snapshot-env.ts           |  1 +
 test/browser/vitest.config.mts                |  2 +
 .../snapshots/test/custom-environment.test.ts | 33 ++++++++++++++
 .../snapshot-environment.ts                   | 45 +++++++++++++++++++
 .../test/snapshots.test.ts                    |  9 ++++
 .../vitest.config.ts                          |  7 +++
 22 files changed, 195 insertions(+), 21 deletions(-)
 create mode 100644 packages/vitest/snapshot.d.ts
 create mode 100644 packages/vitest/src/integrations/snapshot/environments/resolveSnapshotEnvironment.ts
 create mode 100644 packages/vitest/src/snapshot.ts
 create mode 100644 test/browser/custom-snapshot-env.ts
 create mode 100644 test/snapshots/test/custom-environment.test.ts
 create mode 100644 test/snapshots/test/fixtures/custom-snapshot-environment/snapshot-environment.ts
 create mode 100644 test/snapshots/test/fixtures/custom-snapshot-environment/test/snapshots.test.ts
 create mode 100644 test/snapshots/test/fixtures/custom-snapshot-environment/vitest.config.ts

diff --git a/docs/config/index.md b/docs/config/index.md
index 45fe17789f58..3a778920b3cd 100644
--- a/docs/config/index.md
+++ b/docs/config/index.md
@@ -2200,3 +2200,31 @@ The `location` property has `column` and `line` values that correspond to the `t
 ::: tip
 This option has no effect if you do not use custom code that relies on this.
 :::
+
+### snapshotEnvironment <Version>1.6.0</Version> {#snapshotEnvironment}
+
+- **Type:** `string`
+
+Path to a custom snapshot environment implementation. This is useful if you are running your tests in an environment that doesn't support Node.js APIs. This option doesn't have any effect on a browser runner.
+
+This object should have the shape of `SnapshotEnvironment` and is used to resolve and read/write snapshot files:
+
+```ts
+export interface SnapshotEnvironment {
+  getVersion: () => string
+  getHeader: () => string
+  resolvePath: (filepath: string) => Promise<string>
+  resolveRawPath: (testPath: string, rawPath: string) => Promise<string>
+  saveSnapshotFile: (filepath: string, snapshot: string) => Promise<void>
+  readSnapshotFile: (filepath: string) => Promise<string | null>
+  removeSnapshotFile: (filepath: string) => Promise<void>
+}
+```
+
+You can extend default `VitestSnapshotEnvironment` from `vitest/snapshot` entry point if you need to overwrite only a part of the API.
+
+::: warning
+This is a low-level option and should be used only for advanced cases where you don't have access to default Node.js APIs.
+
+If you just need to configure snapshots feature, use [`snapshotFormat`](#snapshotformat) or [`resolveSnapshotPath`](#resolvesnapshotpath) options.
+:::
diff --git a/packages/browser/src/client/runner.ts b/packages/browser/src/client/runner.ts
index 2f762ff61159..e25921ad7193 100644
--- a/packages/browser/src/client/runner.ts
+++ b/packages/browser/src/client/runner.ts
@@ -3,7 +3,7 @@ import type { ResolvedConfig } from 'vitest'
 import type { VitestExecutor } from 'vitest/execute'
 import { rpc } from './rpc'
 import { getConfig, importId } from './utils'
-import { BrowserSnapshotEnvironment } from './snapshot'
+import { VitestBrowserSnapshotEnvironment } from './snapshot'
 
 interface BrowserRunnerOptions {
   config: ResolvedConfig
@@ -92,7 +92,10 @@ export async function initiateRunner() {
   if (cachedRunner)
     return cachedRunner
   const config = getConfig()
-  const [{ VitestTestRunner, NodeBenchmarkRunner }, { takeCoverageInsideWorker, loadDiffConfig, loadSnapshotSerializers }] = await Promise.all([
+  const [
+    { VitestTestRunner, NodeBenchmarkRunner },
+    { takeCoverageInsideWorker, loadDiffConfig, loadSnapshotSerializers },
+  ] = await Promise.all([
     importId('vitest/runners') as Promise<typeof import('vitest/runners')>,
     importId('vitest/browser') as Promise<typeof import('vitest/browser')>,
   ])
@@ -101,7 +104,7 @@ export async function initiateRunner() {
     takeCoverage: () => takeCoverageInsideWorker(config.coverage, { executeId: importId }),
   })
   if (!config.snapshotOptions.snapshotEnvironment)
-    config.snapshotOptions.snapshotEnvironment = new BrowserSnapshotEnvironment()
+    config.snapshotOptions.snapshotEnvironment = new VitestBrowserSnapshotEnvironment()
   const runner = new BrowserRunner({
     config,
   })
diff --git a/packages/browser/src/client/snapshot.ts b/packages/browser/src/client/snapshot.ts
index adf9180c20e2..a441423b81a7 100644
--- a/packages/browser/src/client/snapshot.ts
+++ b/packages/browser/src/client/snapshot.ts
@@ -1,7 +1,7 @@
-import type { SnapshotEnvironment } from 'vitest'
-import { rpc } from './rpc'
+import type { VitestClient } from '@vitest/ws-client'
+import type { SnapshotEnvironment } from 'vitest/snapshot'
 
-export class BrowserSnapshotEnvironment implements SnapshotEnvironment {
+export class VitestBrowserSnapshotEnvironment implements SnapshotEnvironment {
   getVersion(): string {
     return '1'
   }
@@ -30,3 +30,8 @@ export class BrowserSnapshotEnvironment implements SnapshotEnvironment {
     return rpc().removeSnapshotFile(filepath)
   }
 }
+
+function rpc(): VitestClient['rpc'] {
+  // @ts-expect-error not typed global
+  return globalThis.__vitest_worker__.rpc
+}
diff --git a/packages/vitest/package.json b/packages/vitest/package.json
index 9c22a758b82f..5ae171420400 100644
--- a/packages/vitest/package.json
+++ b/packages/vitest/package.json
@@ -91,6 +91,10 @@
     "./reporters": {
       "types": "./dist/reporters.d.ts",
       "default": "./dist/reporters.js"
+    },
+    "./snapshot": {
+      "types": "./dist/snapshot.d.ts",
+      "default": "./dist/snapshot.js"
     }
   },
   "main": "./dist/index.js",
diff --git a/packages/vitest/rollup.config.js b/packages/vitest/rollup.config.js
index 8030ae847fd8..06cd3da01cec 100644
--- a/packages/vitest/rollup.config.js
+++ b/packages/vitest/rollup.config.js
@@ -41,6 +41,8 @@ const entries = {
   'workers/vmForks': 'src/runtime/workers/vmForks.ts',
 
   'workers/runVmTests': 'src/runtime/runVmTests.ts',
+
+  'snapshot': 'src/snapshot.ts',
 }
 
 const dtsEntries = {
@@ -56,6 +58,7 @@ const dtsEntries = {
   execute: 'src/public/execute.ts',
   reporters: 'src/public/reporters.ts',
   workers: 'src/workers.ts',
+  snapshot: 'src/snapshot.ts',
 }
 
 const external = [
diff --git a/packages/vitest/snapshot.d.ts b/packages/vitest/snapshot.d.ts
new file mode 100644
index 000000000000..e032fa789764
--- /dev/null
+++ b/packages/vitest/snapshot.d.ts
@@ -0,0 +1 @@
+export * from './dist/snapshot.js'
diff --git a/packages/vitest/src/index.ts b/packages/vitest/src/index.ts
index aab37704962d..4b298149dc56 100644
--- a/packages/vitest/src/index.ts
+++ b/packages/vitest/src/index.ts
@@ -17,6 +17,7 @@ export * from './integrations/chai'
 export * from './integrations/vi'
 export * from './integrations/utils'
 export { inject } from './integrations/inject'
+// TODO: remove in 2.0.0, import from vitest/snapshot directly
 export type { SnapshotEnvironment } from '@vitest/snapshot/environment'
 
 export * from './types'
diff --git a/packages/vitest/src/integrations/snapshot/environments/node.ts b/packages/vitest/src/integrations/snapshot/environments/node.ts
index 739501a67f29..938a85836ac2 100644
--- a/packages/vitest/src/integrations/snapshot/environments/node.ts
+++ b/packages/vitest/src/integrations/snapshot/environments/node.ts
@@ -1,16 +1,13 @@
 import { NodeSnapshotEnvironment } from '@vitest/snapshot/environment'
-import type { WorkerRPC } from '../../../types/worker'
-
-export class VitestSnapshotEnvironment extends NodeSnapshotEnvironment {
-  constructor(private rpc: WorkerRPC) {
-    super()
-  }
+import { getWorkerState } from '../../../utils'
 
+export class VitestNodeSnapshotEnvironment extends NodeSnapshotEnvironment {
   getHeader(): string {
     return `// Vitest Snapshot v${this.getVersion()}, https://vitest.dev/guide/snapshot.html`
   }
 
   resolvePath(filepath: string): Promise<string> {
-    return this.rpc.resolveSnapshotPath(filepath)
+    const rpc = getWorkerState().rpc
+    return rpc.resolveSnapshotPath(filepath)
   }
 }
diff --git a/packages/vitest/src/integrations/snapshot/environments/resolveSnapshotEnvironment.ts b/packages/vitest/src/integrations/snapshot/environments/resolveSnapshotEnvironment.ts
new file mode 100644
index 000000000000..8dce8c769013
--- /dev/null
+++ b/packages/vitest/src/integrations/snapshot/environments/resolveSnapshotEnvironment.ts
@@ -0,0 +1,18 @@
+import type { SnapshotEnvironment } from '@vitest/snapshot/environment'
+import type { VitestExecutor } from '../../../runtime/execute'
+import type { ResolvedConfig } from '../../../types'
+
+export async function resolveSnapshotEnvironment(
+  config: ResolvedConfig,
+  executor: VitestExecutor,
+): Promise<SnapshotEnvironment> {
+  if (!config.snapshotEnvironment) {
+    const { VitestNodeSnapshotEnvironment } = await import('./node')
+    return new VitestNodeSnapshotEnvironment()
+  }
+
+  const mod = await executor.executeId(config.snapshotEnvironment)
+  if (typeof mod.default !== 'object' || !mod.default)
+    throw new Error('Snapshot environment module must have a default export object with a shape of `SnapshotEnvironment`')
+  return mod.default
+}
diff --git a/packages/vitest/src/node/cli/cli-config.ts b/packages/vitest/src/node/cli/cli-config.ts
index 459db90a5472..0ae8b529b75d 100644
--- a/packages/vitest/src/node/cli/cli-config.ts
+++ b/packages/vitest/src/node/cli/cli-config.ts
@@ -637,4 +637,5 @@ export const cliOptionsConfig: VitestCLIOptions = {
   deps: null,
   name: null,
   includeTaskLocation: null,
+  snapshotEnvironment: null,
 }
diff --git a/packages/vitest/src/node/config.ts b/packages/vitest/src/node/config.ts
index 37b13f416d6f..66bbdcb6e899 100644
--- a/packages/vitest/src/node/config.ts
+++ b/packages/vitest/src/node/config.ts
@@ -251,6 +251,9 @@ export function resolveConfig(
   if (resolved.runner)
     resolved.runner = resolvePath(resolved.runner, resolved.root)
 
+  if (resolved.snapshotEnvironment)
+    resolved.snapshotEnvironment = resolvePath(resolved.snapshotEnvironment, resolved.root)
+
   resolved.testNamePattern = resolved.testNamePattern
     ? resolved.testNamePattern instanceof RegExp
       ? resolved.testNamePattern
diff --git a/packages/vitest/src/runtime/runBaseTests.ts b/packages/vitest/src/runtime/runBaseTests.ts
index d9775aa308b0..9d8ae41b2414 100644
--- a/packages/vitest/src/runtime/runBaseTests.ts
+++ b/packages/vitest/src/runtime/runBaseTests.ts
@@ -14,7 +14,7 @@ import { closeInspector } from './inspector'
 export async function run(files: string[], config: ResolvedConfig, environment: ResolvedTestEnvironment, executor: VitestExecutor): Promise<void> {
   const workerState = getWorkerState()
 
-  await setupGlobalEnv(config, environment)
+  await setupGlobalEnv(config, environment, executor)
   await startCoverageInsideWorker(config.coverage, executor)
 
   if (config.chaiConfig)
diff --git a/packages/vitest/src/runtime/runVmTests.ts b/packages/vitest/src/runtime/runVmTests.ts
index 3121b7e6ffa3..34dcb3361c42 100644
--- a/packages/vitest/src/runtime/runVmTests.ts
+++ b/packages/vitest/src/runtime/runVmTests.ts
@@ -10,8 +10,8 @@ import { setupChaiConfig } from '../integrations/chai/config'
 import { startCoverageInsideWorker, stopCoverageInsideWorker } from '../integrations/coverage'
 import type { ResolvedConfig } from '../types'
 import { getWorkerState } from '../utils/global'
-import { VitestSnapshotEnvironment } from '../integrations/snapshot/environments/node'
 import * as VitestIndex from '../index'
+import { resolveSnapshotEnvironment } from '../integrations/snapshot/environments/resolveSnapshotEnvironment'
 import type { VitestExecutor } from './execute'
 import { resolveTestRunner } from './runners'
 import { setupCommonEnv } from './setup-common'
@@ -27,8 +27,6 @@ export async function run(files: string[], config: ResolvedConfig, executor: Vit
     enumerable: false,
   })
 
-  config.snapshotOptions.snapshotEnvironment = new VitestSnapshotEnvironment(workerState.rpc)
-
   setupColors(createColors(isatty(1)))
 
   if (workerState.environment.transformMode === 'web') {
@@ -55,7 +53,12 @@ export async function run(files: string[], config: ResolvedConfig, executor: Vit
   if (config.chaiConfig)
     setupChaiConfig(config.chaiConfig)
 
-  const runner = await resolveTestRunner(config, executor)
+  const [runner, snapshotEnvironment] = await Promise.all([
+    resolveTestRunner(config, executor),
+    resolveSnapshotEnvironment(config, executor),
+  ])
+
+  config.snapshotOptions.snapshotEnvironment = snapshotEnvironment
 
   workerState.onCancel.then((reason) => {
     closeInspector(config)
diff --git a/packages/vitest/src/runtime/setup-node.ts b/packages/vitest/src/runtime/setup-node.ts
index 809eaccfe70a..0c52f5b8e5d9 100644
--- a/packages/vitest/src/runtime/setup-node.ts
+++ b/packages/vitest/src/runtime/setup-node.ts
@@ -5,15 +5,16 @@ import { isatty } from 'node:tty'
 import { installSourcemapsSupport } from 'vite-node/source-map'
 import { createColors, setupColors } from '@vitest/utils'
 import type { EnvironmentOptions, ResolvedConfig, ResolvedTestEnvironment } from '../types'
-import { VitestSnapshotEnvironment } from '../integrations/snapshot/environments/node'
 import { getSafeTimers, getWorkerState } from '../utils'
 import * as VitestIndex from '../index'
 import { expect } from '../integrations/chai'
+import { resolveSnapshotEnvironment } from '../integrations/snapshot/environments/resolveSnapshotEnvironment'
 import { setupCommonEnv } from './setup-common'
+import type { VitestExecutor } from './execute'
 
 // this should only be used in Node
 let globalSetup = false
-export async function setupGlobalEnv(config: ResolvedConfig, { environment }: ResolvedTestEnvironment) {
+export async function setupGlobalEnv(config: ResolvedConfig, { environment }: ResolvedTestEnvironment, executor: VitestExecutor) {
   await setupCommonEnv(config)
 
   Object.defineProperty(globalThis, '__vitest_index__', {
@@ -24,7 +25,7 @@ export async function setupGlobalEnv(config: ResolvedConfig, { environment }: Re
   const state = getWorkerState()
 
   if (!state.config.snapshotOptions.snapshotEnvironment)
-    state.config.snapshotOptions.snapshotEnvironment = new VitestSnapshotEnvironment(state.rpc)
+    state.config.snapshotOptions.snapshotEnvironment = await resolveSnapshotEnvironment(config, executor)
 
   if (globalSetup)
     return
diff --git a/packages/vitest/src/snapshot.ts b/packages/vitest/src/snapshot.ts
new file mode 100644
index 000000000000..12ca4245fee7
--- /dev/null
+++ b/packages/vitest/src/snapshot.ts
@@ -0,0 +1,4 @@
+export type { SnapshotEnvironment } from '@vitest/snapshot/environment'
+export {
+  VitestNodeSnapshotEnvironment as VitestSnapshotEnvironment,
+} from './integrations/snapshot/environments/node'
diff --git a/packages/vitest/src/types/config.ts b/packages/vitest/src/types/config.ts
index efbd238904bd..a45546afa43c 100644
--- a/packages/vitest/src/types/config.ts
+++ b/packages/vitest/src/types/config.ts
@@ -560,6 +560,11 @@ export interface InlineConfig {
    */
   resolveSnapshotPath?: (path: string, extension: string) => string
 
+  /**
+   * Path to a custom snapshot environment module that has a defualt export of `SnapshotEnvironment` object.
+   */
+  snapshotEnvironment?: string
+
   /**
    * Pass with no tests
    */
diff --git a/test/browser/custom-snapshot-env.ts b/test/browser/custom-snapshot-env.ts
new file mode 100644
index 000000000000..2265376b05c1
--- /dev/null
+++ b/test/browser/custom-snapshot-env.ts
@@ -0,0 +1 @@
+throw new Error('This file should not be executed')
diff --git a/test/browser/vitest.config.mts b/test/browser/vitest.config.mts
index 3056fad93a08..97bd6e9180e2 100644
--- a/test/browser/vitest.config.mts
+++ b/test/browser/vitest.config.mts
@@ -20,6 +20,8 @@ export default defineConfig({
   },
   test: {
     include: ['test/**.test.{ts,js}'],
+    // having a snapshot environment doesn't affect browser tests
+    snapshotEnvironment: './custom-snapshot-env.ts',
     browser: {
       enabled: true,
       name: browser,
diff --git a/test/snapshots/test/custom-environment.test.ts b/test/snapshots/test/custom-environment.test.ts
new file mode 100644
index 000000000000..215de0b1bc08
--- /dev/null
+++ b/test/snapshots/test/custom-environment.test.ts
@@ -0,0 +1,33 @@
+import { readFileSync, rmSync, writeFileSync } from 'node:fs'
+import { afterEach, expect, test } from 'vitest'
+import { dirname, resolve } from 'pathe'
+import { runVitest } from '../../test-utils'
+
+const testFileName = resolve(import.meta.dirname, './fixtures/custom-snapshot-environment/test/snapshots.test.ts')
+const snapshotFile = resolve(dirname(testFileName), './__snapshots__/snapshots.test.ts.snap')
+const testFile = readFileSync(testFileName, 'utf-8')
+
+afterEach(() => {
+  writeFileSync(testFileName, testFile)
+  rmSync(snapshotFile)
+})
+
+test('custom environment resolved correctly', async () => {
+  const { stdout, stderr } = await runVitest({
+    root: 'test/fixtures/custom-snapshot-environment',
+    update: true,
+  })
+
+  const snapshotLogs = stdout.split('\n').filter(i => i.startsWith('## ')).join('\n')
+  expect(stderr).toBe('')
+  expect(snapshotLogs).toMatchInlineSnapshot(`
+    "## resolvePath test/fixtures/custom-snapshot-environment/test/snapshots.test.ts
+    ## readSnapshotFile test/fixtures/custom-snapshot-environment/test/__snapshots__/snapshots.test.ts.snap
+    ## getHeader
+    ## getVersion
+    ## readSnapshotFile test/fixtures/custom-snapshot-environment/test/__snapshots__/snapshots.test.ts.snap
+    ## saveSnapshotFile test/fixtures/custom-snapshot-environment/test/__snapshots__/snapshots.test.ts.snap
+    ## readSnapshotFile test/fixtures/custom-snapshot-environment/test/snapshots.test.ts
+    ## saveSnapshotFile test/fixtures/custom-snapshot-environment/test/snapshots.test.ts"
+  `)
+})
diff --git a/test/snapshots/test/fixtures/custom-snapshot-environment/snapshot-environment.ts b/test/snapshots/test/fixtures/custom-snapshot-environment/snapshot-environment.ts
new file mode 100644
index 000000000000..1675ee9395cf
--- /dev/null
+++ b/test/snapshots/test/fixtures/custom-snapshot-environment/snapshot-environment.ts
@@ -0,0 +1,45 @@
+import { relative as _relative } from 'pathe'
+import { VitestSnapshotEnvironment } from 'vitest/snapshot'
+
+function relative(file: string) {
+  return _relative(process.cwd(), file)
+}
+
+class CustomSnapshotEnvironment extends VitestSnapshotEnvironment {
+  getVersion(): string {
+    console.log('## getVersion')
+    return super.getVersion()
+  }
+
+  getHeader() {
+    console.log('## getHeader')
+    return super.getHeader()
+  }
+
+  resolvePath(filepath: string) {
+    console.log('## resolvePath', relative(filepath))
+    return super.resolvePath(filepath)
+  }
+
+  resolveRawPath(testPath: string, rawPath: string) {
+    console.log('## resolveRawPath', relative(testPath), relative(rawPath))
+    return super.resolveRawPath(testPath, rawPath)
+  }
+
+  saveSnapshotFile(filepath: string, snapshot: string) {
+    console.log('## saveSnapshotFile', relative(filepath))
+    return super.saveSnapshotFile(filepath, snapshot)
+  }
+
+  readSnapshotFile(filepath: string) {
+    console.log('## readSnapshotFile', relative(filepath))
+    return super.readSnapshotFile(filepath)
+  }
+
+  removeSnapshotFile(filepath: string) {
+    console.log('## removeSnapshotFile', relative(filepath))
+    return super.removeSnapshotFile(filepath)
+  }
+}
+
+export default new CustomSnapshotEnvironment()
diff --git a/test/snapshots/test/fixtures/custom-snapshot-environment/test/snapshots.test.ts b/test/snapshots/test/fixtures/custom-snapshot-environment/test/snapshots.test.ts
new file mode 100644
index 000000000000..d6198ffce8ef
--- /dev/null
+++ b/test/snapshots/test/fixtures/custom-snapshot-environment/test/snapshots.test.ts
@@ -0,0 +1,9 @@
+import {test, expect} from 'vitest'
+
+test('regular snapshot', () => {
+  expect({ a: 1 }).toMatchSnapshot()
+})
+
+test('inline snapshot', () => {
+  expect({ a: 1 }).toMatchInlineSnapshot()
+})
diff --git a/test/snapshots/test/fixtures/custom-snapshot-environment/vitest.config.ts b/test/snapshots/test/fixtures/custom-snapshot-environment/vitest.config.ts
new file mode 100644
index 000000000000..9c804c75d2ef
--- /dev/null
+++ b/test/snapshots/test/fixtures/custom-snapshot-environment/vitest.config.ts
@@ -0,0 +1,7 @@
+import { defineConfig } from 'vitest/config';
+
+export default defineConfig({
+  test: {
+    snapshotEnvironment: './snapshot-environment.ts'
+  }
+})
\ No newline at end of file