feat(runner): support test.extend (#3554)

Author: fenghan34

docs/api/index.md

@@ -51,6 +51,36 @@ In Jest, `TestFunction` can also be of type `(done: DoneCallback) => void`. If t
   })
   ```
 
+### test.extend
+
+- **Type:** `<T extends Record<string, any>>(fixtures: Fixtures<T>): TestAPI<ExtraContext & T>`
+- **Alias:** `it.extend`
+
+  Use `test.extend` to extend the test context with custom fixtures. This will return a new `test` and it's also extendable, so you can compose more fixtures or override existing ones by extending it as you need. See [Extend Test Context](/guide/test-context.html#test-extend) for more information.
+
+  ```ts
+  import { expect, test } from 'vitest'
+
+  const todos = []
+  const archive = []
+
+  const myTest = test.extend({
+    todos: async (use) => {
+      todos.push(1, 2, 3)
+      await use(todos)
+      todos.length = 0
+    },
+    archive
+  })
+
+  myTest('add item', ({ todos }) => {
+    expect(todos.length).toBe(3)
+
+    todos.push(4)
+    expect(todos.length).toBe(4)
+  })
+  ```
+
 ### test.skip
 
 - **Type:** `(name: string | Function, fn: TestFunction, timeout?: number | TestOptions) => void`

docs/guide/test-context.md

@@ -31,6 +31,118 @@ The `expect` API bound to the current test.
 
 ## Extend Test Context
 
+Vitest provides two diffident ways to help you extend the test context.
+
+### `test.extend`
+
+Like [Playwright](https://playwright.dev/docs/api/class-test#test-extend), you can use this method to define your own `test` API with custom fixtures and reuse it anywhere.
+
+For example, we first create `myTest` with two fixtures, `todos` and `archive`.
+
+```ts
+// my-test.ts
+import { test } from 'vitest'
+
+const todos = []
+const archive = []
+
+export const myTest = test.extend({
+  todos: async (use) => {
+    // setup the fixture before each test function
+    todos.push(1, 2, 3)
+
+    // use the fixture value
+    await use(todos)
+
+    // cleanup the fixture after each test function
+    todos.length = 0
+  },
+  archive
+})
+```
+
+Then we can import and use it.
+
+```ts
+import { expect } from 'vitest'
+import { myTest } from './my-test.ts'
+
+myTest('add items to todos', ({ todos }) => {
+  expect(todos.length).toBe(3)
+
+  todos.add(4)
+  expect(todos.length).toBe(4)
+})
+
+myTest('move items from todos to archive', ({ todos, archive }) => {
+  expect(todos.length).toBe(3)
+  expect(archive.length).toBe(0)
+
+  archive.push(todos.pop())
+  expect(todos.length).toBe(2)
+  expect(archive.length).toBe(1)
+})
+```
+
+We can also add more fixtures or override existing fixtures by extending `myTest`.
+
+```ts
+export const myTest2 = myTest.extend({
+  settings: {
+    // ...
+  }
+})
+```
+
+#### Fixture initialization
+
+Vitest runner will smartly initialize your fixtures and inject them into the test context based on usage.
+
+```ts
+import { test } from 'vitest'
+
+async function todosFn(use) {
+  await use([1, 2, 3])
+}
+
+const myTest = test.extend({
+  todos: todosFn,
+  archive: []
+})
+
+// todosFn will not run
+myTest('', () => {}) // no fixture is available
+myTets('', ({ archive }) => {}) // only archive is available
+
+// todosFn will run
+myTest('', ({ todos }) => {}) // only todos is available
+myTest('', (context) => {}) // both are available
+myTest('', ({ archive, ...rest }) => {}) // both are available
+```
+
+#### TypeScript
+
+To provide fixture types for all your custom contexts, you can pass the fixtures type as a generic.
+
+```ts
+interface MyFixtures {
+  todos: number[]
+  archive: number[]
+}
+
+const myTest = test.extend<MyFixtures>({
+  todos: [],
+  archive: []
+})
+
+myTest('', (context) => {
+  expectTypeOf(context.todos).toEqualTypeOf<number[]>()
+  expectTypeOf(context.archive).toEqualTypeOf<number[]>()
+})
+```
+
+### `beforeEach` and `afterEach`
+
 The contexts are different for each test. You can access and extend them within the `beforeEach` and `afterEach` hooks.
 
 ```ts
@@ -46,7 +158,7 @@ it('should work', ({ foo }) => {
 })
 ```
 
-### TypeScript
+#### TypeScript
 
 To provide property types for all your custom contexts, you can aggregate the `TestContext` type by adding
 
@@ -74,4 +186,4 @@ it<LocalTestContext>('should work', ({ foo }) => {
   // typeof foo is 'string'
   console.log(foo) // 'bar'
 })
-```
+```

packages/runner/src/fixture.ts

@@ -0,0 +1,66 @@
+import type { Fixtures, Test } from './types'
+
+export function withFixtures(fn: Function, fixtures: Fixtures<Record<string, any>>, context: Test<Record<string, any>>['context']) {
+  const props = getUsedFixtureProps(fn, Object.keys(fixtures))
+
+  if (props.length === 0)
+    return () => fn(context)
+
+  let cursor = 0
+
+  async function use(fixtureValue: any) {
+    context[props[cursor++]] = fixtureValue
+
+    if (cursor < props.length)
+      await next()
+    else await fn(context)
+  }
+
+  async function next() {
+    const fixtureValue = fixtures[props[cursor]]
+    typeof fixtureValue === 'function'
+      ? await fixtureValue(use)
+      : await use(fixtureValue)
+  }
+
+  return () => next()
+}
+
+function getUsedFixtureProps(fn: Function, fixtureProps: string[]) {
+  if (!fixtureProps.length || !fn.length)
+    return []
+
+  const paramsStr = fn.toString().match(/[^(]*\(([^)]*)/)![1]
+
+  if (paramsStr[0] === '{' && paramsStr.at(-1) === '}') {
+    // ({...}) => {}
+    const props = paramsStr.slice(1, -1).split(',')
+    const filteredProps = []
+
+    for (const prop of props) {
+      if (!prop)
+        continue
+
+      let _prop = prop.trim()
+
+      if (_prop.startsWith('...')) {
+        // ({ a, b, ...rest }) => {}
+        return fixtureProps
+      }
+
+      const colonIndex = _prop.indexOf(':')
+      if (colonIndex > 0)
+        _prop = _prop.slice(0, colonIndex).trim()
+
+      if (fixtureProps.includes(_prop))
+        filteredProps.push(_prop)
+    }
+
+    // ({}) => {}
+    // ({ a, b, c}) => {}
+    return filteredProps
+  }
+
+  // (ctx) => {}
+  return fixtureProps
+}

packages/runner/src/suite.ts

@@ -1,9 +1,10 @@
 import { format, isObject, noop, objDisplay, objectAttr } from '@vitest/utils'
-import type { File, RunMode, Suite, SuiteAPI, SuiteCollector, SuiteFactory, SuiteHooks, Task, TaskCustom, Test, TestAPI, TestFunction, TestOptions } from './types'
+import type { File, Fixtures, RunMode, Suite, SuiteAPI, SuiteCollector, SuiteFactory, SuiteHooks, Task, TaskCustom, Test, TestAPI, TestFunction, TestOptions } from './types'
 import type { VitestRunner } from './types/runner'
 import { createChainable } from './utils/chain'
 import { collectTask, collectorContext, createTestContext, runWithSuite, withTimeout } from './context'
 import { getHooks, setFn, setHooks } from './map'
+import { withFixtures } from './fixture'
 
 // apis
 export const suite = createSuite()
@@ -95,7 +96,9 @@ function createSuiteCollector(name: string, factory: SuiteFactory = () => { }, m
     })
 
     setFn(test, withTimeout(
-      () => fn(context),
+      this.fixtures
+        ? withFixtures(fn, this.fixtures, context)
+        : () => fn(context),
       options?.timeout ?? runner.config.testTimeout,
     ))
 
@@ -229,12 +232,12 @@ function createSuite() {
 
 function createTest(fn: (
   (
-    this: Record<'concurrent' | 'skip' | 'only' | 'todo' | 'fails' | 'each', boolean | undefined>,
+    this: Record<'concurrent' | 'skip' | 'only' | 'todo' | 'fails' | 'each', boolean | undefined> & { fixtures?: Fixtures<Record<string, any>> },
     title: string,
     fn?: TestFunction,
     options?: number | TestOptions
   ) => void
-)) {
+), context?: Record<string, any>) {
   const testFn = fn as any
 
   testFn.each = function<T>(this: { withContext: () => SuiteAPI; setContext: (key: string, value: boolean | undefined) => SuiteAPI }, cases: ReadonlyArray<T>, ...args: any[]) {
@@ -262,9 +265,20 @@ function createTest(fn: (
   testFn.skipIf = (condition: any) => (condition ? test.skip : test) as TestAPI
   testFn.runIf = (condition: any) => (condition ? test : test.skip) as TestAPI
 
+  testFn.extend = function (fixtures: Fixtures<Record<string, any>>) {
+    const _context = context
+      ? { ...context, fixtures: { ...context.fixtures, ...fixtures } }
+      : { fixtures }
+
+    return createTest(function fn(name: string | Function, fn?: TestFunction, options?: number | TestOptions) {
+      getCurrentSuite().test.fn.call(this, formatName(name), fn, options)
+    }, _context)
+  }
+
   return createChainable(
     ['concurrent', 'skip', 'only', 'todo', 'fails'],
     testFn,
+    context,
   ) as TestAPI
 }
 

packages/runner/src/types/tasks.ts

@@ -183,6 +183,11 @@ export type TestAPI<ExtraContext = {}> = ChainableTestAPI<ExtraContext> & {
   each: TestEachFunction
   skipIf(condition: any): ChainableTestAPI<ExtraContext>
   runIf(condition: any): ChainableTestAPI<ExtraContext>
+  extend<T extends Record<string, any>>(fixtures: Fixtures<T>): TestAPI<ExtraContext & T>
+}
+
+export type Fixtures<T extends Record<string, any>> = {
+  [K in keyof T]: T[K] | ((use: (fixture: T[K]) => Promise<void>) => Promise<void>)
 }
 
 type ChainableSuiteAPI<ExtraContext = {}> = ChainableFunction<

packages/runner/src/utils/chain.ts

@@ -3,20 +3,21 @@ export type ChainableFunction<T extends string, Args extends any[], R = any, E =
 } & {
   [x in T]: ChainableFunction<T, Args, R, E>
 } & {
-  fn: (this: Record<T, boolean | undefined>, ...args: Args) => R
+  fn: (this: Record<T, any>, ...args: Args) => R
 } & E
 
 export function createChainable<T extends string, Args extends any[], R = any, E = {}>(
   keys: T[],
-  fn: (this: Record<T, boolean | undefined>, ...args: Args) => R,
+  fn: (this: Record<T, any>, ...args: Args) => R,
+  initialContext?: Record<T, any>,
 ): ChainableFunction<T, Args, R, E> {
-  function create(context: Record<T, boolean | undefined>) {
+  function create(context: Record<T, any>) {
     const chain = function (this: any, ...args: Args) {
       return fn.apply(context, args)
     }
     Object.assign(chain, fn)
     chain.withContext = () => chain.bind(context)
-    chain.setContext = (key: T, value: boolean | undefined) => {
+    chain.setContext = (key: T, value: any) => {
       context[key] = value
     }
     for (const key of keys) {
@@ -29,7 +30,7 @@ export function createChainable<T extends string, Args extends any[], R = any, E
     return chain
   }
 
-  const chain = create({} as any) as any
+  const chain = create(initialContext || {} as any) as any
   chain.fn = fn
   return chain
 }

test/core/test/test-extend.test.ts

@@ -0,0 +1,143 @@
+/* eslint-disable no-empty-pattern */
+/* eslint-disable prefer-rest-params */
+import { describe, expect, expectTypeOf, test, vi } from 'vitest'
+
+const todoList: number[] = [1, 2, 3]
+const doneList: number[] = []
+const archiveList: number[] = []
+
+const todoFn = vi.fn().mockImplementation(async (use) => {
+  await use(todoList)
+  // cleanup
+  todoFn.mockClear()
+  todoList.length = 0
+  todoList.push(1, 2, 3)
+})
+
+const doneFn = vi.fn().mockImplementation(async (use) => {
+  await use(doneList)
+  // cleanup
+  doneFn.mockClear()
+  doneList.length = 0
+})
+
+const myTest = test
+  .extend<{ todoList: number[] }>({
+    todoList: todoFn,
+  })
+  .extend<{ doneList: number[]; archiveList: number[] }>({
+    doneList: doneFn,
+    archiveList,
+  })
+
+describe('test.extend()', () => {
+  myTest('todoList and doneList', ({ todoList, doneList, archiveList }) => {
+    expect(todoFn).toBeCalledTimes(1)
+    expect(doneFn).toBeCalledTimes(1)
+
+    expectTypeOf(todoList).toEqualTypeOf<number[]>()
+    expectTypeOf(doneList).toEqualTypeOf<number[]>()
+    expectTypeOf(doneList).toEqualTypeOf<number[]>()
+
+    expect(todoList).toEqual([1, 2, 3])
+    expect(doneList).toEqual([])
+    expect(archiveList).toEqual([])
+
+    doneList.push(todoList.shift()!)
+    expect(todoList).toEqual([2, 3])
+    expect(doneList).toEqual([1])
+
+    doneList.push(todoList.shift()!)
+    expect(todoList).toEqual([3])
+    expect(doneList).toEqual([1, 2])
+
+    archiveList.push(todoList.shift()!)
+    expect(todoList).toEqual([])
+    expect(archiveList).toEqual([3])
+  })
+
+  myTest('should called cleanup functions', ({ todoList, doneList, archiveList }) => {
+    expect(todoList).toEqual([1, 2, 3])
+    expect(doneList).toEqual([])
+    expect(archiveList).toEqual([3])
+  })
+
+  describe('smartly init fixtures', () => {
+    myTest('should not init any fixtures', function () {
+      expect(todoFn).not.toBeCalled()
+      expect(doneFn).not.toBeCalled()
+
+      expectTypeOf(arguments[0].todoList).not.toEqualTypeOf<number[]>()
+      expectTypeOf(arguments[0].doneList).not.toEqualTypeOf<number[]>()
+      expectTypeOf(arguments[0].archiveList).not.toEqualTypeOf<number[]>()
+
+      expect(arguments[0].todoList).toBeUndefined()
+      expect(arguments[0].doneList).toBeUndefined()
+      expect(arguments[0].archiveList).toBeUndefined()
+    })
+
+    myTest('should not init any fixtures', function ({}) {
+      expect(todoFn).not.toBeCalled()
+      expect(doneFn).not.toBeCalled()
+
+      expectTypeOf(arguments[0].todoList).not.toEqualTypeOf<number[]>()
+      expectTypeOf(arguments[0].doneList).not.toEqualTypeOf<number[]>()
+      expectTypeOf(arguments[0].archiveList).not.toEqualTypeOf<number[]>()
+
+      expect(arguments[0].todoList).toBeUndefined()
+      expect(arguments[0].doneList).toBeUndefined()
+      expect(arguments[0].archiveList).toBeUndefined()
+    })
+
+    myTest('should only init todoList', function ({ todoList }) {
+      expect(todoFn).toBeCalledTimes(1)
+      expect(doneFn).not.toBeCalled()
+
+      expectTypeOf(todoList).toEqualTypeOf<number[]>()
+      expectTypeOf(arguments[0].doneList).not.toEqualTypeOf<number[]>()
+      expectTypeOf(arguments[0].archiveList).not.toEqualTypeOf<number[]>()
+
+      expect(arguments[0].doneList).toBeUndefined()
+      expect(arguments[0].archiveList).toBeUndefined()
+    })
+
+    myTest('should only init todoList and doneList', function ({ todoList, doneList }) {
+      expect(todoFn).toBeCalledTimes(1)
+      expect(doneFn).toBeCalledTimes(1)
+
+      expectTypeOf(todoList).toEqualTypeOf<number[]>()
+      expectTypeOf(doneList).toEqualTypeOf<number[]>()
+      expectTypeOf(arguments[0].archiveList).not.toEqualTypeOf<number[]>()
+
+      expect(todoList).toEqual([1, 2, 3])
+      expect(doneList).toEqual([])
+      expect(arguments[0].archiveList).toBeUndefined()
+    })
+
+    myTest('should init all fixtures', ({ todoList, ...rest }) => {
+      expect(todoFn).toBeCalledTimes(1)
+      expect(doneFn).toBeCalledTimes(1)
+
+      expectTypeOf(todoList).toEqualTypeOf<number[]>()
+      expectTypeOf(rest.doneList).toEqualTypeOf<number[]>()
+      expectTypeOf(rest.archiveList).toEqualTypeOf<number[]>()
+
+      expect(todoList).toEqual([1, 2, 3])
+      expect(rest.doneList).toEqual([])
+      expect(rest.archiveList).toEqual([3])
+    })
+
+    myTest('should init all fixtures', (context) => {
+      expect(todoFn).toBeCalledTimes(1)
+      expect(doneFn).toBeCalledTimes(1)
+
+      expectTypeOf(context.todoList).toEqualTypeOf<number[]>()
+      expectTypeOf(context.doneList).toEqualTypeOf<number[]>()
+      expectTypeOf(context.archiveList).toEqualTypeOf<number[]>()
+
+      expect(context.todoList).toEqual([1, 2, 3])
+      expect(context.doneList).toEqual([])
+      expect(context.archiveList).toEqual([3])
+    })
+  })
+})