fix!(runner): correctly process custom tasks, update runner hooks nam…

…ing (vitest-dev#4076)

docs/advanced/runner.md

@@ -20,26 +20,26 @@ export interface VitestRunner {
   /**
    * Called when test runner should cancel next test runs.
    * Runner should listen for this method and mark tests and suites as skipped in
-   * "onBeforeRunSuite" and "onBeforeRunTest" when called.
+   * "onBeforeRunSuite" and "onBeforeRunTask" when called.
    */
   onCancel?(reason: CancelReason): unknown
 
   /**
    * Called before running a single test. Doesn't have "result" yet.
    */
-  onBeforeRunTest?(test: Test): unknown
+  onBeforeRunTask?(test: TaskPopulated): unknown
   /**
    * Called before actually running the test function. Already has "result" with "state" and "startTime".
    */
-  onBeforeTryTest?(test: Test, retryCount: number): unknown
+  onBeforeTryTask?(test: TaskPopulated, options: { retry: number; repeats: number }): unknown
   /**
    * Called after result and state are set.
    */
-  onAfterRunTest?(test: Test): unknown
+  onAfterRunTask?(test: TaskPopulated): unknown
   /**
    * Called right after running the test function. Doesn't have new state yet. Will not be called, if the test function throws.
    */
-  onAfterTryTest?(test: Test, retryCount: number): unknown
+  onAfterTryTask?(test: TaskPopulated, options: { retry: number; repeats: number }): unknown
 
   /**
    * Called before running a single suite. Doesn't have "result" yet.
@@ -59,7 +59,7 @@ export interface VitestRunner {
    * If defined, will be called instead of usual Vitest handling. Useful, if you have your custom test function.
    * "before" and "after" hooks will not be ignored.
    */
-  runTest?(test: Test): Promise<void>
+  runTask?(test: TaskPopulated): Promise<void>
 
   /**
    * Called, when a task is updated. The same as "onTaskUpdate" in a reporter, but this is running in the same thread as tests.
@@ -69,16 +69,20 @@ export interface VitestRunner {
   /**
    * Called before running all tests in collected paths.
    */
-  onBeforeRun?(files: File[]): unknown
+  onBeforeRunFiles?(files: File[]): unknown
   /**
    * Called right after running all tests in collected paths.
    */
-  onAfterRun?(files: File[]): unknown
+  onAfterRunFiles?(files: File[]): unknown
   /**
    * Called when new context for a test is defined. Useful, if you want to add custom properties to the context.
    * If you only want to define custom context with a runner, consider using "beforeAll" in "setupFiles" instead.
+   *
+   * This method is called for both "test" and "custom" handlers.
+   *
+   * @see https://vitest.dev/advanced/runner.html#your-task-function
    */
-  extendTestContext?(context: TestContext): TestContext
+  extendTaskContext?<T extends Test | Custom>(context: TaskContext<T>): TaskContext<T>
   /**
    * Called, when certain files are imported. Can be called in two situations: when collecting tests and when importing setup files.
    */
@@ -108,18 +112,23 @@ You can extend Vitest task system with your tasks. A task is an object that is p
 
 ```js
 // ./utils/custom.js
-import { getCurrentSuite, setFn } from 'vitest/suite'
+import { createTaskCollector, getCurrentSuite, setFn } from 'vitest/suite'
 
 export { describe, beforeAll, afterAll } from 'vitest'
 
-// this function will be called, when Vitest collects tasks
-export const myCustomTask = function (name, fn) {
-  const task = getCurrentSuite().custom(name)
-  task.meta = {
-    customPropertyToDifferentiateTask: true
-  }
-  setFn(task, fn || (() => {}))
-}
+// this function will be called when Vitest collects tasks
+// createTaskCollector just provides all "todo"/"each"/... support, you don't have to use it
+// To support custom tasks, you just need to call "getCurrentSuite().task()"
+export const myCustomTask = createTaskCollector(function (name, fn, timeout) {
+  getCurrentSuite().task(name, {
+    ...this, // so "todo"/"skip" is tracked correctly
+    meta: {
+      customPropertyToDifferentiateTask: true
+    },
+    handler: fn,
+    timeout,
+  })
+})
 ```
 
 ```js
@@ -135,6 +144,9 @@ describe('take care of the garden', () => {
   myCustomTask('weed the grass', () => {
     gardener.weedTheGrass()
   })
+  myCustomTask.todo('mow the lawn', () => {
+    gardener.mowerTheLawn()
+  })
   myCustomTask('water flowers', () => {
     gardener.waterFlowers()
   })

packages/browser/src/client/runner.ts

@@ -22,7 +22,7 @@ export function createBrowserRunner(original: any, coverageModule: CoverageHandl
       this.hashMap = options.browserHashMap
     }
 
-    async onAfterRunTest(task: Test) {
+    async onAfterRunTask(task: Test) {
       await super.onAfterRunTest?.(task)
       task.result?.errors?.forEach((error) => {
         console.error(error.message)
@@ -39,7 +39,7 @@ export function createBrowserRunner(original: any, coverageModule: CoverageHandl
       }
     }
 
-    async onAfterRun() {
+    async onAfterRunFiles() {
       await super.onAfterRun?.()
       const coverage = await coverageModule?.takeCoverage?.()
       if (coverage)

packages/runner/src/context.ts

@@ -1,6 +1,6 @@
 import type { Awaitable } from '@vitest/utils'
 import { getSafeTimers } from '@vitest/utils'
-import type { RuntimeContext, SuiteCollector, Test, TestContext } from './types'
+import type { Custom, ExtendedContext, RuntimeContext, SuiteCollector, TaskContext, Test } from './types'
 import type { VitestRunner } from './types/runner'
 import { PendingError } from './errors'
 
@@ -42,12 +42,11 @@ export function withTimeout<T extends((...args: any[]) => any)>(
   }) as T
 }
 
-export function createTestContext(test: Test, runner: VitestRunner): TestContext {
+export function createTestContext<T extends Test | Custom>(test: T, runner: VitestRunner): ExtendedContext<T> {
   const context = function () {
     throw new Error('done() callback is deprecated, use promise instead')
-  } as unknown as TestContext
+  } as unknown as TaskContext<T>
 
-  context.meta = test
   context.task = test
 
   context.skip = () => {
@@ -60,7 +59,7 @@ export function createTestContext(test: Test, runner: VitestRunner): TestContext
     test.onFailed.push(fn)
   }
 
-  return runner.extendTestContext?.(context) || context
+  return runner.extendTaskContext?.(context) as ExtendedContext<T> || context
 }
 
 function makeTimeoutMsg(isHook: boolean, timeout: number) {

packages/runner/src/hooks.ts

@@ -1,4 +1,4 @@
-import type { OnTestFailedHandler, SuiteHooks, Test } from './types'
+import type { OnTestFailedHandler, SuiteHooks, TaskPopulated } from './types'
 import { getCurrentSuite, getRunner } from './suite'
 import { getCurrentTest } from './test-state'
 import { withTimeout } from './context'
@@ -27,7 +27,7 @@ export const onTestFailed = createTestHook<OnTestFailedHandler>('onTestFailed',
   test.onFailed.push(handler)
 })
 
-function createTestHook<T>(name: string, handler: (test: Test, handler: T) => void) {
+function createTestHook<T>(name: string, handler: (test: TaskPopulated, handler: T) => void) {
   return (fn: T) => {
     const current = getCurrentTest()
 

packages/runner/src/index.ts

@@ -1,5 +1,5 @@
 export { startTests, updateTask } from './run'
-export { test, it, describe, suite, getCurrentSuite } from './suite'
+export { test, it, describe, suite, getCurrentSuite, createTaskCollector } from './suite'
 export { beforeAll, beforeEach, afterAll, afterEach, onTestFailed } from './hooks'
 export { setFn, getFn } from './map'
 export { getCurrentTest } from './test-state'

packages/runner/src/map.ts

@@ -1,17 +1,17 @@
 import type { Awaitable } from '@vitest/utils'
-import type { Suite, SuiteHooks, Test, TestContext } from './types'
+import type { Custom, Suite, SuiteHooks, Test, TestContext } from './types'
 import type { FixtureItem } from './fixture'
 
 // use WeakMap here to make the Test and Suite object serializable
 const fnMap = new WeakMap()
 const fixtureMap = new WeakMap()
 const hooksMap = new WeakMap()
 
-export function setFn(key: Test, fn: (() => Awaitable<void>)) {
+export function setFn(key: Test | Custom, fn: (() => Awaitable<void>)) {
   fnMap.set(key, fn)
 }
 
-export function getFn<Task = Test>(key: Task): (() => Awaitable<void>) {
+export function getFn<Task = Test | Custom>(key: Task): (() => Awaitable<void>) {
   return fnMap.get(key as any)
 }
 

packages/runner/src/run.ts

@@ -3,7 +3,7 @@ import { getSafeTimers, shuffle } from '@vitest/utils'
 import { processError } from '@vitest/utils/error'
 import type { DiffOptions } from '@vitest/utils/diff'
 import type { VitestRunner } from './types/runner'
-import type { File, HookCleanupCallback, HookListener, SequenceHooks, Suite, SuiteHooks, Task, TaskMeta, TaskResult, TaskResultPack, TaskState, Test } from './types'
+import type { Custom, File, HookCleanupCallback, HookListener, SequenceHooks, Suite, SuiteHooks, Task, TaskMeta, TaskResult, TaskResultPack, TaskState, Test } from './types'
 import { partitionSuiteChildren } from './utils/suite'
 import { getFn, getHooks } from './map'
 import { collectTests } from './collect'
@@ -114,8 +114,8 @@ async function callCleanupHooks(cleanups: HookCleanupCallback[]) {
   }))
 }
 
-export async function runTest(test: Test, runner: VitestRunner) {
-  await runner.onBeforeRunTest?.(test)
+export async function runTest(test: Test | Custom, runner: VitestRunner) {
+  await runner.onBeforeRunTask?.(test)
 
   if (test.mode !== 'run')
     return
@@ -142,14 +142,14 @@ export async function runTest(test: Test, runner: VitestRunner) {
     for (let retryCount = 0; retryCount <= retry; retryCount++) {
       let beforeEachCleanups: HookCleanupCallback[] = []
       try {
-        await runner.onBeforeTryTest?.(test, { retry: retryCount, repeats: repeatCount })
+        await runner.onBeforeTryTask?.(test, { retry: retryCount, repeats: repeatCount })
 
         test.result.repeatCount = repeatCount
 
         beforeEachCleanups = await callSuiteHook(test.suite, test, 'beforeEach', runner, [test.context, test.suite])
 
-        if (runner.runTest) {
-          await runner.runTest(test)
+        if (runner.runTask) {
+          await runner.runTask(test)
         }
         else {
           const fn = getFn(test)
@@ -165,7 +165,7 @@ export async function runTest(test: Test, runner: VitestRunner) {
             throw errors
         }
 
-        await runner.onAfterTryTest?.(test, { retry: retryCount, repeats: repeatCount })
+        await runner.onAfterTryTask?.(test, { retry: retryCount, repeats: repeatCount })
 
         if (test.result.state !== 'fail') {
           if (!test.repeats)
@@ -230,7 +230,7 @@ export async function runTest(test: Test, runner: VitestRunner) {
 
   test.result.duration = now() - start
 
-  await runner.onAfterRunTest?.(test)
+  await runner.onAfterRunTask?.(test)
 
   updateTask(test, runner)
 }
@@ -356,7 +356,7 @@ export async function runSuite(suite: Suite, runner: VitestRunner) {
 }
 
 async function runSuiteChild(c: Task, runner: VitestRunner) {
-  if (c.type === 'test')
+  if (c.type === 'test' || c.type === 'custom')
     return runTest(c, runner)
 
   else if (c.type === 'suite')
@@ -385,11 +385,11 @@ export async function startTests(paths: string[], runner: VitestRunner) {
   const files = await collectTests(paths, runner)
 
   runner.onCollected?.(files)
-  await runner.onBeforeRun?.(files)
+  await runner.onBeforeRunFiles?.(files)
 
   await runFiles(files, runner)
 
-  await runner.onAfterRun?.(files)
+  await runner.onAfterRunFiles?.(files)
 
   await sendTasksUpdate(runner)