feat: allow coverage providers to store data in fs module cache (#11747)

Author: AriPerkkio

packages/vitest/src/node/cache/fsModuleCache.ts

@@ -43,7 +43,7 @@ export class FileSystemModuleCache {
   private rootCache: string
   private metadataFilePath: string
 
-  private version = '1.0.0-beta.8'
+  private version = '1.0.0-beta.9'
   private fsCacheRoots = new WeakMap<ResolvedConfig, string>()
   private fsEnvironmentHashMap = new WeakMap<DevEnvironment, string>()
   private fsCacheKeyGenerators = new WeakMap<ResolvedConfig, Set<CacheKeyIdGenerator>>()
@@ -152,6 +152,8 @@ export class FileSystemModuleCache {
       deps: meta.deps,
       dynamicDeps: meta.dynamicDeps,
       staticMocks: meta.staticMocks,
+      hasCoverageData: meta.hasCoverageData,
+      coverageData: meta.hasCoverageData ? await readCoverageData(cachedFilePath) : undefined,
     }
   }
 
@@ -161,6 +163,7 @@ export class FileSystemModuleCache {
     transformResult: TransformResult | null,
     imports: CachedModuleImports = { urls: [], ids: {} },
     mappings: boolean = false,
+    coverageData?: unknown,
   ): Promise<void> {
     if ('code' in fetchResult) {
       const result = {
@@ -172,15 +175,40 @@ export class FileSystemModuleCache {
         deps: transformResult?.deps,
         dynamicDeps: transformResult?.dynamicDeps,
         staticMocks: transformResult?.__vitestStaticMocks,
-      } satisfies Omit<CachedInlineModuleMeta, 'code'>
+        hasCoverageData: coverageData !== undefined,
+      } satisfies Omit<CachedInlineModuleMeta, 'code' | 'coverageData'>
       debugFs?.(`${c.yellow('[write]')} ${fetchResult.id} is cached in ${cachedFilePath}`)
-      await atomicWriteFile(
-        cachedFilePath,
-        `${fetchResult.code}${cacheComment}${this.toBase64(result)}`,
-      )
+
+      await Promise.all([
+        // Coverage data is stored in separate file so that it's not transferred to the workers
+        result.hasCoverageData &&
+          atomicWriteFile(getCoverageDataPath(cachedFilePath), JSON.stringify(coverageData)),
+
+        atomicWriteFile(
+          cachedFilePath,
+          `${fetchResult.code}${cacheComment}${this.toBase64(result)}`,
+        ),
+      ])
     }
   }
 
+  getCoverageData(environment: DevEnvironment, id: string): unknown {
+    const coverageProvider = this.vitest.coverageProvider
+
+    if (!this.vitest.config.coverage.enabled || !coverageProvider?.requiresTransform?.(id)) {
+      return
+    }
+
+    return coverageProvider.getModuleCacheData?.(environment, id)
+  }
+
+  /**
+   * Passes the data of the coverage provider back to it when the cached module is used
+   */
+  restoreCoverageData(environment: DevEnvironment, id: string, data: unknown): void {
+    this.vitest.coverageProvider?.restoreModuleCacheData?.(environment, id, data)
+  }
+
   private toBase64(obj: unknown) {
     const json = stringify(obj)
     return Buffer.from(json).toString('base64')
@@ -242,9 +270,13 @@ export class FileSystemModuleCache {
     const config = environment.config
     // coverage provider is dynamic, so we also clear the whole cache if
     // vitest.enableCoverage/vitest.disableCoverage is called
-    const coverageAffectsCache = String(
-      this.vitest.config.coverage.enabled && this.vitest.coverageProvider?.requiresTransform?.(id),
-    )
+    const coverageProvider = this.vitest.coverageProvider
+
+    const coverageAffectsCache =
+      this.vitest.config.coverage.enabled && coverageProvider?.requiresTransform?.(id)
+        ? `${coverageProvider.name}@${coverageProvider.version}`
+        : 'false'
+
     let environmentHash = this.fsEnvironmentHashMap.get(environment)
     if (!environmentHash) {
       const cacheConfig = JSON.stringify(
@@ -441,6 +473,8 @@ export interface CachedInlineModuleMeta {
   deps?: string[]
   dynamicDeps?: string[]
   staticMocks?: StaticMockCall[] | null
+  hasCoverageData?: boolean
+  coverageData?: unknown
 }
 
 export interface CachedModuleImports {
@@ -591,3 +625,13 @@ function tryStatSync(file: string): fs.Stats | undefined {
     // Ignore errors
   }
 }
+
+async function readCoverageData(cachedFilePath: string): Promise<unknown> {
+  try {
+    return JSON.parse(await readFile(getCoverageDataPath(cachedFilePath), 'utf-8'))
+  } catch {}
+}
+
+function getCoverageDataPath(cachedFilePath: string): string {
+  return `${cachedFilePath}.coverage.json`
+}

packages/vitest/src/node/environments/transformService.ts

@@ -220,13 +220,15 @@ export class ModuleTransformService {
     const imports = this.getCachedImports(environment, moduleGraphModule)
     const map = moduleGraphModule.transformResult?.map
     const mappings = map && !('version' in map) && map.mappings === ''
+    const coverageData = this.fsCache.getCoverageData(environment, moduleGraphModule.id!)
 
     const cachedResult = await this.cacheResult(
       result,
       cachePath,
       moduleGraphModule.transformResult,
       imports,
       !!mappings,
+      coverageData,
     )
     // remember where the code is stored on disk so that repeat fetches and the
     // `fetchWarmModules` snapshot can point at it in this session already, not
@@ -363,6 +365,11 @@ export class ModuleTransformService {
       return
     }
 
+    // Notify coverage provider that cached results were read
+    if (cachedModule.hasCoverageData) {
+      this.fsCache.restoreCoverageData(environment, cachedModule.id, cachedModule.coverageData)
+    }
+
     // keep the module graph in sync
     let map: Rollup.SourceMap | null | { mappings: '' } = extractSourceMap(cachedModule.code)
     if (map && cachedModule.file) {
@@ -508,6 +515,7 @@ export class ModuleTransformService {
     transformResult: TransformResult | null,
     imports?: CachedModuleImports,
     mappings = false,
+    coverageData?: unknown,
   ): Promise<FetchResult | FetchCachedFileSystemResult> {
     const returnResult = 'code' in result ? getCachedResult(result, cachePath) : result
 
@@ -516,7 +524,7 @@ export class ModuleTransformService {
     }
 
     const savePromise = this.fsCache
-      .saveCachedModule(cachePath, result, transformResult, imports, mappings)
+      .saveCachedModule(cachePath, result, transformResult, imports, mappings, coverageData)
       .then(() => returnResult)
       .catch((error) => {
         debugFs?.(`failed to cache ${cachePath}, serving it inline: ${error}`)

packages/vitest/src/node/types/coverage.ts

@@ -1,6 +1,6 @@
 import type { ReportOptions } from '@vitest/istanbul-lib-report'
 import type { Arrayable } from '@vitest/utils'
-import type { ESTree, TransformResult as ViteTransformResult } from 'vite'
+import type { DevEnvironment, ESTree, TransformResult as ViteTransformResult } from 'vite'
 import type { AfterSuiteRunMeta } from '../../types/general'
 import type {
   RuntimeCoverageModuleLoader,
@@ -57,6 +57,27 @@ export interface CoverageProvider {
    */
   requiresTransform?: (id: string) => boolean
 
+  /**
+   * Version of the coverage provider. Can be used as cache key in various caching mechanisms.
+   * @experimental
+   */
+  version?: string
+
+  /**
+   * Return data that should be stored with the transformed module in `fsModuleCache`,
+   * e.g. data collected in `onFileTransform`. The data must be JSON serializable.
+   * When the module is later loaded from `fsModuleCache`, `onFileTransform` is not called.
+   * The data is passed to `restoreModuleCacheData` instead.
+   * @experimental
+   */
+  getModuleCacheData?: (environment: DevEnvironment, id: string) => unknown
+
+  /**
+   * Called with the data returned by `getModuleCacheData` when the module is loaded from `fsModuleCache`.
+   * @experimental
+   */
+  restoreModuleCacheData?: (environment: DevEnvironment, id: string, data: unknown) => void
+
   /** Callback that's called when the coverage is enabled via a programmatic `enableCoverage` API. */
   onEnabled?: () => void | Promise<void>
 }

test/coverage-test/fixtures/custom-provider.ts

@@ -48,6 +48,7 @@ class CustomCoverageProvider implements CoverageProvider {
   calls: Set<string> = new Set()
   coverageReports: Set<string> = new Set()
   transformedFiles: Set<string> = new Set()
+  restoredFiles: Set<string> = new Set()
 
   initialize(ctx: Vitest) {
     this.options = ctx.config.coverage
@@ -83,6 +84,7 @@ class CustomCoverageProvider implements CoverageProvider {
       calls: Array.from(this.calls.values()),
       coverageReports: Array.from(this.coverageReports.values()).sort(),
       transformedFiles: Array.from(this.transformedFiles.values()).sort(),
+      restoredFiles: Array.from(this.restoredFiles.values()).sort(),
     }, null, 2)
 
     if (existsSync('./coverage')) {
@@ -103,6 +105,18 @@ class CustomCoverageProvider implements CoverageProvider {
     return { code }
   }
 
+  requiresTransform(id: string) {
+    return /\/fixtures\/src\//.test(normalizeFilename(id))
+  }
+
+  getModuleCacheData(_environment: unknown, id: string) {
+    return { cachedBy: this.name, filename: normalizeFilename(id).split('?')[0] }
+  }
+
+  restoreModuleCacheData(_environment: unknown, id: string, data: unknown) {
+    this.restoredFiles.add(`${normalizeFilename(id).split('?')[0]} ${JSON.stringify(data)}`)
+  }
+
   resolveOptions(): ResolvedCoverageOptions {
     this.calls.add('resolveOptions')
 

test/coverage-test/test/custom-provider.custom.test.ts

@@ -1,4 +1,4 @@
-import { readFileSync } from 'node:fs'
+import { readFileSync, rmSync } from 'node:fs'
 import { expect } from 'vitest'
 import { runVitest, test } from '../utils'
 
@@ -26,7 +26,43 @@ test('custom provider', async () => {
       "transformedFiles": [
         "<process-cwd>/fixtures/src/even.ts",
         "<process-cwd>/fixtures/src/math.ts"
-      ]
+      ],
+      "restoredFiles": []
     }"
   `)
 })
+
+test('custom provider restores its data from fs module cache', async () => {
+  const fsModuleCachePath = './node_modules/.vitest-custom-provider-cache'
+  rmSync(fsModuleCachePath, { recursive: true, force: true })
+
+  const config = {
+    include: ['fixtures/test/math.test.ts', 'fixtures/test/even.test.ts'],
+    fsModuleCache: true,
+    fsModuleCachePath,
+  }
+
+  // first run transforms the files and stores their data in the cache
+  await runVitest(config)
+  const first = JSON.parse(readFileSync('./coverage/custom-coverage-provider-report.json', 'utf-8'))
+
+  expect(first.transformedFiles).toEqual([
+    '<process-cwd>/fixtures/src/even.ts',
+    '<process-cwd>/fixtures/src/math.ts',
+  ])
+  expect(first.restoredFiles).toEqual([])
+
+  // second run loads the files from the cache and doesn't transform them
+  await runVitest(config)
+  const second = JSON.parse(
+    readFileSync('./coverage/custom-coverage-provider-report.json', 'utf-8'),
+  )
+
+  expect(second.transformedFiles).toEqual([])
+  expect(second.restoredFiles).toEqual([
+    '<process-cwd>/fixtures/src/even.ts {"cachedBy":"custom-coverage-provider","filename":"<process-cwd>/fixtures/src/even.ts"}',
+    '<process-cwd>/fixtures/src/math.ts {"cachedBy":"custom-coverage-provider","filename":"<process-cwd>/fixtures/src/math.ts"}',
+  ])
+
+  rmSync(fsModuleCachePath, { recursive: true, force: true })
+})