feat: test context (#900)

* feat: support concurrent snapshot
* feat: cleanup callback for `beforeAll` and `beforeEach`

docs/api/index.md

@@ -3,14 +3,15 @@
 The following types are used in the type signatures below
 
 ```ts
-type DoneCallback = (error?: any) => void
 type Awaitable<T> = T | PromiseLike<T>
-type TestFunction = () => Awaitable<void> | ((done: DoneCallback) => void)
+type TestFunction = () => Awaitable<void>
 ```
 
 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.
 
-For compatibility with Jest, `TestFunction` can also be of type `(done: DoneCallback) => void`. If this form is used, the test will not be concluded until `done` is called (with zero arguments or a falsy value for a successful test, and with a truthy error value as argument to trigger a fail). We don't recommend using this form, as you can achieve the same using an `async` function.
+::: tip 
+In Jest, `TestFunction` can also be of type `(done: DoneCallback) => void`. If this form is used, the test will not be concluded until `done` is called. You can achieve the same using an `async` function, see the [Migration guide Done Callback section](../guide/migration#done-callback).
+:::
 
 ## test
 
@@ -90,6 +91,17 @@ For compatibility with Jest, `TestFunction` can also be of type `(done: DoneCall
   test.todo.concurrent(/* ... */) // or test.concurrent.todo(/* ... */)
   ```
 
+  When using Snapshots with async concurrent tests, due to the limitation of JavaScript, you need to use the `expect` from the [Test Context](/guide/test-context.md) to ensure the right test is being detected.
+
+  ```ts
+  test.concurrent('test 1', async ({ expect }) => {
+    expect(foo).toMatchSnapshot()
+  })
+  test.concurrent('test 2', async ({ expect }) => {
+    expect(foo).toMatchSnapshot()
+  })
+  ```
+
 ### test.todo
 
 - **Type:** `(name: string) => void`
@@ -1321,6 +1333,22 @@ These functions allow you to hook into the life cycle of tests to avoid repeatin
 
   Here, the `beforeEach` ensures that user is added for each test.
 
+  Since Vitest v0.10.0, `beforeEach` also accepts an optional cleanup function (equivalent to `afterEach`).
+
+  ```ts
+  import { beforeEach } from 'vitest'
+
+  beforeEach(async () => {
+    // called once before all tests run
+    await prepareSomething()
+  
+    // clean up function, called once after all tests run
+    return async () => {
+      await resetSomething()
+    }
+  })
+  ```
+
 ### afterEach
 
 - **Type:** `afterEach(fn: () => Awaitable<void>, timeout?: number)`
@@ -1356,7 +1384,23 @@ These functions allow you to hook into the life cycle of tests to avoid repeatin
   })
   ```
 
-  Here the `beforeAll` ensures that the mock data is set up before tests run
+  Here the `beforeAll` ensures that the mock data is set up before tests run.
+
+  Since Vitest v0.10.0, `beforeAll` also accepts an optional cleanup function (equivalent to `afterAll`).
+
+  ```ts
+  import { beforeAll } from 'vitest'
+
+  beforeAll(async () => {
+    // called once before all tests run
+    await startMocking()
+  
+    // clean up function, called once after all tests run
+    return async () => {
+      await stopMocking()
+    }
+  })
+  ```
 
 ### afterAll
 

docs/guide/migration.md

@@ -25,3 +25,16 @@ Jest exports various [`jasmine`](https://jasmine.github.io/) globals (such as `j
 **Envs**
 
 Just like Jest, Vitest sets `NODE_ENV` to `test`, if it wasn't set before. Vitest also has a counterpart for `JEST_WORKER_ID` called `VITEST_WORKER_ID`, so if you rely on it, don't forget to rename it.
+
+**Done Callback**
+
+From Vitest v0.10.0, the callback style of declaring tests is deprecated. You can rewrite them to use `async`/`await` functions, or use Promise to mimic the callback style.
+
+```diff
+- it('should work', (done) => {
++ it('should work', () => new Promise(done => {
+    // ...
+    done()
+- })
++ }))
+```

docs/guide/test-context.md

@@ -0,0 +1,56 @@
+# Test Context
+
+Inspired by [Playwright Fixtures](https://playwright.dev/docs/test-fixtures), Vitest's test context allows you to define utils, states, and fixtures that can be used in your tests.
+
+## Usage
+
+The first argument or each test callback is a test context.
+
+```ts
+import { it } from 'vitest'
+
+it('should work', (ctx) => {
+  // prints name of the test
+  console.log(ctx.meta.name)
+})
+```
+
+## Built-in Test Context
+
+#### `context.meta`
+
+A readonly object containing metadata about the test.
+
+#### `context.expect`
+
+The `expect` API which bound to the current test.
+
+## Extend Test Context
+
+The contexts are different for each test. You can access and extend them within the `beforeEach` and `afterEach` hooks.
+
+```ts
+import { beforeEach, it } from 'vitest'
+
+beforeEach(async (context) => {
+  // extend context
+  context.foo = 'bar'
+})
+
+it('should work', ({ foo }) => {
+  console.log(foo) // 'bar'
+})
+```
+
+### TypeScript
+
+To provide type for your custom context properties, you can aggregate the type `TestContext` by adding
+
+```ts
+declare module 'vitest' {
+  export interface TestContext {
+    foo?: string
+  }
+}
+```
+

packages/vitest/src/index.ts

@@ -4,10 +4,11 @@ import type { Constructable, InlineConfig } from './types'
 
 export { suite, test, describe, it } from './runtime/suite'
 export * from './runtime/hooks'
+export * from './runtime/utils'
 
 export { runOnce, isFirstRun } from './integrations/run-once'
-export * from './integrations/chai'
 export type { EnhancedSpy, MockedFunction, MockedObject, SpyInstance, SpyInstanceFn, SpyContext } from './integrations/spy'
+export * from './integrations/chai'
 export * from './integrations/vi'
 export * from './integrations/utils'
 

packages/vitest/src/integrations/chai/index.ts

@@ -1,19 +1,31 @@
 import chai from 'chai'
 import './setup'
+import type { Test } from '../../types'
 import { getState, setState } from './jest-expect'
 
-const expect = ((value: any, message?: string): Vi.Assertion => {
-  const { assertionCalls } = getState()
-  setState({ assertionCalls: assertionCalls + 1 })
-  return chai.expect(value, message) as unknown as Vi.Assertion
-}) as Vi.ExpectStatic
+export function createExpect(test?: Test) {
+  const expect = ((value: any, message?: string): Vi.Assertion => {
+    const { assertionCalls } = getState()
+    setState({ assertionCalls: assertionCalls + 1 })
+    const assert = chai.expect(value, message) as unknown as Vi.Assertion
+    if (test)
+      // @ts-expect-error internal
+      return assert.withTest(test) as Vi.Assertion
+    else
+      return assert
+  }) as Vi.ExpectStatic
+  Object.assign(expect, chai.expect)
 
-Object.assign(expect, chai.expect)
+  expect.getState = getState
+  expect.setState = setState
 
-expect.getState = getState
-expect.setState = setState
-// @ts-expect-error untyped
-expect.extend = matchers => chai.expect.extend(expect, matchers)
+  // @ts-expect-error untyped
+  expect.extend = matchers => chai.expect.extend(expect, matchers)
+
+  return expect
+}
+
+const expect = createExpect()
 
 export { assert, should } from 'chai'
 export { chai, expect }

packages/vitest/src/integrations/chai/jest-expect.ts

@@ -2,7 +2,7 @@ import c from 'picocolors'
 import type { EnhancedSpy } from '../spy'
 import { isMockFunction } from '../spy'
 import { addSerializer } from '../snapshot/port/plugins'
-import type { Constructable } from '../../types'
+import type { Constructable, Test } from '../../types'
 import { assertTypes } from '../../utils'
 import { unifiedDiff } from '../../node/diff'
 import type { ChaiPlugin, MatcherState } from './types'
@@ -65,6 +65,12 @@ export const JestChaiExpect: ChaiPlugin = (chai, utils) => {
     })
   })
 
+  // @ts-expect-error @internal
+  def('withTest', function (test: Test) {
+    utils.flag(this, 'vitest-test', test)
+    return this
+  })
+
   def('toEqual', function (expected) {
     const actual = utils.flag(this, 'object')
     const equal = jestEquals(

packages/vitest/src/integrations/snapshot/chai.ts

@@ -31,11 +31,12 @@ export const SnapshotPlugin: ChaiPlugin = (chai, utils) => {
       key,
       function (this: Record<string, unknown>, properties?: object, message?: string) {
         const expected = utils.flag(this, 'object')
+        const test = utils.flag(this, 'vitest-test')
         if (typeof properties === 'string' && typeof message === 'undefined') {
           message = properties
           properties = undefined
         }
-        getSnapshotClient().assert(expected, message, false, properties)
+        getSnapshotClient().assert(expected, test, message, false, properties)
       },
     )
   }
@@ -45,6 +46,7 @@ export const SnapshotPlugin: ChaiPlugin = (chai, utils) => {
     function __VITEST_INLINE_SNAPSHOT__(this: Record<string, unknown>, properties?: object, inlineSnapshot?: string, message?: string) {
       const expected = utils.flag(this, 'object')
       const error = utils.flag(this, 'error')
+      const test = utils.flag(this, 'vitest-test')
       if (typeof properties === 'string') {
         message = inlineSnapshot
         inlineSnapshot = properties
@@ -52,7 +54,7 @@ export const SnapshotPlugin: ChaiPlugin = (chai, utils) => {
       }
       if (inlineSnapshot)
         inlineSnapshot = stripSnapshotIndentation(inlineSnapshot)
-      getSnapshotClient().assert(expected, message, true, properties, inlineSnapshot, error)
+      getSnapshotClient().assert(expected, test, message, true, properties, inlineSnapshot, error)
     },
   )
   utils.addMethod(
@@ -60,7 +62,8 @@ export const SnapshotPlugin: ChaiPlugin = (chai, utils) => {
     'toThrowErrorMatchingSnapshot',
     function (this: Record<string, unknown>, message?: string) {
       const expected = utils.flag(this, 'object')
-      getSnapshotClient().assert(getErrorString(expected), message)
+      const test = utils.flag(this, 'vitest-test')
+      getSnapshotClient().assert(getErrorString(expected), test, message)
     },
   )
   utils.addMethod(
@@ -69,7 +72,8 @@ export const SnapshotPlugin: ChaiPlugin = (chai, utils) => {
     function __VITEST_INLINE_SNAPSHOT__(this: Record<string, unknown>, inlineSnapshot: string, message: string) {
       const expected = utils.flag(this, 'object')
       const error = utils.flag(this, 'error')
-      getSnapshotClient().assert(getErrorString(expected), message, true, undefined, inlineSnapshot, error)
+      const test = utils.flag(this, 'vitest-test')
+      getSnapshotClient().assert(getErrorString(expected), test, message, true, undefined, inlineSnapshot, error)
     },
   )
 }

packages/vitest/src/integrations/snapshot/client.ts

@@ -1,5 +1,5 @@
 import { expect } from 'chai'
-import type { SnapshotResult, Test } from '../../types'
+import type { Test } from '../../types'
 import { rpc } from '../../runtime/rpc'
 import { getNames, getWorkerState } from '../../utils'
 import { equals, iterableEquality, subsetEquality } from '../chai/jest-utils'
@@ -14,30 +14,40 @@ export interface Context {
 
 export class SnapshotClient {
   test: Test | undefined
-  testFile = ''
   snapshotState: SnapshotState | undefined
+  snapshotStateMap = new Map<string, SnapshotState>()
 
   async setTest(test: Test) {
     this.test = test
 
-    if (this.testFile !== this.test.file!.filepath) {
-      if (this.snapshotState)
-        this.saveSnap()
+    if (this.snapshotState?.testFilePath !== this.test.file!.filepath) {
+      this.saveCurrent()
 
-      this.testFile = this.test!.file!.filepath
-      this.snapshotState = new SnapshotState(
-        await rpc().resolveSnapshotPath(this.testFile),
-        getWorkerState().config.snapshotOptions,
-      )
+      const filePath = this.test!.file!.filepath
+      if (!this.getSnapshotState(test)) {
+        this.snapshotStateMap.set(
+          filePath,
+          new SnapshotState(
+            filePath,
+            await rpc().resolveSnapshotPath(filePath),
+            getWorkerState().config.snapshotOptions,
+          ),
+        )
+      }
+      this.snapshotState = this.getSnapshotState(test)
     }
   }
 
+  getSnapshotState(test: Test) {
+    return this.snapshotStateMap.get(test.file!.filepath)!
+  }
+
   clearTest() {
     this.test = undefined
   }
 
-  assert(received: unknown, message?: string, isInline = false, properties?: object, inlineSnapshot?: string, error?: Error): void {
-    if (!this.test)
+  assert(received: unknown, test = this.test, message?: string, isInline = false, properties?: object, inlineSnapshot?: string, error?: Error): void {
+    if (!test)
       throw new Error('Snapshot cannot be used outside of test')
 
     if (typeof properties === 'object') {
@@ -58,11 +68,13 @@ export class SnapshotClient {
     }
 
     const testName = [
-      ...getNames(this.test).slice(1),
+      ...getNames(test).slice(1),
       ...(message ? [message] : []),
     ].join(' > ')
 
-    const { actual, expected, key, pass } = this.snapshotState!.match({
+    const snapshotState = this.getSnapshotState(test)
+
+    const { actual, expected, key, pass } = snapshotState.match({
       testName,
       received,
       isInline,
@@ -81,41 +93,16 @@ export class SnapshotClient {
     }
   }
 
-  async saveSnap() {
-    if (!this.testFile || !this.snapshotState)
+  async saveCurrent() {
+    if (!this.snapshotState)
       return
-    const result = await packSnapshotState(this.testFile, this.snapshotState)
+    const result = await this.snapshotState.pack()
     await rpc().snapshotSaved(result)
 
-    this.testFile = ''
     this.snapshotState = undefined
   }
-}
 
-export async function packSnapshotState(filepath: string, state: SnapshotState): Promise<SnapshotResult> {
-  const snapshot: SnapshotResult = {
-    filepath,
-    added: 0,
-    fileDeleted: false,
-    matched: 0,
-    unchecked: 0,
-    uncheckedKeys: [],
-    unmatched: 0,
-    updated: 0,
+  clear() {
+    this.snapshotStateMap.clear()
   }
-  const uncheckedCount = state.getUncheckedCount()
-  const uncheckedKeys = state.getUncheckedKeys()
-  if (uncheckedCount)
-    state.removeUncheckedKeys()
-
-  const status = await state.save()
-  snapshot.fileDeleted = status.deleted
-  snapshot.added = state.added
-  snapshot.matched = state.matched
-  snapshot.unmatched = state.unmatched
-  snapshot.updated = state.updated
-  snapshot.unchecked = !status.deleted ? uncheckedCount : 0
-  snapshot.uncheckedKeys = Array.from(uncheckedKeys)
-
-  return snapshot
 }

packages/vitest/src/integrations/snapshot/port/state.ts

@@ -9,7 +9,7 @@ import fs from 'fs'
 import type { Config } from '@jest/types'
 // import { getStackTraceLines, getTopFrame } from 'jest-message-util'
 import type { OptionsReceived as PrettyFormatOptions } from 'pretty-format'
-import type { ParsedStack, SnapshotData, SnapshotMatchOptions, SnapshotStateOptions } from '../../../types'
+import type { ParsedStack, SnapshotData, SnapshotMatchOptions, SnapshotResult, SnapshotStateOptions } from '../../../types'
 import { slash } from '../../../utils'
 import { parseStacktrace } from '../../../utils/source-map'
 import type { InlineSnapshot } from './inlineSnapshot'
@@ -45,7 +45,6 @@ export default class SnapshotState {
   private _updateSnapshot: Config.SnapshotUpdateState
   private _snapshotData: SnapshotData
   private _initialData: SnapshotData
-  private _snapshotPath: string
   private _inlineSnapshots: Array<InlineSnapshot>
   private _uncheckedKeys: Set<string>
   private _snapshotFormat: PrettyFormatOptions
@@ -56,10 +55,13 @@ export default class SnapshotState {
   unmatched: number
   updated: number
 
-  constructor(snapshotPath: string, options: SnapshotStateOptions) {
-    this._snapshotPath = snapshotPath
+  constructor(
+    public testFilePath: string,
+    public snapshotPath: string,
+    options: SnapshotStateOptions,
+  ) {
     const { data, dirty } = getSnapshotData(
-      this._snapshotPath,
+      this.snapshotPath,
       options.updateSnapshot,
     )
     this._initialData = data
@@ -133,6 +135,7 @@ export default class SnapshotState {
     this.matched = 0
     this.unmatched = 0
     this.updated = 0
+    this._dirty = false
   }
 
   async save(): Promise<SaveStatus> {
@@ -147,15 +150,15 @@ export default class SnapshotState {
 
     if ((this._dirty || this._uncheckedKeys.size) && !isEmpty) {
       if (hasExternalSnapshots)
-        await saveSnapshotFile(this._snapshotData, this._snapshotPath)
+        await saveSnapshotFile(this._snapshotData, this.snapshotPath)
       if (hasInlineSnapshots)
         await saveInlineSnapshots(this._inlineSnapshots)
 
       status.saved = true
     }
-    else if (!hasExternalSnapshots && fs.existsSync(this._snapshotPath)) {
+    else if (!hasExternalSnapshots && fs.existsSync(this.snapshotPath)) {
       if (this._updateSnapshot === 'all')
-        fs.unlinkSync(this._snapshotPath)
+        fs.unlinkSync(this.snapshotPath)
 
       status.deleted = true
     }
@@ -204,7 +207,7 @@ export default class SnapshotState {
     const expectedTrimmed = prepareExpected(expected)
     const pass = expectedTrimmed === prepareExpected(receivedSerialized)
     const hasSnapshot = expected !== undefined
-    const snapshotIsPersisted = isInline || fs.existsSync(this._snapshotPath)
+    const snapshotIsPersisted = isInline || fs.existsSync(this.snapshotPath)
 
     if (pass && !isInline) {
       // Executing a snapshot file as JavaScript and writing the strings back
@@ -280,4 +283,32 @@ export default class SnapshotState {
       }
     }
   }
+
+  async pack(): Promise<SnapshotResult> {
+    const snapshot: SnapshotResult = {
+      filepath: this.testFilePath,
+      added: 0,
+      fileDeleted: false,
+      matched: 0,
+      unchecked: 0,
+      uncheckedKeys: [],
+      unmatched: 0,
+      updated: 0,
+    }
+    const uncheckedCount = this.getUncheckedCount()
+    const uncheckedKeys = this.getUncheckedKeys()
+    if (uncheckedCount)
+      this.removeUncheckedKeys()
+
+    const status = await this.save()
+    snapshot.fileDeleted = status.deleted
+    snapshot.added = this.added
+    snapshot.matched = this.matched
+    snapshot.unmatched = this.unmatched
+    snapshot.updated = this.updated
+    snapshot.unchecked = !status.deleted ? uncheckedCount : 0
+    snapshot.uncheckedKeys = Array.from(uncheckedKeys)
+
+    return snapshot
+  }
 }

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

@@ -78,6 +78,7 @@ export async function startVitest(cliFilters: string[], options: CliOptions, vit
     ctx.error(`\n${c.red(divider(c.bold(c.inverse(' Unhandled Error '))))}`)
     await ctx.printError(e)
     ctx.error('\n\n')
+    return false
   }
 
   if (!ctx.config.watch) {

packages/vitest/src/runtime/collect.ts

@@ -1,10 +1,10 @@
 import { createHash } from 'crypto'
 import { relative } from 'pathe'
 import type { File, ResolvedConfig, Suite, TaskBase } from '../types'
-import { clearContext, defaultSuite } from './suite'
+import { clearCollectorContext, defaultSuite } from './suite'
 import { getHooks, setHooks } from './map'
 import { processError } from './error'
-import { context } from './context'
+import { collectorContext } from './context'
 import { runSetupFiles } from './setup'
 
 const now = Date.now
@@ -30,7 +30,7 @@ export async function collectTests(paths: string[], config: ResolvedConfig) {
       tasks: [],
     }
 
-    clearContext()
+    clearCollectorContext()
     try {
       await runSetupFiles(config)
       await import(filepath)
@@ -39,7 +39,7 @@ export async function collectTests(paths: string[], config: ResolvedConfig) {
 
       setHooks(file, getHooks(defaultTasks))
 
-      for (const c of [...defaultTasks.tasks, ...context.tasks]) {
+      for (const c of [...defaultTasks.tasks, ...collectorContext.tasks]) {
         if (c.type === 'test') {
           file.tasks.push(c)
         }

packages/vitest/src/runtime/context.ts

@@ -1,20 +1,21 @@
-import type { Awaitable, DoneCallback, RuntimeContext, SuiteCollector, TestFunction } from '../types'
+import type { Awaitable, RuntimeContext, SuiteCollector, Test, TestContext } from '../types'
+import { createExpect } from '../integrations/chai'
 import { clearTimeout, getWorkerState, setTimeout } from '../utils'
 
-export const context: RuntimeContext = {
+export const collectorContext: RuntimeContext = {
   tasks: [],
   currentSuite: null,
 }
 
 export function collectTask(task: SuiteCollector) {
-  context.currentSuite?.tasks.push(task)
+  collectorContext.currentSuite?.tasks.push(task)
 }
 
 export async function runWithSuite(suite: SuiteCollector, fn: (() => Awaitable<void>)) {
-  const prev = context.currentSuite
-  context.currentSuite = suite
+  const prev = collectorContext.currentSuite
+  collectorContext.currentSuite = suite
   await fn()
-  context.currentSuite = prev
+  collectorContext.currentSuite = prev
 }
 
 export function getDefaultTestTimeout() {
@@ -44,22 +45,25 @@ export function withTimeout<T extends((...args: any[]) => any)>(
   }) as T
 }
 
+export function createTestContext(test: Test): TestContext {
+  const context = function () {
+    throw new Error('done() callback is deperated, use promise instead')
+  } as unknown as TestContext
+
+  context.meta = test
+
+  let _expect: Vi.ExpectStatic | undefined
+  Object.defineProperty(context, 'expect', {
+    get() {
+      if (!_expect)
+        _expect = createExpect(test)
+      return _expect
+    },
+  })
+
+  return context
+}
+
 function makeTimeoutMsg(isHook: boolean, timeout: number) {
   return `${isHook ? 'Hook' : 'Test'} timed out in ${timeout}ms.\nIf this is a long-running test, pass a timeout value as the last argument or configure it globally with "${isHook ? 'hookTimeout' : 'testTimeout'}".`
 }
-
-function ensureAsyncTest(fn: TestFunction): () => Awaitable<void> {
-  if (!fn.length)
-    return fn as () => Awaitable<void>
-
-  return () => new Promise((resolve, reject) => {
-    const done: DoneCallback = (...args: any[]) => args[0] // reject on truthy values
-      ? reject(args[0])
-      : resolve()
-    fn(done)
-  })
-}
-
-export function normalizeTest(fn: TestFunction, timeout?: number): () => Awaitable<void> {
-  return withTimeout(ensureAsyncTest(fn), timeout)
-}

packages/vitest/src/runtime/run.ts

@@ -1,4 +1,4 @@
-import type { File, HookListener, ResolvedConfig, Suite, SuiteHooks, Task, TaskResult, TaskState, Test } from '../types'
+import type { File, HookCleanupCallback, HookListener, ResolvedConfig, Suite, SuiteHooks, Task, TaskResult, TaskState, Test } from '../types'
 import { vi } from '../integrations/vi'
 import { getSnapshotClient } from '../integrations/snapshot/chai'
 import { clearTimeout, getFullName, getWorkerState, hasFailed, hasTests, partitionSuiteChildren, setTimeout } from '../utils'
@@ -23,16 +23,32 @@ function updateSuiteHookState(suite: Task, name: keyof SuiteHooks, state: TaskSt
   }
 }
 
-export async function callSuiteHook<T extends keyof SuiteHooks>(suite: Suite, currentTask: Task, name: T, args: SuiteHooks[T][0] extends HookListener<infer A> ? A : never) {
-  if (name === 'beforeEach' && suite.suite)
-    await callSuiteHook(suite.suite, currentTask, name, args)
+export async function callSuiteHook<T extends keyof SuiteHooks>(
+  suite: Suite,
+  currentTask: Task,
+  name: T,
+  args: SuiteHooks[T][0] extends HookListener<infer A, any> ? A : never,
+): Promise<HookCleanupCallback[]> {
+  const callbacks: HookCleanupCallback[] = []
+  if (name === 'beforeEach' && suite.suite) {
+    callbacks.push(
+      ...await callSuiteHook(suite.suite, currentTask, name, args),
+    )
+  }
 
   updateSuiteHookState(currentTask, name, 'run')
-  await Promise.all(getHooks(suite)[name].map(fn => fn(...(args as any))))
+  callbacks.push(
+    ...await Promise.all(getHooks(suite)[name].map(fn => fn(...(args as any)))),
+  )
   updateSuiteHookState(currentTask, name, 'pass')
 
-  if (name === 'afterEach' && suite.suite)
-    await callSuiteHook(suite.suite, currentTask, name, args)
+  if (name === 'afterEach' && suite.suite) {
+    callbacks.push(
+      ...await callSuiteHook(suite.suite, currentTask, name, args),
+    )
+  }
+
+  return callbacks
 }
 
 const packs = new Map<string, TaskResult|undefined>()
@@ -84,8 +100,9 @@ export async function runTest(test: Test) {
 
   workerState.current = test
 
+  let beforeEachCleanups: HookCleanupCallback[] = []
   try {
-    await callSuiteHook(test.suite, test, 'beforeEach', [test, test.suite])
+    beforeEachCleanups = await callSuiteHook(test.suite, test, 'beforeEach', [test.context, test.suite])
     setState({
       assertionCalls: 0,
       isExpectingAssertions: false,
@@ -110,7 +127,8 @@ export async function runTest(test: Test) {
   }
 
   try {
-    await callSuiteHook(test.suite, test, 'afterEach', [test, test.suite])
+    await callSuiteHook(test.suite, test, 'afterEach', [test.context, test.suite])
+    await Promise.all(beforeEachCleanups.map(i => i?.()))
   }
   catch (e) {
     test.result.state = 'fail'
@@ -172,7 +190,7 @@ export async function runSuite(suite: Suite) {
   }
   else {
     try {
-      await callSuiteHook(suite, suite, 'beforeAll', [suite])
+      const beforeAllCleanups = await callSuiteHook(suite, suite, 'beforeAll', [suite])
 
       for (const tasksGroup of partitionSuiteChildren(suite)) {
         if (tasksGroup[0].concurrent === true) {
@@ -183,7 +201,9 @@ export async function runSuite(suite: Suite) {
             await runSuiteChild(c)
         }
       }
+
       await callSuiteHook(suite, suite, 'afterAll', [suite])
+      await Promise.all(beforeAllCleanups.map(i => i?.()))
     }
     catch (e) {
       suite.result.state = 'fail'
@@ -234,12 +254,13 @@ export async function startTests(paths: string[], config: ResolvedConfig) {
   const files = await collectTests(paths, config)
 
   rpc().onCollected(files)
+  getSnapshotClient().clear()
 
   await runFiles(files, config)
 
   takeCoverage()
 
-  await getSnapshotClient().saveSnap()
+  await getSnapshotClient().saveCurrent()
 
   await sendTasksUpdate()
 }

packages/vitest/src/runtime/suite.ts

@@ -2,7 +2,7 @@ import { format } from 'util'
 import type { File, RunMode, Suite, SuiteAPI, SuiteCollector, SuiteFactory, SuiteHooks, Task, Test, TestAPI, TestFunction } from '../types'
 import { isObject, noop, toArray } from '../utils'
 import { createChainable } from './chain'
-import { collectTask, context, normalizeTest, runWithSuite } from './context'
+import { collectTask, collectorContext, createTestContext, runWithSuite, withTimeout } from './context'
 import { getHooks, setFn, setHooks } from './map'
 
 // apis
@@ -40,14 +40,14 @@ export const it = test
 // implementations
 export const defaultSuite = suite('')
 
-export function clearContext() {
-  context.tasks.length = 0
+export function clearCollectorContext() {
+  collectorContext.tasks.length = 0
   defaultSuite.clear()
-  context.currentSuite = defaultSuite
+  collectorContext.currentSuite = defaultSuite
 }
 
 export function getCurrentSuite() {
-  return context.currentSuite || defaultSuite
+  return collectorContext.currentSuite || defaultSuite
 }
 
 export function createSuiteHooks() {
@@ -67,7 +67,7 @@ function createSuiteCollector(name: string, factory: SuiteFactory = () => { }, m
 
   initSuite()
 
-  const test = createTest(function (name: string, fn?: TestFunction, timeout?: number) {
+  const test = createTest(function (name: string, fn = noop, timeout?: number) {
     const mode = this.only ? 'only' : this.skip ? 'skip' : this.todo ? 'todo' : 'run'
 
     const test: Test = {
@@ -77,10 +77,22 @@ function createSuiteCollector(name: string, factory: SuiteFactory = () => { }, m
       mode,
       suite: undefined!,
       fails: this.fails,
-    }
+    } as Omit<Test, 'context'> as Test
     if (this.concurrent || concurrent)
       test.concurrent = true
-    setFn(test, normalizeTest(fn || noop, timeout))
+
+    const context = createTestContext(test)
+    // create test context
+    Object.defineProperty(test, 'context', {
+      value: context,
+      enumerable: false,
+    })
+
+    setFn(test, withTimeout(
+      () => fn(context),
+      timeout,
+    ))
+
     tasks.push(test)
   })
 

packages/vitest/src/runtime/utils.ts

@@ -0,0 +1,15 @@
+import type { DoneCallback } from '../types'
+
+/**
+ * A simple wrapper for converting callback style to promise
+ */
+export function withCallback(fn: (done: DoneCallback) => void): Promise<void> {
+  return new Promise((resolve, reject) =>
+    fn((err) => {
+      if (err)
+        reject(err)
+      else
+        resolve()
+    }),
+  )
+}

packages/vitest/src/types/tasks.ts

@@ -4,6 +4,7 @@ import type { UserConsoleLog } from '.'
 
 export type RunMode = 'run' | 'skip' | 'only' | 'todo'
 export type TaskState = RunMode | 'pass' | 'fail'
+
 export interface TaskBase {
   id: string
   name: string
@@ -36,17 +37,18 @@ export interface File extends Suite {
   collectDuration?: number
 }
 
-export interface Test extends TaskBase {
+export interface Test<ExtraContext = {}> extends TaskBase {
   type: 'test'
   suite: Suite
   result?: TaskResult
   fails?: boolean
+  context: TestContext & ExtraContext
 }
 
 export type Task = Test | Suite | File
 
 export type DoneCallback = (error?: any) => void
-export type TestFunction = (done: DoneCallback) => Awaitable<void>
+export type TestFunction<ExtraContext = {}> = (context: TestContext & ExtraContext) => Awaitable<void>
 
 // jest's ExtractEachCallbackArgs
 type ExtractEachCallbackArgs<T extends ReadonlyArray<any>> = {
@@ -98,33 +100,35 @@ interface EachFunction {
   ) => void
 }
 
-export type TestAPI = ChainableFunction<
+export type TestAPI<ExtraContext = {}> = ChainableFunction<
 'concurrent' | 'only' | 'skip' | 'todo' | 'fails',
-[name: string, fn?: TestFunction, timeout?: number],
+[name: string, fn?: TestFunction<ExtraContext>, timeout?: number],
 void
 > & { each: EachFunction }
 
-export type SuiteAPI = ChainableFunction<
+export type SuiteAPI<ExtraContext = {}> = ChainableFunction<
 'concurrent' | 'only' | 'skip' | 'todo',
 [name: string, factory?: SuiteFactory],
-SuiteCollector
+SuiteCollector<ExtraContext>
 > & { each: EachFunction }
 
-export type HookListener<T extends any[]> = (...args: T) => Awaitable<void>
+export type HookListener<T extends any[], Return = void> = (...args: T) => Awaitable<Return | void>
 
 export interface SuiteHooks {
-  beforeAll: HookListener<[Suite]>[]
+  beforeAll: HookListener<[Suite], () => Awaitable<void>>[]
   afterAll: HookListener<[Suite]>[]
-  beforeEach: HookListener<[Test, Suite]>[]
-  afterEach: HookListener<[Test, Suite]>[]
+  beforeEach: HookListener<[TestContext, Suite], () => Awaitable<void>>[]
+  afterEach: HookListener<[TestContext, Suite]>[]
 }
 
-export interface SuiteCollector {
+export type HookCleanupCallback = (() => Awaitable<void>) | void
+
+export interface SuiteCollector<ExtraContext = {}> {
   readonly name: string
   readonly mode: RunMode
   type: 'collector'
-  test: TestAPI
-  tasks: (Suite | Test | SuiteCollector)[]
+  test: TestAPI<ExtraContext>
+  tasks: (Suite | Test | SuiteCollector<ExtraContext>)[]
   collect: (file?: File) => Promise<Suite>
   clear: () => void
   on: <T extends keyof SuiteHooks>(name: T, ...fn: SuiteHooks[T]) => void
@@ -136,3 +140,20 @@ export interface RuntimeContext {
   tasks: (SuiteCollector | Test)[]
   currentSuite: SuiteCollector | null
 }
+
+export interface TestContext {
+  /**
+   * @deprecated Use promise instead
+   */
+  (error?: any): void
+
+  /**
+   * Metadata of the current test
+   */
+  meta: Readonly<Test>
+
+  /**
+   * A expect instance bound to the test
+   */
+  expect: Vi.ExpectStatic
+}

test/core/test/basic.test.ts

@@ -54,22 +54,6 @@ test.skip('async with timeout', async () => {
 
 it('timeout', () => new Promise(resolve => setTimeout(resolve, timeout)))
 
-function callbackTest(name: string, doneValue: any) {
-  let callbackAwaited = false
-
-  it(`callback setup ${name}`, (done) => {
-    setTimeout(() => {
-      expect({}).toBeTruthy()
-      callbackAwaited = true
-      done(doneValue)
-    }, 20)
-  })
-
-  it(`callback test ${name}`, () => {
-    expect(callbackAwaited).toBe(true)
-  })
-}
-
-callbackTest('success ', undefined)
-
-callbackTest('success done(false)', false)
+it.fails('deprecated done callback', (done) => {
+  done()
+})

test/core/test/hooks.test.ts

@@ -53,3 +53,31 @@ suite('level1', () => {
     expect(count).toBe(2)
   })
 })
+
+suite('hooks cleanup', () => {
+  let cleanUpCount = 0
+  suite('run', () => {
+    beforeAll(() => {
+      cleanUpCount += 10
+      return () => {
+        cleanUpCount -= 10
+      }
+    })
+    beforeEach(() => {
+      cleanUpCount += 1
+      return () => {
+        cleanUpCount -= 1
+      }
+    })
+
+    it('one', () => {
+      expect(cleanUpCount).toBe(11)
+    })
+    it('two', () => {
+      expect(cleanUpCount).toBe(11)
+    })
+  })
+  it('end', () => {
+    expect(cleanUpCount).toBe(0)
+  })
+})

test/reporters/src/data.ts

@@ -66,6 +66,7 @@ const innerTasks: Task[] = [
       error,
       duration: 1.4422860145568848,
     },
+    context: null as any,
   },
 ]
 
@@ -82,6 +83,7 @@ const tasks: Task[] = [
     fails: undefined,
     file,
     result: { state: 'pass', duration: 1.0237109661102295 },
+    context: null as any,
   },
   {
     id: '1223128da3_3',
@@ -92,6 +94,7 @@ const tasks: Task[] = [
     fails: undefined,
     file,
     result: undefined,
+    context: null as any,
   },
   {
     id: '1223128da3_4',
@@ -102,6 +105,7 @@ const tasks: Task[] = [
     fails: undefined,
     file,
     result: { state: 'pass', duration: 100.50598406791687 },
+    context: null as any,
   },
   {
     id: '1223128da3_5',
@@ -112,6 +116,7 @@ const tasks: Task[] = [
     fails: undefined,
     file,
     result: { state: 'pass', duration: 20.184875011444092 },
+    context: null as any,
   },
   {
     id: '1223128da3_6',
@@ -122,6 +127,7 @@ const tasks: Task[] = [
     fails: undefined,
     file,
     result: { state: 'pass', duration: 0.33245420455932617 },
+    context: null as any,
   },
   {
     id: '1223128da3_7',
@@ -132,6 +138,7 @@ const tasks: Task[] = [
     fails: undefined,
     file,
     result: { state: 'pass', duration: 19.738605976104736 },
+    context: null as any,
   },
   {
     id: '1223128da3_8',
@@ -142,6 +149,7 @@ const tasks: Task[] = [
     fails: undefined,
     file,
     result: { state: 'pass', duration: 0.1923508644104004 },
+    context: null as any,
     logs: [
       {
         content: 'error',

test/reporters/tsconfig.json

@@ -1,9 +1,3 @@
 {
-  "compilerOptions": {
-    "paths": {
-      "vitest": ["./packages/vitest/index.d.ts"],
-      "vitest/node": ["./packages/vitest/node.d.ts"],
-      "vitest/config": ["./packages/vitest/config.d.ts"],
-    }
-  }
+  "extends": "../../tsconfig.json"
 }

test/snapshots/test/__snapshots__/shapshots-concurrent.test.ts.snap

@@ -0,0 +1,9 @@
+// Vitest Snapshot v1
+
+exports[`four 1`] = `"four"`;
+
+exports[`one 1`] = `"one"`;
+
+exports[`three 1`] = `"three"`;
+
+exports[`two 1`] = `"two"`;

test/snapshots/test/shapshots-concurrent.test.ts

@@ -0,0 +1,15 @@
+import { test } from 'vitest'
+
+const data = [
+  'one',
+  'two',
+  'three',
+  'four',
+]
+
+data.forEach((i) => {
+  test.concurrent(i, async ({ expect }) => {
+    await new Promise(resolve => setTimeout(resolve, Math.random() * 100))
+    expect(i).toMatchSnapshot()
+  })
+})