test: onTestFinished instead of afterEach for file restoration (#221)

Author: sapphi-red

packages/vite/src/node/ssr/runtime/__tests__/server-source-maps.spec.ts

@@ -1,6 +1,10 @@
 import { describe, expect } from 'vitest'
 import type { ViteDevServer } from '../../..'
-import { createModuleRunnerTester, editFile, resolvePath } from './utils'
+import {
+  createFixtureEditor,
+  createModuleRunnerTester,
+  resolvePath,
+} from './utils'
 
 describe('module runner initialization', async () => {
   const it = await createModuleRunnerTester(
@@ -47,8 +51,10 @@ describe('module runner initialization', async () => {
       '    at Module.throwError (<root>/fixtures/throws-error-method.ts:6:9)',
     )
 
+    const fixtureEditor = createFixtureEditor()
+
     // simulate HMR
-    editFile(
+    fixtureEditor.editFile(
       resolvePath(import.meta.url, './fixtures/throws-error-method.ts'),
       (code) => '\n\n\n\n\n' + code + '\n',
     )

packages/vite/src/node/ssr/runtime/__tests__/utils.ts

@@ -2,7 +2,7 @@ import fs from 'node:fs'
 import { dirname, resolve } from 'node:path'
 import { fileURLToPath } from 'node:url'
 import type { TestAPI } from 'vitest'
-import { afterEach, beforeEach, test } from 'vitest'
+import { afterEach, beforeEach, onTestFinished, test } from 'vitest'
 import type { ModuleRunner } from 'vite/module-runner'
 import type { ServerModuleRunnerOptions } from '../serverModuleRunner'
 import type { ViteDevServer } from '../../../server'
@@ -103,32 +103,26 @@ export async function createModuleRunnerTester(
   return test as TestAPI<TestClient>
 }
 
-const originalFiles = new Map<string, string>()
-const createdFiles = new Set<string>()
-afterEach(() => {
-  originalFiles.forEach((content, file) => {
-    fs.writeFileSync(file, content, 'utf-8')
-  })
-  createdFiles.forEach((file) => {
-    if (fs.existsSync(file)) fs.unlinkSync(file)
-  })
-  originalFiles.clear()
-  createdFiles.clear()
-})
-
-export function createFile(file: string, content: string): void {
-  createdFiles.add(file)
-  fs.mkdirSync(dirname(file), { recursive: true })
-  fs.writeFileSync(file, content, 'utf-8')
+type FixtureEditor = {
+  editFile: (file: string, callback: (content: string) => string) => void
 }
 
-export function editFile(
-  file: string,
-  callback: (content: string) => string,
-): void {
-  const content = fs.readFileSync(file, 'utf-8')
-  if (!originalFiles.has(file)) originalFiles.set(file, content)
-  fs.writeFileSync(file, callback(content), 'utf-8')
+export function createFixtureEditor(): FixtureEditor {
+  const originalFiles = new Map<string, string>()
+  onTestFinished(() => {
+    originalFiles.forEach((content, file) => {
+      fs.writeFileSync(file, content, 'utf-8')
+    })
+    originalFiles.clear()
+  })
+
+  return {
+    editFile(file, callback) {
+      const content = fs.readFileSync(file, 'utf-8')
+      if (!originalFiles.has(file)) originalFiles.set(file, content)
+      fs.writeFileSync(file, callback(content), 'utf-8')
+    },
+  }
 }
 
 export function resolvePath(baseUrl: string, path: string): string {