From 801d5c01f0c54edf87d593b009461615f3f7f735 Mon Sep 17 00:00:00 2001
From: Tom Mrazauskas <tom@mrazauskas.de>
Date: Fri, 2 Sep 2022 14:20:57 +0300
Subject: [PATCH] fix(jest-mock): `generateFromMetadata` return type

---
 .../__typetests__/ModuleMocker.test.ts        |  75 +++++++++++
 packages/jest-mock/src/index.ts               | 116 ++++++++----------
 2 files changed, 126 insertions(+), 65 deletions(-)
 create mode 100644 packages/jest-mock/__typetests__/ModuleMocker.test.ts

diff --git a/packages/jest-mock/__typetests__/ModuleMocker.test.ts b/packages/jest-mock/__typetests__/ModuleMocker.test.ts
new file mode 100644
index 000000000000..c7b275bfebe3
--- /dev/null
+++ b/packages/jest-mock/__typetests__/ModuleMocker.test.ts
@@ -0,0 +1,75 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import {expectType} from 'tsd-lite';
+import {MockMetadata, Mocked, ModuleMocker} from 'jest-mock';
+
+class ExampleClass {
+  memberA: Array<number>;
+
+  constructor() {
+    this.memberA = [1, 2, 3];
+  }
+  memberB() {}
+}
+
+const exampleModule = {
+  instance: new ExampleClass(),
+
+  methodA: function square(a: number, b: number) {
+    return a * b;
+  },
+  methodB: async function asyncSquare(a: number, b: number) {
+    const result = (await a) * b;
+    return result;
+  },
+
+  propertyA: {
+    one: 'foo',
+    three: {
+      nine: 1,
+      ten: [1, 2, 3],
+    },
+    two() {},
+  },
+  propertyB: [1, 2, 3],
+  propertyC: 123,
+  propertyD: 'baz',
+  propertyE: true,
+  propertyF: Symbol.for('a.b.c'),
+};
+
+const moduleMocker = new ModuleMocker(globalThis);
+
+// getMetadata
+
+const exampleMetadata = moduleMocker.getMetadata(exampleModule);
+
+expectType<MockMetadata<typeof exampleModule> | null>(exampleMetadata);
+
+// generateFromMetadata
+
+const exampleMock = moduleMocker.generateFromMetadata(exampleMetadata!);
+
+expectType<Mocked<typeof exampleModule>>(exampleMock);
+
+expectType<Array<[a: number, b: number]>>(exampleMock.methodA.mock.calls);
+expectType<Array<[a: number, b: number]>>(exampleMock.methodB.mock.calls);
+
+expectType<Array<number>>(exampleMock.instance.memberA);
+expectType<Array<[]>>(exampleMock.instance.memberB.mock.calls);
+
+expectType<string>(exampleMock.propertyA.one);
+expectType<Array<[]>>(exampleMock.propertyA.two.mock.calls);
+expectType<number>(exampleMock.propertyA.three.nine);
+expectType<Array<number>>(exampleMock.propertyA.three.ten);
+
+expectType<Array<number>>(exampleMock.propertyB);
+expectType<number>(exampleMock.propertyC);
+expectType<string>(exampleMock.propertyD);
+expectType<boolean>(exampleMock.propertyE);
+expectType<symbol>(exampleMock.propertyF);
diff --git a/packages/jest-mock/src/index.ts b/packages/jest-mock/src/index.ts
index 305699b0a50c..224ac7e9c70a 100644
--- a/packages/jest-mock/src/index.ts
+++ b/packages/jest-mock/src/index.ts
@@ -7,7 +7,7 @@
 
 /* eslint-disable local/ban-types-eventually, local/prefer-rest-params-eventually */
 
-export type MockFunctionMetadataType =
+export type MockMetadataType =
   | 'object'
   | 'array'
   | 'regexp'
@@ -17,20 +17,26 @@ export type MockFunctionMetadataType =
   | 'null'
   | 'undefined';
 
-export type MockFunctionMetadata<
-  T extends UnknownFunction = UnknownFunction,
-  MetadataType = MockFunctionMetadataType,
-> = {
+// TODO remove re-export in Jest 30
+export type MockFunctionMetadataType = MockMetadataType;
+
+export type MockMetadata<T, MetadataType = MockMetadataType> = {
   ref?: number;
-  members?: Record<string, MockFunctionMetadata<T>>;
+  members?: Record<string, MockMetadata<T>>;
   mockImpl?: T;
   name?: string;
   refID?: number;
   type?: MetadataType;
-  value?: ReturnType<T>;
+  value?: T;
   length?: number;
 };
 
+// TODO remove re-export in Jest 30
+export type MockFunctionMetadata<
+  T = unknown,
+  MetadataType = MockMetadataType,
+> = MockMetadata<T, MetadataType>;
+
 export type ClassLike = {new (...args: any): any};
 export type FunctionLike = (...args: any) => any;
 
@@ -75,7 +81,7 @@ type MockedObjectShallow<T extends object> = {
     : T[K];
 } & T;
 
-export type Mocked<T extends object> = T extends ClassLike
+export type Mocked<T> = T extends ClassLike
   ? MockedClass<T>
   : T extends FunctionLike
   ? MockedFunction<T>
@@ -83,7 +89,7 @@ export type Mocked<T extends object> = T extends ClassLike
   ? MockedObject<T>
   : T;
 
-export type MockedShallow<T extends object> = T extends ClassLike
+export type MockedShallow<T> = T extends ClassLike
   ? MockedClass<T>
   : T extends FunctionLike
   ? MockedFunctionShallow<T>
@@ -386,7 +392,7 @@ function getObjectType(value: unknown): string {
   return Object.prototype.toString.apply(value).slice(8, -1);
 }
 
-function getType(ref?: unknown): MockFunctionMetadataType | null {
+function getType(ref?: unknown): MockMetadataType | null {
   const typeName = getObjectType(ref);
   if (
     typeName === 'Function' ||
@@ -560,39 +566,30 @@ export class ModuleMocker {
     };
   }
 
-  private _makeComponent<T extends UnknownFunction>(
-    metadata: MockFunctionMetadata<T, 'object'>,
+  private _makeComponent<T extends Record<string, any>>(
+    metadata: MockMetadata<T, 'object'>,
     restore?: () => void,
-  ): Record<string, any>;
-  private _makeComponent<T extends UnknownFunction>(
-    metadata: MockFunctionMetadata<T, 'array'>,
+  ): T;
+  private _makeComponent<T extends Array<unknown>>(
+    metadata: MockMetadata<T, 'array'>,
     restore?: () => void,
-  ): Array<unknown>;
-  private _makeComponent<T extends UnknownFunction>(
-    metadata: MockFunctionMetadata<T, 'regexp'>,
+  ): T;
+  private _makeComponent<T extends RegExp>(
+    metadata: MockMetadata<T, 'regexp'>,
     restore?: () => void,
-  ): RegExp;
-  private _makeComponent<T extends UnknownFunction>(
-    metadata: MockFunctionMetadata<
-      T,
-      'constant' | 'collection' | 'null' | 'undefined'
-    >,
+  ): T;
+  private _makeComponent<T>(
+    metadata: MockMetadata<T, 'constant' | 'collection' | 'null' | 'undefined'>,
     restore?: () => void,
   ): T;
   private _makeComponent<T extends UnknownFunction>(
-    metadata: MockFunctionMetadata<T, 'function'>,
+    metadata: MockMetadata<T, 'function'>,
     restore?: () => void,
   ): Mock<T>;
   private _makeComponent<T extends UnknownFunction>(
-    metadata: MockFunctionMetadata<T>,
+    metadata: MockMetadata<T>,
     restore?: () => void,
-  ):
-    | Record<string, any>
-    | Array<unknown>
-    | RegExp
-    | ReturnType<T>
-    | undefined
-    | Mock<T> {
+  ): Record<string, any> | Array<unknown> | RegExp | T | Mock | undefined {
     if (metadata.type === 'object') {
       return new this._environmentGlobal.Object();
     } else if (metadata.type === 'array') {
@@ -808,7 +805,7 @@ export class ModuleMocker {
   }
 
   private _createMockFunction<T extends UnknownFunction>(
-    metadata: MockFunctionMetadata<T>,
+    metadata: MockMetadata<T>,
     mockConstructor: Function,
   ): Function {
     let name = metadata.name;
@@ -828,15 +825,13 @@ export class ModuleMocker {
       } while (name && name.startsWith(boundFunctionPrefix));
     }
 
-    // Special case functions named `mockConstructor` to guard for infinite
-    // loops.
+    // Special case functions named `mockConstructor` to guard for infinite loops
     if (name === MOCK_CONSTRUCTOR_NAME) {
       return mockConstructor;
     }
 
     if (
-      // It's a syntax error to define functions with a reserved keyword
-      // as name.
+      // It's a syntax error to define functions with a reserved keyword as name
       RESERVED_KEYWORDS.has(name) ||
       // It's also a syntax error to define functions with a name that starts with a number
       /^\d/.test(name)
@@ -862,19 +857,14 @@ export class ModuleMocker {
     return createConstructor(mockConstructor);
   }
 
-  private _generateMock<T extends UnknownFunction>(
-    metadata: MockFunctionMetadata<T>,
+  private _generateMock<T>(
+    metadata: MockMetadata<T>,
     callbacks: Array<Function>,
-    refs: {
-      [key: string]:
-        | Record<string, any>
-        | Array<unknown>
-        | RegExp
-        | UnknownFunction
-        | undefined
-        | Mock<T>;
-    },
-  ): Mock<T> {
+    refs: Record<
+      number,
+      Record<string, any> | Array<unknown> | RegExp | T | Mock | undefined
+    >,
+  ): Mocked<T> {
     // metadata not compatible but it's the same type, maybe problem with
     // overloading of _makeComponent and not _generateMock?
     // @ts-expect-error - unsure why TSC complains here?
@@ -905,7 +895,7 @@ export class ModuleMocker {
       mock.prototype.constructor = mock;
     }
 
-    return mock as Mock<T>;
+    return mock as Mocked<T>;
   }
 
   /**
@@ -913,12 +903,10 @@ export class ModuleMocker {
    * @param metadata Metadata for the mock in the schema returned by the
    * getMetadata method of this module.
    */
-  generateFromMetadata<T extends UnknownFunction>(
-    metadata: MockFunctionMetadata<T>,
-  ): Mock<T> {
+  generateFromMetadata<T>(metadata: MockMetadata<T>): Mocked<T> {
     const callbacks: Array<Function> = [];
     const refs = {};
-    const mock = this._generateMock(metadata, callbacks, refs);
+    const mock = this._generateMock<T>(metadata, callbacks, refs);
     callbacks.forEach(setter => setter());
     return mock;
   }
@@ -927,11 +915,11 @@ export class ModuleMocker {
    * @see README.md
    * @param component The component for which to retrieve metadata.
    */
-  getMetadata<T extends UnknownFunction>(
-    component: ReturnType<T>,
-    _refs?: Map<ReturnType<T>, number>,
-  ): MockFunctionMetadata<T> | null {
-    const refs = _refs || new Map<ReturnType<T>, number>();
+  getMetadata<T = unknown>(
+    component: T,
+    _refs?: Map<T, number>,
+  ): MockMetadata<T> | null {
+    const refs = _refs || new Map<T, number>();
     const ref = refs.get(component);
     if (ref != null) {
       return {ref};
@@ -942,7 +930,7 @@ export class ModuleMocker {
       return null;
     }
 
-    const metadata: MockFunctionMetadata<T> = {type};
+    const metadata: MockMetadata<T> = {type};
     if (
       type === 'constant' ||
       type === 'collection' ||
@@ -966,9 +954,7 @@ export class ModuleMocker {
     metadata.refID = refs.size;
     refs.set(component, metadata.refID);
 
-    let members: {
-      [key: string]: MockFunctionMetadata<T>;
-    } | null = null;
+    let members: Record<string, MockMetadata<T>> | null = null;
     // Leave arrays alone
     if (type !== 'array') {
       // @ts-expect-error component is object
@@ -1007,7 +993,7 @@ export class ModuleMocker {
   ): fn is Mock<(...args: P) => R>;
   isMockFunction(fn: unknown): fn is Mock<UnknownFunction>;
   isMockFunction(fn: unknown): fn is Mock<UnknownFunction> {
-    return fn != null && (fn as any)._isMockFunction === true;
+    return fn != null && (fn as Mock)._isMockFunction === true;
   }
 
   fn<T extends FunctionLike = UnknownFunction>(implementation?: T): Mock<T> {