From 02eb750a82f89588b6d93dc4faaf55d3d5eba479 Mon Sep 17 00:00:00 2001
From: skovhus <kenneth.skovhus@gmail.com>
Date: Mon, 27 Feb 2023 10:59:30 +0100
Subject: [PATCH] Correct pass fallbackMockImplementation

---
 src/CalledWithFn.ts | 17 ++++++++++++-----
 src/Mock.spec.ts    |  8 +++++++-
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/CalledWithFn.ts b/src/CalledWithFn.ts
index 9b8a165..4690fe2 100644
--- a/src/CalledWithFn.ts
+++ b/src/CalledWithFn.ts
@@ -13,8 +13,12 @@ function isJestAsymmetricMatcher(obj: any): obj is JestAsymmetricMatcher {
     return !!obj && typeof obj === 'object' && 'asymmetricMatch' in obj && typeof obj.asymmetricMatch === 'function';
 }
 
-const checkCalledWith = <T, Y extends any[]>(calledWithStack: CalledWithStackItem<T, Y>[], actualArgs: Y): T => {
-    const calledWithInstance = calledWithStack.find(instance =>
+const checkCalledWith = <T, Y extends any[]>(
+    calledWithStack: CalledWithStackItem<T, Y>[],
+    actualArgs: Y,
+    fallbackMockImplementation?: (...args: Y) => T
+): T => {
+    const calledWithInstance = calledWithStack.find((instance) =>
         instance.args.every((matcher, i) => {
             if (matcher instanceof Matcher) {
                 return matcher.asymmetricMatch(actualArgs[i]);
@@ -29,7 +33,9 @@ const checkCalledWith = <T, Y extends any[]>(calledWithStack: CalledWithStackIte
     );
 
     // @ts-ignore cannot return undefined, but this will fail the test if there is an expectation which is what we want
-    return calledWithInstance ? calledWithInstance.calledWithFn(...actualArgs) : undefined;
+    return calledWithInstance
+        ? calledWithInstance.calledWithFn(...actualArgs)
+        : fallbackMockImplementation && fallbackMockImplementation(...actualArgs);
 };
 
 export const calledWithFn = <T, Y extends any[]>({
@@ -42,9 +48,10 @@ export const calledWithFn = <T, Y extends any[]>({
         // We create new function to delegate any interactions (mockReturnValue etc.) to for this set of args.
         // If that set of args is matched, we just call that jest.fn() for the result.
         const calledWithFn = jest.fn(fallbackMockImplementation);
-        if (!fn.getMockImplementation()) {
+        const mockImplementation = fn.getMockImplementation();
+        if (!mockImplementation || mockImplementation === fallbackMockImplementation) {
             // Our original function gets a mock implementation which handles the matching
-            fn.mockImplementation((...args: Y) => checkCalledWith(calledWithStack, args));
+            fn.mockImplementation((...args: Y) => checkCalledWith(calledWithStack, args, fallbackMockImplementation));
             calledWithStack = [];
         }
         calledWithStack.unshift({ args, calledWithFn });
diff --git a/src/Mock.spec.ts b/src/Mock.spec.ts
index 7514c18..0f5729a 100644
--- a/src/Mock.spec.ts
+++ b/src/Mock.spec.ts
@@ -244,7 +244,7 @@ describe('jest-mock-extended', () => {
             mockObj.getSomethingWithArgs.calledWith(1, 2).mockReturnValue(3);
 
             expect(mockObj.getSomethingWithArgs(1, 2)).toBe(3);
-        })
+        });
 
         test('Support jest matcher', () => {
             const mockObj = mock<MockInt>();
@@ -328,7 +328,9 @@ describe('jest-mock-extended', () => {
                     throw new Error('not mocked');
                 },
             });
+            mockObj.deepProp.getAnotherString.calledWith('foo'); // no mock implementation
             expect(() => mockObj.getNumber()).toThrowError('not mocked');
+            expect(() => mockObj.deepProp.getAnotherString('foo')).toThrowError('not mocked');
         });
 
         test('fallback mock implementation can be overridden while also providing a mock implementation', () => {
@@ -344,8 +346,12 @@ describe('jest-mock-extended', () => {
                     },
                 }
             );
+            mockObj.deepProp.getAnotherString.calledWith('?').mockReturnValue('mocked');
             expect(mockObj.getNumber()).toBe(150);
+            expect(mockObj.deepProp.getAnotherString('?')).toBe('mocked');
+
             expect(() => mockObj.deepProp.getNumber(1)).toThrowError('not mocked');
+            expect(() => mockObj.deepProp.getAnotherString('!')).toThrowError('not mocked');
         });
     });