From 3fbd136e7f59361b608151663e5596e4d84450b0 Mon Sep 17 00:00:00 2001
From: antoinebrault <antoinebrault@gmail.com>
Date: Tue, 5 Feb 2019 21:17:57 -0500
Subject: [PATCH] [jest] fix type inference for Mocked &
 mockResolvedValue/mockRejectedValue

---
 types/jest/index.d.ts    | 10 ++++----
 types/jest/jest-tests.ts | 55 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 5 deletions(-)

diff --git a/types/jest/index.d.ts b/types/jest/index.d.ts
index cfc9be2070..6cee72a6d1 100644
--- a/types/jest/index.d.ts
+++ b/types/jest/index.d.ts
@@ -804,7 +804,7 @@ declare namespace jest {
      *  myApi.myApiMethod.mockImplementation(() => "test");
      */
     type Mocked<T> = {
-        [P in keyof T]: T[P] & MockInstance<T[P], ArgsType<T[P]>>;
+        [P in keyof T]: T[P] extends (...args: any[]) => any ? MockInstance<ReturnType<T[P]>, ArgsType<T[P]>>: T[P];
     } & T;
 
     interface MockInstance<T, Y extends any[]> {
@@ -912,7 +912,7 @@ declare namespace jest {
         /**
          * Simple sugar function for: `jest.fn().mockImplementation(() => Promise.resolve(value));`
          */
-        mockResolvedValue(value: T | PromiseLike<T>): Mock<Promise<T>, Y>;
+        mockResolvedValue(value: T extends PromiseLike<infer U> ? U | T : never): Mock<T, Y>;
         /**
          * Simple sugar function for: `jest.fn().mockImplementationOnce(() => Promise.resolve(value));`
          *
@@ -932,7 +932,7 @@ declare namespace jest {
          * });
          *
          */
-        mockResolvedValueOnce(value: T | PromiseLike<T>): Mock<Promise<T>, Y>;
+        mockResolvedValueOnce(value: T extends PromiseLike<infer U> ? U | T : never): Mock<T, Y>;
         /**
          * Simple sugar function for: `jest.fn().mockImplementation(() => Promise.reject(value));`
          *
@@ -944,7 +944,7 @@ declare namespace jest {
          *   await asyncMock(); // throws "Async error"
          * });
          */
-        mockRejectedValue(value: any): Mock<Promise<T>, Y>;
+        mockRejectedValue(value: T extends PromiseLike<any> ? any : never): Mock<T, Y>;
 
         /**
          * Simple sugar function for: `jest.fn().mockImplementationOnce(() => Promise.reject(value));`
@@ -962,7 +962,7 @@ declare namespace jest {
          * });
          *
          */
-        mockRejectedValueOnce(value: any): Mock<Promise<T>, Y>;
+        mockRejectedValueOnce(value: T extends PromiseLike<any> ? any : never): Mock<T, Y>;
     }
 
     /**
diff --git a/types/jest/jest-tests.ts b/types/jest/jest-tests.ts
index c2e930c26f..00ccd38106 100644
--- a/types/jest/jest-tests.ts
+++ b/types/jest/jest-tests.ts
@@ -412,6 +412,61 @@ const spy6 = jest.spyOn(spiedTarget2, "value", "set");
 
 // should compile
 jest.fn().mockImplementation((test: number) => test);
+jest.fn().mockResolvedValue(1);
+
+interface Type1 { a: number; }
+interface Type2 { b: number; }
+class TestMocked {
+    field: string;
+    test1(x: Type1): Promise<Type1> {
+        return Promise.resolve(x);
+    }
+    test2(x: Promise<Type1>): Promise<Type1> {
+        return x;
+    }
+    test3(x: Promise<Type1>): Promise<Type2> {
+        return x.then(() => ({ b: 1 }));
+    }
+    test4(x: Type1): Type1 {
+        return x;
+    }
+}
+
+const mocked: jest.Mocked<TestMocked> = new TestMocked() as any;
+mocked.test1.mockImplementation(() => Promise.resolve({ a: 1 }));
+mocked.test1.mockReturnValue(Promise.resolve({ a: 1 }));
+// $ExpectType Mock<Promise<Type1>, [Type1]>
+mocked.test1.mockResolvedValue({ a: 1 });
+mocked.test1.mockResolvedValueOnce({ a: 1 });
+// $ExpectType Mock<Promise<Type1>, [Type1]>
+mocked.test1.mockResolvedValue(Promise.resolve({ a: 1 }));
+mocked.test1.mockResolvedValueOnce(Promise.resolve({ a: 1 }));
+// $ExpectType Mock<Promise<Type1>, [Promise<Type1>]>
+mocked.test2.mockResolvedValue({ a: 1 });
+mocked.test2.mockResolvedValueOnce({ a: 1 });
+// $ExpectType Mock<Promise<Type1>, [Promise<Type1>]>
+mocked.test2.mockResolvedValue(Promise.resolve({ a: 1 }));
+mocked.test2.mockResolvedValueOnce(Promise.resolve({ a: 1 }));
+// $ExpectType Mock<Promise<Type2>, [Promise<Type1>]>
+mocked.test3.mockResolvedValue({ b: 1 });
+mocked.test3.mockResolvedValueOnce({ b: 1 });
+// $ExpectType Mock<Promise<Type2>, [Promise<Type1>]>
+mocked.test3.mockResolvedValue(Promise.resolve({ b: 1 }));
+mocked.test3.mockResolvedValueOnce(Promise.resolve({ b: 1 }));
+mocked.test3.mockRejectedValue(new Error());
+mocked.test3.mockRejectedValueOnce(new Error());
+// $ExpectError
+mocked.test4.mockResolvedValue({ a: 1 });
+// $ExpectError
+mocked.test4.mockResolvedValueOnce({ a: 1 });
+// $ExpectError
+mocked.test4.mockResolvedValue(Promise.resolve({ a: 1 }));
+// $ExpectError
+mocked.test4.mockResolvedValueOnce(Promise.resolve({ a: 1 }));
+// $ExpectError
+mocked.test4.mockRejectedValue(new Error());
+// $ExpectError
+mocked.test4.mockRejectedValueOnce(new Error());
 
 /* Snapshot serialization */