From b9e6335d9a821441e5d6bad92809a1897284d764 Mon Sep 17 00:00:00 2001
From: Vladimir Sheremet <sleuths.slews0s@icloud.com>
Date: Thu, 4 Aug 2022 19:05:27 +0300
Subject: [PATCH] feat: allow timeout in test.each

---
 packages/vitest/src/runtime/suite.ts          |  6 +--
 packages/vitest/src/types/tasks.ts            | 42 +++++++++++++------
 test/fails/fixtures/each-timeout.test.ts      |  5 +++
 .../test/__snapshots__/runner.test.ts.snap    |  2 +
 4 files changed, 40 insertions(+), 15 deletions(-)
 create mode 100644 test/fails/fixtures/each-timeout.test.ts

diff --git a/packages/vitest/src/runtime/suite.ts b/packages/vitest/src/runtime/suite.ts
index 1082484076e6..1cb931d2f7de 100644
--- a/packages/vitest/src/runtime/suite.ts
+++ b/packages/vitest/src/runtime/suite.ts
@@ -166,7 +166,7 @@ function createSuite() {
   const suite = createChainable(
     ['concurrent', 'shuffle', 'skip', 'only', 'todo'],
     function (name: string, factory?: SuiteFactory) {
-      const mode = this.only ? 'only' : this.skip ? 'skip' : this.todo ? 'todo' : 'run'
+      const mode: RunMode = this.only ? 'only' : this.skip ? 'skip' : this.todo ? 'todo' : 'run'
       return createSuiteCollector(name, factory, mode, this.concurrent, this.shuffle)
     },
   ) as SuiteAPI
@@ -200,10 +200,10 @@ function createTest(fn: (
   ) as TestAPI
 
   test.each = <T>(cases: ReadonlyArray<T>) => {
-    return (name: string, fn: (...args: T[]) => void) => {
+    return (name: string, fn: (...args: T[]) => void, timeout?: number) => {
       cases.forEach((i, idx) => {
         const items = Array.isArray(i) ? i : [i]
-        test(formatTitle(name, items, idx), () => fn(...items))
+        test(formatTitle(name, items, idx), () => fn(...items), timeout)
       })
     }
   }
diff --git a/packages/vitest/src/types/tasks.ts b/packages/vitest/src/types/tasks.ts
index ef3ecafb108a..5221a9f3383c 100644
--- a/packages/vitest/src/types/tasks.ts
+++ b/packages/vitest/src/types/tasks.ts
@@ -89,41 +89,59 @@ type ExtractEachCallbackArgs<T extends ReadonlyArray<any>> = {
                     ? 10
                     : 'fallback']
 
-interface EachFunction {
+interface SuiteEachFunction {
   <T extends any[] | [any]>(cases: ReadonlyArray<T>): (
     name: string,
-    fn: (...args: T) => Awaitable<void>
+    fn: (...args: T) => Awaitable<void>,
   ) => void
   <T extends ReadonlyArray<any>>(cases: ReadonlyArray<T>): (
     name: string,
-    fn: (...args: ExtractEachCallbackArgs<T>) => Awaitable<void>
+    fn: (...args: ExtractEachCallbackArgs<T>) => Awaitable<void>,
   ) => void
   <T>(cases: ReadonlyArray<T>): (
     name: string,
-    fn: (...args: T[]) => Awaitable<void>
+    fn: (...args: T[]) => Awaitable<void>,
+  ) => void
+}
+
+interface TestEachFunction {
+  <T extends any[] | [any]>(cases: ReadonlyArray<T>): (
+    name: string,
+    fn: (...args: T) => Awaitable<void>,
+    timeout?: number,
+  ) => void
+  <T extends ReadonlyArray<any>>(cases: ReadonlyArray<T>): (
+    name: string,
+    fn: (...args: ExtractEachCallbackArgs<T>) => Awaitable<void>,
+    timeout?: number,
+  ) => void
+  <T>(cases: ReadonlyArray<T>): (
+    name: string,
+    fn: (...args: T[]) => Awaitable<void>,
+    timeout?: number,
   ) => void
 }
 
 type ChainableTestAPI<ExtraContext = {}> = ChainableFunction<
-'concurrent' | 'only' | 'skip' | 'todo' | 'fails',
-[name: string, fn?: TestFunction<ExtraContext>, timeout?: number],
-void
+  'concurrent' | 'only' | 'skip' | 'todo' | 'fails',
+  [name: string, fn?: TestFunction<ExtraContext>, timeout?: number],
+  void
 >
 
 export type TestAPI<ExtraContext = {}> = ChainableTestAPI<ExtraContext> & {
-  each: EachFunction
+  each: TestEachFunction
   skipIf(condition: any): ChainableTestAPI<ExtraContext>
   runIf(condition: any): ChainableTestAPI<ExtraContext>
 }
 
 type ChainableSuiteAPI<ExtraContext = {}> = ChainableFunction<
-'concurrent' | 'only' | 'skip' | 'todo' | 'shuffle',
-[name: string, factory?: SuiteFactory],
-SuiteCollector<ExtraContext>
+  'concurrent' | 'only' | 'skip' | 'todo' | 'shuffle',
+  [name: string, factory?: SuiteFactory],
+  SuiteCollector<ExtraContext>
 >
 
 export type SuiteAPI<ExtraContext = {}> = ChainableSuiteAPI & {
-  each: EachFunction
+  each: SuiteEachFunction
   skipIf(condition: any): ChainableSuiteAPI<ExtraContext>
   runIf(condition: any): ChainableSuiteAPI<ExtraContext>
 }
diff --git a/test/fails/fixtures/each-timeout.test.ts b/test/fails/fixtures/each-timeout.test.ts
new file mode 100644
index 000000000000..fac22e14c14b
--- /dev/null
+++ b/test/fails/fixtures/each-timeout.test.ts
@@ -0,0 +1,5 @@
+import { test } from 'vitest'
+
+test.each([1])('test each timeout', async () => {
+  await new Promise(resolve => setTimeout(resolve, 20))
+}, 10)
diff --git a/test/fails/test/__snapshots__/runner.test.ts.snap b/test/fails/test/__snapshots__/runner.test.ts.snap
index ed4784ca5e38..46844e79c850 100644
--- a/test/fails/test/__snapshots__/runner.test.ts.snap
+++ b/test/fails/test/__snapshots__/runner.test.ts.snap
@@ -1,5 +1,7 @@
 // Vitest Snapshot v1
 
+exports[`should fails > each-timeout.test.ts > each-timeout.test.ts 1`] = `"Error: Test timed out in 10ms."`;
+
 exports[`should fails > empty.test.ts > empty.test.ts 1`] = `"Error: No test suite found in file <rootDir>/empty.test.ts"`;
 
 exports[`should fails > expect.test.ts > expect.test.ts 1`] = `"AssertionError: expected 2 to deeply equal 3"`;