From f7c1bdd49b404de14fd02f841f27766ae45b4ee8 Mon Sep 17 00:00:00 2001
From: regevbr <regevbr@gmail.com>
Date: Sun, 29 Mar 2020 22:09:31 +0300
Subject: [PATCH] feat: add functionality to safe mock functions #17

---
 src/Mock.spec.ts | 18 ++++++++++++++----
 src/Mock.ts      |  8 ++++++++
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/Mock.spec.ts b/src/Mock.spec.ts
index 7ac5956..c70f6c8 100644
--- a/src/Mock.spec.ts
+++ b/src/Mock.spec.ts
@@ -1,4 +1,4 @@
-import mock, { mockClear, mockDeep, mockReset } from './Mock';
+import mock, { mockClear, mockDeep, mockReset, mockFn } from './Mock';
 import { anyNumber } from './Matchers';
 import calledWithFn from './CalledWithFn';
 
@@ -107,9 +107,9 @@ describe('jest-mock-extended', () => {
 
     describe('calledWith', () => {
         test('can use calledWith without mock', () => {
-            const mockFn = calledWithFn();
-            mockFn.calledWith(anyNumber(), anyNumber()).mockReturnValue(3);
-            expect(mockFn(1, 2)).toBe(3);
+            const mockFunc = calledWithFn();
+            mockFunc.calledWith(anyNumber(), anyNumber()).mockReturnValue(3);
+            expect(mockFunc(1, 2)).toBe(3);
         });
 
         test('Can specify matchers', () => {
@@ -349,4 +349,14 @@ describe('jest-mock-extended', () => {
             expect(mockObj.deepProp.getNumber(1)).toBe(4);
         });
     });
+
+    describe('function mock', () => {
+        test('should mock function', async () => {
+            type MyFn = (x: number, y: number) => Promise<string>;
+            const mockFunc = mockFn<MyFn>();
+            mockFunc.mockResolvedValue(`str`);
+            const result: string = await mockFunc(1, 2);
+            expect(result).toBe(`str`);
+        });
+    })
 });
diff --git a/src/Mock.ts b/src/Mock.ts
index cb15d41..fc8f8d0 100644
--- a/src/Mock.ts
+++ b/src/Mock.ts
@@ -116,4 +116,12 @@ const mock = <T>(mockImplementation: DeepPartial<T> = {} as DeepPartial<T>, opts
     return overrideMockImp(mockImplementation, opts);
 };
 
+export const mockFn = <T extends Function>(): T extends (...args: infer A) => infer B ? CalledWithMock<B, A> & T : T => {
+    const mockImplementation = calledWithFn();
+    // @ts-ignore private
+    mockImplementation!._isMockFunction = true;
+    // @ts-ignore
+    return mockImplementation;
+};
+
 export default mock;