From 426503cb19bd360413b7964e995861404cf64ad5 Mon Sep 17 00:00:00 2001
From: Tanimodori <unknowner2014@gmail.com>
Date: Sat, 6 Aug 2022 18:43:51 +0800
Subject: [PATCH] feat(types): better local test context support (#1805)

* feat(types): better local test context support

* feat(types): fix beforeEach and afterEach

* test: add unit test for local context
---
 packages/vitest/src/runtime/hooks.ts |  4 ++--
 packages/vitest/src/runtime/suite.ts |  4 ++--
 packages/vitest/src/types/tasks.ts   | 24 +++++++++++++--------
 test/core/test/local-context.test.ts | 31 ++++++++++++++++++++++++++++
 4 files changed, 50 insertions(+), 13 deletions(-)
 create mode 100644 test/core/test/local-context.test.ts

diff --git a/packages/vitest/src/runtime/hooks.ts b/packages/vitest/src/runtime/hooks.ts
index 4abfd59a9f14..df20f8605ec1 100644
--- a/packages/vitest/src/runtime/hooks.ts
+++ b/packages/vitest/src/runtime/hooks.ts
@@ -5,5 +5,5 @@ import { getCurrentSuite } from './suite'
 // suite hooks
 export const beforeAll = (fn: SuiteHooks['beforeAll'][0], timeout?: number) => getCurrentSuite().on('beforeAll', withTimeout(fn, timeout ?? getDefaultHookTimeout(), true))
 export const afterAll = (fn: SuiteHooks['afterAll'][0], timeout?: number) => getCurrentSuite().on('afterAll', withTimeout(fn, timeout ?? getDefaultHookTimeout(), true))
-export const beforeEach = (fn: SuiteHooks['beforeEach'][0], timeout?: number) => getCurrentSuite().on('beforeEach', withTimeout(fn, timeout ?? getDefaultHookTimeout(), true))
-export const afterEach = (fn: SuiteHooks['afterEach'][0], timeout?: number) => getCurrentSuite().on('afterEach', withTimeout(fn, timeout ?? getDefaultHookTimeout(), true))
+export const beforeEach = <ExtraContext = {}>(fn: SuiteHooks<ExtraContext>['beforeEach'][0], timeout?: number) => getCurrentSuite<ExtraContext>().on('beforeEach', withTimeout(fn, timeout ?? getDefaultHookTimeout(), true))
+export const afterEach = <ExtraContext = {}>(fn: SuiteHooks<ExtraContext>['afterEach'][0], timeout?: number) => getCurrentSuite<ExtraContext>().on('afterEach', withTimeout(fn, timeout ?? getDefaultHookTimeout(), true))
diff --git a/packages/vitest/src/runtime/suite.ts b/packages/vitest/src/runtime/suite.ts
index 3cda46e3e359..887bb69eecf3 100644
--- a/packages/vitest/src/runtime/suite.ts
+++ b/packages/vitest/src/runtime/suite.ts
@@ -50,8 +50,8 @@ export function clearCollectorContext() {
   collectorContext.currentSuite = defaultSuite
 }
 
-export function getCurrentSuite() {
-  return collectorContext.currentSuite || defaultSuite
+export function getCurrentSuite<ExtraContext = {}>() {
+  return (collectorContext.currentSuite || defaultSuite) as SuiteCollector<ExtraContext>
 }
 
 export function createSuiteHooks() {
diff --git a/packages/vitest/src/types/tasks.ts b/packages/vitest/src/types/tasks.ts
index 991081a74bca..ac96e21653af 100644
--- a/packages/vitest/src/types/tasks.ts
+++ b/packages/vitest/src/types/tasks.ts
@@ -126,7 +126,10 @@ type ChainableTestAPI<ExtraContext = {}> = ChainableFunction<
   'concurrent' | 'only' | 'skip' | 'todo' | 'fails',
   [name: string, fn?: TestFunction<ExtraContext>, timeout?: number],
   void,
-  { each: TestEachFunction }
+  {
+    each: TestEachFunction
+    <T extends ExtraContext>(name: string, fn?: TestFunction<T>, timeout?: number): void
+  }
 >
 
 export type TestAPI<ExtraContext = {}> = ChainableTestAPI<ExtraContext> & {
@@ -137,12 +140,15 @@ export type TestAPI<ExtraContext = {}> = ChainableTestAPI<ExtraContext> & {
 
 type ChainableSuiteAPI<ExtraContext = {}> = ChainableFunction<
   'concurrent' | 'only' | 'skip' | 'todo' | 'shuffle',
-  [name: string, factory?: SuiteFactory],
+  [name: string, factory?: SuiteFactory<ExtraContext>],
   SuiteCollector<ExtraContext>,
-  { each: TestEachFunction }
+  {
+    each: TestEachFunction
+    <T extends ExtraContext>(name: string, factory?: SuiteFactory<T>): SuiteCollector<T>
+  }
 >
 
-export type SuiteAPI<ExtraContext = {}> = ChainableSuiteAPI & {
+export type SuiteAPI<ExtraContext = {}> = ChainableSuiteAPI<ExtraContext> & {
   each: SuiteEachFunction
   skipIf(condition: any): ChainableSuiteAPI<ExtraContext>
   runIf(condition: any): ChainableSuiteAPI<ExtraContext>
@@ -152,11 +158,11 @@ export type HookListener<T extends any[], Return = void> = (...args: T) => Await
 
 export type HookCleanupCallback = (() => Awaitable<unknown>) | void
 
-export interface SuiteHooks {
+export interface SuiteHooks<ExtraContext = {}> {
   beforeAll: HookListener<[Suite | File], HookCleanupCallback>[]
   afterAll: HookListener<[Suite | File]>[]
-  beforeEach: HookListener<[TestContext, Suite], HookCleanupCallback>[]
-  afterEach: HookListener<[TestContext, Suite]>[]
+  beforeEach: HookListener<[TestContext & ExtraContext, Suite], HookCleanupCallback>[]
+  afterEach: HookListener<[TestContext & ExtraContext, Suite]>[]
 }
 
 export interface SuiteCollector<ExtraContext = {}> {
@@ -167,10 +173,10 @@ export interface SuiteCollector<ExtraContext = {}> {
   tasks: (Suite | Test | SuiteCollector<ExtraContext>)[]
   collect: (file?: File) => Promise<Suite>
   clear: () => void
-  on: <T extends keyof SuiteHooks>(name: T, ...fn: SuiteHooks[T]) => void
+  on: <T extends keyof SuiteHooks<ExtraContext>>(name: T, ...fn: SuiteHooks<ExtraContext>[T]) => void
 }
 
-export type SuiteFactory = (test: (name: string, fn: TestFunction) => void) => Awaitable<void>
+export type SuiteFactory<ExtraContext = {}> = (test: (name: string, fn: TestFunction<ExtraContext>) => void) => Awaitable<void>
 
 export interface RuntimeContext {
   tasks: (SuiteCollector | Test)[]
diff --git a/test/core/test/local-context.test.ts b/test/core/test/local-context.test.ts
new file mode 100644
index 000000000000..a475840dfc51
--- /dev/null
+++ b/test/core/test/local-context.test.ts
@@ -0,0 +1,31 @@
+import { beforeEach, describe, expect, it } from 'vitest'
+
+describe('local test context works with explicit type', () => {
+  interface LocalTestContext {
+    foo: string
+  }
+  beforeEach<LocalTestContext>((context) => {
+    context.foo = 'foo'
+  })
+  it<LocalTestContext>('works with explicit type', (context) => {
+    expect(context.foo).toBe('foo')
+  })
+  it.todo.skip<LocalTestContext>('is chainable with explicit type', (context) => {
+    expect(context.foo).toBe('foo')
+  })
+})
+
+describe('local test context works with implicit type', () => {
+  interface LocalTestContext {
+    bar: string
+  }
+  beforeEach((context: LocalTestContext) => {
+    context.bar = 'bar'
+  })
+  it('works with implicit type', (context: LocalTestContext) => {
+    expect(context.bar).toBe('bar')
+  })
+  it.only('is chainable with implicit type', (context: LocalTestContext) => {
+    expect(context.bar).toBe('bar')
+  })
+})