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

Author: sheremet-va

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.
+:::

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)
+}

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()

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)

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,
+      }
+    }
+  }
+}

packages/runner/src/types/runner.ts

@@ -33,6 +33,7 @@ export interface VitestRunnerConfig {
   testTimeout: number
   hookTimeout: number
   retry: number
+  includeTaskLocation?: boolean
   diffOptions?: DiffOptions
 }
 

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 {

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) {

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>

packages/vitest/src/node/cli/cli-config.ts

@@ -615,4 +615,5 @@ export const cliOptionsConfig: VitestCLIOptions = {
   poolMatchGlobs: null,
   deps: null,
   name: null,
+  includeTaskLocation: null,
 }