From aac50234c9f14d39db189b381c40c9f2a9e03f74 Mon Sep 17 00:00:00 2001
From: Victor Villas <mail@victor.villas>
Date: Wed, 16 Sep 2020 22:21:41 -0300
Subject: [PATCH] Add generator function to accepted hook types

---
 packages/jest-circus/src/utils.ts    | 14 +++++++++++---
 packages/jest-types/src/Global.ts    |  7 ++++++-
 test-types/top-level-globals.test.ts |  5 +++++
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/packages/jest-circus/src/utils.ts b/packages/jest-circus/src/utils.ts
index b1e9f8aeb95d..0c4b1abf74a2 100644
--- a/packages/jest-circus/src/utils.ts
+++ b/packages/jest-circus/src/utils.ts
@@ -214,17 +214,25 @@ export const callAsyncCircusFn = (
         });
       };
 
-      returnedValue = fn.call(testContext, done);
+      returnedValue = fn.call<
+        Circus.TestContext | undefined,
+        Array<typeof done>,
+        void | Promise<unknown> | Generator | undefined
+      >(testContext, done);
 
       return;
     }
 
-    let returnedValue;
+    let returnedValue: any;
     if (isGeneratorFn(fn)) {
       returnedValue = co.wrap(fn).call({});
     } else {
       try {
-        returnedValue = fn.call(testContext);
+        returnedValue = fn.call<
+          Circus.TestContext | undefined,
+          [],
+          void | Promise<unknown> | Generator | undefined
+        >(testContext);
       } catch (error) {
         reject(error);
         return;
diff --git a/packages/jest-types/src/Global.ts b/packages/jest-types/src/Global.ts
index 411e61191d23..2dd190717654 100644
--- a/packages/jest-types/src/Global.ts
+++ b/packages/jest-types/src/Global.ts
@@ -7,6 +7,11 @@
 
 import type {CoverageMapData} from 'istanbul-lib-coverage';
 
+// TypeScript's GeneratorFunction is too narrow
+interface GeneratorFn {
+  (...args: Array<any>): Generator;
+}
+
 export type DoneFn = (reason?: string | Error) => void;
 export type TestName = string;
 export type TestFn = (
@@ -17,7 +22,7 @@ export type ConcurrentTestFn = (
 ) => Promise<void | undefined | unknown>;
 export type BlockFn = () => void;
 export type BlockName = string;
-export type HookFn = TestFn;
+export type HookFn = TestFn | GeneratorFn;
 
 export type Col = unknown;
 export type Row = Array<Col>;
diff --git a/test-types/top-level-globals.test.ts b/test-types/top-level-globals.test.ts
index 9b6499411362..de98c468623e 100644
--- a/test-types/top-level-globals.test.ts
+++ b/test-types/top-level-globals.test.ts
@@ -20,6 +20,7 @@ import {
 
 const fn = () => {};
 const asyncFn = async () => {};
+const genFn = function* () {};
 const timeout = 5;
 const testName = 'Test name';
 const testTable = [[1, 2]];
@@ -27,15 +28,19 @@ const testTable = [[1, 2]];
 // https://jestjs.io/docs/en/api#methods
 expectType<void>(afterAll(fn));
 expectType<void>(afterAll(asyncFn));
+expectType<void>(afterAll(genFn));
 expectType<void>(afterAll(fn, timeout));
 expectType<void>(afterEach(fn));
 expectType<void>(afterEach(asyncFn));
+expectType<void>(afterEach(genFn));
 expectType<void>(afterEach(fn, timeout));
 expectType<void>(beforeAll(fn));
 expectType<void>(beforeAll(asyncFn));
+expectType<void>(beforeAll(genFn));
 expectType<void>(beforeAll(fn, timeout));
 expectType<void>(beforeEach(fn));
 expectType<void>(beforeEach(asyncFn));
+expectType<void>(beforeEach(genFn));
 expectType<void>(beforeEach(fn, timeout));
 
 expectType<void>(test.each(testTable)(testName, fn));