feat: introduce retry option for tests (vitest-dev#1929)

docs/api/index.md

@@ -9,6 +9,11 @@ The following types are used in the type signatures below
 ```ts
 type Awaitable<T> = T | PromiseLike<T>
 type TestFunction = () => Awaitable<void>
+
+interface TestOptions {
+  timeout?: number
+  retry?: number
+}
 ```
 
 When a test function returns a promise, the runner will wait until it is resolved to collect async expectations. If the promise is rejected, the test will fail.
@@ -19,7 +24,7 @@ In Jest, `TestFunction` can also be of type `(done: DoneCallback) => void`. If t
 
 ## test
 
-- **Type:** `(name: string, fn: TestFunction, timeout?: number) => void`
+- **Type:** `(name: string, fn: TestFunction, timeout?: number | TestOptions) => void`
 - **Alias:** `it`
 
   `test` defines a set of related expectations. It receives the test name and a function that holds the expectations to test.
@@ -36,7 +41,7 @@ In Jest, `TestFunction` can also be of type `(done: DoneCallback) => void`. If t
 
 ### test.skip
 
-- **Type:** `(name: string, fn: TestFunction, timeout?: number) => void`
+- **Type:** `(name: string, fn: TestFunction, timeout?: number | TestOptions) => void`
 - **Alias:** `it.skip`
 
   If you want to skip running certain tests, but you don't want to delete the code due to any reason, you can use `test.skip` to avoid running them.
@@ -330,7 +335,7 @@ When you use `test` in the top level of file, they are collected as part of the
 
 ### describe.shuffle
 
-- **Type:** `(name: string, fn: TestFunction, timeout?: number) => void`
+- **Type:** `(name: string, fn: TestFunction, timeout?: number | TestOptions) => void`
 
   Vitest provides a way to run all tests in random order via CLI flag [`--sequence.shuffle`](/guide/cli) or config option [`sequence.shuffle`](/config/#sequence-shuffle), but if you want to have only part of your test suite to run tests in random order, you can mark it with this flag.
 

packages/vitest/src/node/reporters/renderers/listRenderer.ts

@@ -92,6 +92,9 @@ export function renderTree(tasks: Task[], options: ListRendererOptions, level =
     let suffix = ''
     const prefix = ` ${getStateSymbol(task)} `
 
+    if (task.type === 'test' && task.result?.retryCount && task.result.retryCount > 1)
+      suffix += c.yellow(` (retry x${task.result.retryCount})`)
+
     if (task.type === 'suite')
       suffix += c.dim(` (${getTests(task).length})`)
 

packages/vitest/src/runtime/run.ts

@@ -108,48 +108,60 @@ export async function runTest(test: Test) {
 
   workerState.current = test
 
-  let beforeEachCleanups: HookCleanupCallback[] = []
-  try {
-    beforeEachCleanups = await callSuiteHook(test.suite, test, 'beforeEach', [test.context, test.suite])
-    setState({
-      assertionCalls: 0,
-      isExpectingAssertions: false,
-      isExpectingAssertionsError: null,
-      expectedAssertionsNumber: null,
-      expectedAssertionsNumberErrorGen: null,
-      testPath: test.suite.file?.filepath,
-      currentTestName: getFullName(test),
-    }, (globalThis as any)[GLOBAL_EXPECT])
-    await getFn(test)()
-    const {
-      assertionCalls,
-      expectedAssertionsNumber,
-      expectedAssertionsNumberErrorGen,
-      isExpectingAssertions,
-      isExpectingAssertionsError,
+  const retry = test.retry || 1
+  for (let retryCount = 0; retryCount < retry; retryCount++) {
+    let beforeEachCleanups: HookCleanupCallback[] = []
+    try {
+      beforeEachCleanups = await callSuiteHook(test.suite, test, 'beforeEach', [test.context, test.suite])
+      setState({
+        assertionCalls: 0,
+        isExpectingAssertions: false,
+        isExpectingAssertionsError: null,
+        expectedAssertionsNumber: null,
+        expectedAssertionsNumberErrorGen: null,
+        testPath: test.suite.file?.filepath,
+        currentTestName: getFullName(test),
+      }, (globalThis as any)[GLOBAL_EXPECT])
+
+      test.result.retryCount = retryCount
+
+      await getFn(test)()
+      const {
+        assertionCalls,
+        expectedAssertionsNumber,
+        expectedAssertionsNumberErrorGen,
+        isExpectingAssertions,
+        isExpectingAssertionsError,
       // @ts-expect-error local is private
-    } = test.context._local
-      ? test.context.expect.getState()
-      : getState((globalThis as any)[GLOBAL_EXPECT])
-    if (expectedAssertionsNumber !== null && assertionCalls !== expectedAssertionsNumber)
-      throw expectedAssertionsNumberErrorGen!()
-    if (isExpectingAssertions === true && assertionCalls === 0)
-      throw isExpectingAssertionsError
-
-    test.result.state = 'pass'
-  }
-  catch (e) {
-    test.result.state = 'fail'
-    test.result.error = processError(e)
-  }
+      } = test.context._local
+        ? test.context.expect.getState()
+        : getState((globalThis as any)[GLOBAL_EXPECT])
+      if (expectedAssertionsNumber !== null && assertionCalls !== expectedAssertionsNumber)
+        throw expectedAssertionsNumberErrorGen!()
+      if (isExpectingAssertions === true && assertionCalls === 0)
+        throw isExpectingAssertionsError
 
-  try {
-    await callSuiteHook(test.suite, test, 'afterEach', [test.context, test.suite])
-    await Promise.all(beforeEachCleanups.map(i => i?.()))
-  }
-  catch (e) {
-    test.result.state = 'fail'
-    test.result.error = processError(e)
+      test.result.state = 'pass'
+    }
+    catch (e) {
+      test.result.state = 'fail'
+      test.result.error = processError(e)
+    }
+
+    try {
+      await callSuiteHook(test.suite, test, 'afterEach', [test.context, test.suite])
+      await Promise.all(beforeEachCleanups.map(i => i?.()))
+    }
+    catch (e) {
+      test.result.state = 'fail'
+      test.result.error = processError(e)
+    }
+
+    if (test.result.state === 'pass')
+      break
+
+    // update retry info
+    updateTask(test)
   }
 
   // if test is marked to be failed, flip the result

packages/vitest/src/runtime/suite.ts

@@ -1,5 +1,5 @@
 import util from 'util'
-import type { BenchFunction, BenchOptions, Benchmark, BenchmarkAPI, File, RunMode, Suite, SuiteAPI, SuiteCollector, SuiteFactory, SuiteHooks, Task, Test, TestAPI, TestFunction } from '../types'
+import type { BenchFunction, BenchOptions, Benchmark, BenchmarkAPI, File, RunMode, Suite, SuiteAPI, SuiteCollector, SuiteFactory, SuiteHooks, Task, Test, TestAPI, TestFunction, TestOptions } from '../types'
 import { getWorkerState, isObject, isRunningInBenchmark, isRunningInTest, noop } from '../utils'
 import { createChainable } from './chain'
 import { collectTask, collectorContext, createTestContext, runWithSuite, withTimeout } from './context'
@@ -8,8 +8,8 @@ import { getHooks, setFn, setHooks } from './map'
 // apis
 export const suite = createSuite()
 export const test = createTest(
-  function (name: string, fn?: TestFunction, timeout?: number) {
-    getCurrentSuite().test.fn.call(this, name, fn, timeout)
+  function (name: string, fn?: TestFunction, options?: number | TestOptions) {
+    getCurrentSuite().test.fn.call(this, name, fn, options)
   },
 )
 
@@ -76,20 +76,25 @@ function createSuiteCollector(name: string, factory: SuiteFactory = () => { }, m
 
   initSuite()
 
-  const test = createTest(function (name: string, fn = noop, timeout?: number) {
+  const test = createTest(function (name: string, fn = noop, options?: number | TestOptions) {
     if (!isRunningInTest())
       throw new Error('`test()` and `it()` is only available in test mode.')
 
     const mode = this.only ? 'only' : this.skip ? 'skip' : this.todo ? 'todo' : 'run'
 
+    if (typeof options === 'number')
+      options = { timeout: options }
+
     const test: Test = {
       id: '',
       type: 'test',
       name,
       mode,
       suite: undefined!,
       fails: this.fails,
+      retry: options?.retry,
     } as Omit<Test, 'context'> as Test
+
     if (this.concurrent || concurrent)
       test.concurrent = true
     if (shuffle)
@@ -104,7 +109,7 @@ function createSuiteCollector(name: string, factory: SuiteFactory = () => { }, m
 
     setFn(test, withTimeout(
       () => fn(context),
-      timeout,
+      options?.timeout,
     ))
 
     tasks.push(test)
@@ -220,18 +225,18 @@ function createTest(fn: (
     this: Record<'concurrent' | 'skip' | 'only' | 'todo' | 'fails', boolean | undefined>,
     title: string,
     fn?: TestFunction,
-    timeout?: number
+    options?: number | TestOptions
   ) => void
 )) {
   const testFn = fn as any
 
   testFn.each = function<T>(this: { withContext: () => TestAPI }, cases: ReadonlyArray<T>) {
     const test = this.withContext()
 
-    return (name: string, fn: (...args: T[]) => void, timeout?: number) => {
+    return (name: string, fn: (...args: T[]) => void, options?: number | TestOptions) => {
       cases.forEach((i, idx) => {
         const items = Array.isArray(i) ? i : [i]
-        test(formatTitle(name, items, idx), () => fn(...items), timeout)
+        test(formatTitle(name, items, idx), () => fn(...items), options)
       })
     }
   }

packages/vitest/src/types/tasks.ts

@@ -15,6 +15,7 @@ export interface TaskBase {
   suite?: Suite
   file?: File
   result?: TaskResult
+  retry?: number
   logs?: UserConsoleLog[]
 }
 
@@ -26,6 +27,7 @@ export interface TaskResult {
   error?: ErrorWithDiff
   htmlError?: string
   hooks?: Partial<Record<keyof SuiteHooks, TaskState>>
+  retryCount?: number
   benchmark?: BenchmarkResult
 }
 
@@ -111,30 +113,44 @@ interface TestEachFunction {
   <T extends any[] | [any]>(cases: ReadonlyArray<T>): (
     name: string,
     fn: (...args: T) => Awaitable<void>,
-    timeout?: number,
+    options?: number | TestOptions,
   ) => void
   <T extends ReadonlyArray<any>>(cases: ReadonlyArray<T>): (
     name: string,
     fn: (...args: ExtractEachCallbackArgs<T>) => Awaitable<void>,
-    timeout?: number,
+    options?: number | TestOptions,
   ) => void
   <T>(cases: ReadonlyArray<T>): (
     name: string,
     fn: (...args: T[]) => Awaitable<void>,
-    timeout?: number,
+    options?: number | TestOptions,
   ) => void
 }
 
 type ChainableTestAPI<ExtraContext = {}> = ChainableFunction<
   'concurrent' | 'only' | 'skip' | 'todo' | 'fails',
-  [name: string, fn?: TestFunction<ExtraContext>, timeout?: number],
+  [name: string, fn?: TestFunction<ExtraContext>, options?: number | TestOptions],
   void,
   {
     each: TestEachFunction
-    <T extends ExtraContext>(name: string, fn?: TestFunction<T>, timeout?: number): void
+    <T extends ExtraContext>(name: string, fn?: TestFunction<T>, options?: number | TestOptions): void
   }
 >
 
+export interface TestOptions {
+  /**
+   * Test timeout.
+   */
+  timeout?: number
+  /**
+   * Times to retry the test if fails. Useful for making flaky tests more stable.
+   * When retries is up, the last test error will be thrown.
+   *
+   * @default 1
+   */
+  retry?: number
+}
+
 export type TestAPI<ExtraContext = {}> = ChainableTestAPI<ExtraContext> & {
   each: TestEachFunction
   skipIf(condition: any): ChainableTestAPI<ExtraContext>

test/core/test/retry.test.ts

@@ -0,0 +1,25 @@
+import { expect, it } from 'vitest'
+
+let count1 = 0
+it('retry test', () => {
+  count1 += 1
+  expect(count1).toBe(3)
+}, { retry: 3 })
+
+let count2 = 0
+it.fails('retry test fails', () => {
+  count2 += 1
+  expect(count2).toBe(3)
+}, { retry: 2 })
+
+let count3 = 0
+it('retry test fails', () => {
+  count3 += 1
+  expect(count3).toBe(3)
+}, { retry: 10 })
+
+it('result', () => {
+  expect(count1).toEqual(3)
+  expect(count2).toEqual(2)
+  expect(count3).toEqual(3)
+})