feat: add a flag to include test location in tasks (#5342)

sheremet-va · web-flow · commit d627e209dd99 · 2024-03-14T15:54:40.000+01:00
diff --git a/docs/config/index.md b/docs/config/index.md
@@ -2108,3 +2108,16 @@ Disabling this option might [improve performance](/guide/improving-performance)
 ::: tip
 You can disable isolation for specific pools by using [`poolOptions`](#pooloptions) property.
 :::
+
+### includeTaskLocation <Badge type="info">1.4.0+</Badge> {#includeTaskLocation}
+
+- **Type:** `boolean`
+- **Default:** `false`
+
+Should `location` property be included when Vitest API receives tasks in [reporters](#reporters). If you have a lot of tests, this might cause a small performance regression.
+
+The `location` property has `column` and `line` values that correspond to the `test` or `describe` position in the original file.
+
+::: tip
+This option has no effect if you do not use custom code that relies on this.
+:::
diff --git a/packages/browser/src/client/runner.ts b/packages/browser/src/client/runner.ts
@@ -55,7 +55,13 @@ export function createBrowserRunner(
       }
     }
 
-    onCollected = (files: File[]): unknown => {
+    onCollected = async (files: File[]): Promise<unknown> => {
+      if (this.config.includeTaskLocation) {
+        try {
+          await updateFilesLocations(files)
+        }
+        catch (_) {}
+      }
       return rpc().onCollected(files)
     }
 
@@ -107,3 +113,28 @@ export async function initiateRunner() {
   cachedRunner = runner
   return runner
 }
+
+async function updateFilesLocations(files: File[]) {
+  const { loadSourceMapUtils } = await importId('vitest/utils') as typeof import('vitest/utils')
+  const { TraceMap, originalPositionFor } = await loadSourceMapUtils()
+
+  const promises = files.map(async (file) => {
+    const result = await rpc().getBrowserFileSourceMap(file.filepath)
+    if (!result)
+      return null
+    const traceMap = new TraceMap(result as any)
+    function updateLocation(task: Task) {
+      if (task.location) {
+        const { line, column } = originalPositionFor(traceMap, task.location)
+        if (line != null && column != null)
+          task.location = { line, column: column + 1 }
+      }
+      if ('tasks' in task)
+        task.tasks.forEach(updateLocation)
+    }
+    file.tasks.forEach(updateLocation)
+    return null
+  })
+
+  await Promise.all(promises)
+}
diff --git a/packages/runner/src/collect.ts b/packages/runner/src/collect.ts
@@ -28,7 +28,7 @@ export async function collectTests(paths: string[], runner: VitestRunner): Promi
       projectName: config.name,
     }
 
-    clearCollectorContext(runner)
+    clearCollectorContext(filepath, runner)
 
     try {
       const setupStart = now()
diff --git a/packages/runner/src/run.ts b/packages/runner/src/run.ts
@@ -393,7 +393,7 @@ export async function startTests(paths: string[], runner: VitestRunner) {
 
   const files = await collectTests(paths, runner)
 
-  runner.onCollected?.(files)
+  await runner.onCollected?.(files)
   await runner.onBeforeRunFiles?.(files)
 
   await runFiles(files, runner)
diff --git a/packages/runner/src/suite.ts b/packages/runner/src/suite.ts
@@ -1,4 +1,5 @@
 import { format, isObject, objDisplay, objectAttr } from '@vitest/utils'
+import { parseSingleStack } from '@vitest/utils/source-map'
 import type { Custom, CustomAPI, File, Fixtures, RunMode, Suite, SuiteAPI, SuiteCollector, SuiteFactory, SuiteHooks, Task, TaskCustomOptions, Test, TestAPI, TestFunction, TestOptions } from './types'
 import type { VitestRunner } from './types/runner'
 import { createChainable } from './utils/chain'
@@ -25,19 +26,25 @@ export const it = test
 
 let runner: VitestRunner
 let defaultSuite: SuiteCollector
+let currentTestFilepath: string
 
 export function getDefaultSuite() {
   return defaultSuite
 }
 
+export function getTestFilepath() {
+  return currentTestFilepath
+}
+
 export function getRunner() {
   return runner
 }
 
-export function clearCollectorContext(currentRunner: VitestRunner) {
+export function clearCollectorContext(filepath: string, currentRunner: VitestRunner) {
   if (!defaultSuite)
     defaultSuite = currentRunner.config.sequence.shuffle ? suite.shuffle('') : currentRunner.config.sequence.concurrent ? suite.concurrent('') : suite('')
   runner = currentRunner
+  currentTestFilepath = filepath
   collectorContext.tasks.length = 0
   defaultSuite.clear()
   collectorContext.currentSuite = defaultSuite
@@ -103,7 +110,7 @@ function createSuiteCollector(name: string, factory: SuiteFactory = () => { }, m
 
   let suite: Suite
 
-  initSuite()
+  initSuite(true)
 
   const task = function (name = '', options: TaskCustomOptions = {}) {
     const task: Custom = {
@@ -140,6 +147,17 @@ function createSuiteCollector(name: string, factory: SuiteFactory = () => { }, m
       ))
     }
 
+    if (runner.config.includeTaskLocation) {
+      const limit = Error.stackTraceLimit
+      // custom can be called from any place, let's assume the limit is 10 stacks
+      Error.stackTraceLimit = 10
+      const error = new Error('stacktrace').stack!
+      Error.stackTraceLimit = limit
+      const stack = findStackTrace(error)
+      if (stack)
+        task.location = stack
+    }
+
     tasks.push(task)
     return task
   }
@@ -183,7 +201,7 @@ function createSuiteCollector(name: string, factory: SuiteFactory = () => { }, m
     getHooks(suite)[name].push(...fn as any)
   }
 
-  function initSuite() {
+  function initSuite(includeLocation: boolean) {
     if (typeof suiteOptions === 'number')
       suiteOptions = { timeout: suiteOptions }
 
@@ -199,13 +217,27 @@ function createSuiteCollector(name: string, factory: SuiteFactory = () => { }, m
       projectName: '',
     }
 
+    if (runner && includeLocation && runner.config.includeTaskLocation) {
+      const limit = Error.stackTraceLimit
+      Error.stackTraceLimit = 5
+      const error = new Error('stacktrace').stack!
+      Error.stackTraceLimit = limit
+      const stack = parseSingleStack(error.split('\n')[5])
+      if (stack) {
+        suite.location = {
+          line: stack.line,
+          column: stack.column,
+        }
+      }
+    }
+
     setHooks(suite, createSuiteHooks())
   }
 
   function clear() {
     tasks.length = 0
     factoryQueue.length = 0
-    initSuite()
+    initSuite(false)
   }
 
   async function collect(file?: File) {
@@ -397,3 +429,18 @@ function formatTemplateString(cases: any[], args: any[]): any[] {
   }
   return res
 }
+
+function findStackTrace(error: string) {
+  // first line is the error message
+  // and the first 3 stacks are always from the collector
+  const lines = error.split('\n').slice(4)
+  for (const line of lines) {
+    const stack = parseSingleStack(line)
+    if (stack && stack.file === getTestFilepath()) {
+      return {
+        line: stack.line,
+        column: stack.column,
+      }
+    }
+  }
+}
diff --git a/packages/runner/src/types/runner.ts b/packages/runner/src/types/runner.ts
@@ -33,6 +33,7 @@ export interface VitestRunnerConfig {
   testTimeout: number
   hookTimeout: number
   retry: number
+  includeTaskLocation?: boolean
   diffOptions?: DiffOptions
 }
 
diff --git a/packages/runner/src/types/tasks.ts b/packages/runner/src/types/tasks.ts
@@ -18,6 +18,10 @@ export interface TaskBase {
   result?: TaskResult
   retry?: number
   repeats?: number
+  location?: {
+    line: number
+    column: number
+  }
 }
 
 export interface TaskPopulated extends TaskBase {
diff --git a/packages/vitest/src/api/setup.ts b/packages/vitest/src/api/setup.ts
@@ -113,6 +113,12 @@ export function setup(vitestOrWorkspace: Vitest | WorkspaceProject, _server?: Vi
         getConfig() {
           return vitestOrWorkspace.config
         },
+        async getBrowserFileSourceMap(id) {
+          if (!('ctx' in vitestOrWorkspace))
+            return undefined
+          const mod = vitestOrWorkspace.browser?.moduleGraph.getModuleById(id)
+          return mod?.transformResult?.map
+        },
         async getTransformResult(id) {
           const result: TransformResultWithSource | null | undefined = await ctx.vitenode.transformRequest(id)
           if (result) {
diff --git a/packages/vitest/src/api/types.ts b/packages/vitest/src/api/types.ts
@@ -20,6 +20,7 @@ export interface WebSocketHandlers {
   resolveSnapshotPath: (testPath: string) => string
   resolveSnapshotRawPath: (testPath: string, rawPath: string) => string
   getModuleGraph: (id: string) => Promise<ModuleGraphData>
+  getBrowserFileSourceMap: (id: string) => Promise<TransformResult['map'] | undefined>
   getTransformResult: (id: string) => Promise<TransformResultWithSource | undefined>
   readSnapshotFile: (id: string) => Promise<string | null>
   readTestFile: (id: string) => Promise<string | null>
diff --git a/packages/vitest/src/node/cli/cli-config.ts b/packages/vitest/src/node/cli/cli-config.ts
@@ -615,4 +615,5 @@ export const cliOptionsConfig: VitestCLIOptions = {
   poolMatchGlobs: null,
   deps: null,
   name: null,
+  includeTaskLocation: null,
 }
diff --git a/packages/vitest/src/node/workspace.ts b/packages/vitest/src/node/workspace.ts
@@ -382,6 +382,7 @@ export class WorkspaceProject {
       inspect: this.ctx.config.inspect,
       inspectBrk: this.ctx.config.inspectBrk,
       alias: [],
+      includeTaskLocation: this.config.includeTaskLocation ?? this.ctx.config.includeTaskLocation,
     }, this.ctx.configOverride || {} as any) as ResolvedConfig
   }
 
diff --git a/packages/vitest/src/public/utils.ts b/packages/vitest/src/public/utils.ts
@@ -1 +1,5 @@
 export * from '@vitest/utils'
+
+export function loadSourceMapUtils() {
+  return import('@vitest/utils/source-map')
+}
diff --git a/packages/vitest/src/types/config.ts b/packages/vitest/src/types/config.ts
@@ -715,6 +715,13 @@ export interface InlineConfig {
    * @default false
    */
   disableConsoleIntercept?: boolean
+
+  /**
+   * Include "location" property inside the test definition
+   *
+   * @default false
+   */
+  includeTaskLocation?: boolean
 }
 
 export interface TypecheckConfig {
diff --git a/test/public-api/fixtures/vitest.config.ts b/test/public-api/fixtures/vitest.config.ts
@@ -0,0 +1 @@
+export default {}
diff --git a/test/public-api/package.json b/test/public-api/package.json
@@ -3,7 +3,8 @@
   "type": "module",
   "private": true,
   "scripts": {
-    "test": "vitest"
+    "test": "vitest",
+    "fixtures": "vitest --root ./fixtures"
   },
   "devDependencies": {
     "@vitest/browser": "workspace:*",
diff --git a/test/public-api/tests/runner.spec.ts b/test/public-api/tests/runner.spec.ts
@@ -15,9 +15,10 @@ it.each([
       headless: true,
     },
   },
-] as UserConfig[])('passes down metadata when $name', async (config) => {
+] as UserConfig[])('passes down metadata when $name', { timeout: 60_000, retry: 3 }, async (config) => {
   const taskUpdate: TaskResultPack[] = []
   const finishedFiles: File[] = []
+  const collectedFiles: File[] = []
   const { vitest, stdout, stderr } = await runVitest({
     root: resolve(__dirname, '..', 'fixtures'),
     include: ['**/*.spec.ts'],
@@ -30,8 +31,12 @@ it.each([
         onFinished(files) {
           finishedFiles.push(...files || [])
         },
+        onCollected(files) {
+          collectedFiles.push(...files || [])
+        },
       },
     ],
+    includeTaskLocation: true,
     ...config,
   })
 
@@ -69,7 +74,17 @@ it.each([
 
   expect(files[0].meta).toEqual(suiteMeta)
   expect(files[0].tasks[0].meta).toEqual(testMeta)
-}, {
-  timeout: 60_000,
-  retry: 3,
+
+  expect(finishedFiles[0].tasks[0].location).toEqual({
+    line: 14,
+    column: 1,
+  })
+  expect(collectedFiles[0].tasks[0].location).toEqual({
+    line: 14,
+    column: 1,
+  })
+  expect(files[0].tasks[0].location).toEqual({
+    line: 14,
+    column: 1,
+  })
 })

Original file line number	Diff line number	Diff line change
`@@ -28,7 +28,7 @@ export async function collectTests(paths: string[], runner: VitestRunner): Promi`
`28`	`28`	`projectName: config.name,`
`29`	`29`	`}`
`30`	`30`
`31`		`- clearCollectorContext(runner)`
	`31`	`+ clearCollectorContext(filepath, runner)`
`32`	`32`
`33`	`33`	`try {`
`34`	`34`	`const setupStart = now()`
Original file line number	Diff line number	Diff line change
`@@ -33,6 +33,7 @@ export interface VitestRunnerConfig {`
`33`	`33`	`testTimeout: number`
`34`	`34`	`hookTimeout: number`
`35`	`35`	`retry: number`
	`36`	`+ includeTaskLocation?: boolean`
`36`	`37`	`diffOptions?: DiffOptions`
`37`	`38`	`}`
`38`	`39`
Original file line number	Diff line number	Diff line change
`@@ -615,4 +615,5 @@ export const cliOptionsConfig: VitestCLIOptions = {`
`615`	`615`	`poolMatchGlobs: null,`
`616`	`616`	`deps: null,`
`617`	`617`	`name: null,`
	`618`	`+ includeTaskLocation: null,`
`618`	`619`	`}`