From 23b5ee987193cd835cda129017f27ac108b60930 Mon Sep 17 00:00:00 2001
From: Henry Cai <huacong.cai@gmail.com>
Date: Thu, 3 Apr 2025 12:47:36 -0700
Subject: [PATCH] feat(utils): add mocked type helpers

similar to Vitest's vi.mocked
---
 README.md        | 45 ++++++++++++++++++++++++++++++++++++++++++++-
 src/Mock.spec.ts | 30 +++++++++++++++++++++++++++++-
 src/Mock.ts      | 12 +++++++++++-
 3 files changed, 84 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 14412f67..93de9664 100644
--- a/README.md
+++ b/README.md
@@ -159,7 +159,7 @@ expect(mockObj.deepProp(1)).toBe(3);
 expect(mockObj.deepProp.getNumber(1)).toBe(4);
 ```
 
-Can can provide a fallback mock implementation used if you do not define a return value using `calledWith`.
+You can also provide a fallback mock implementation to be used if you do not define a return value using `calledWith`.
 
 ```ts
 import { mockDeep } from 'jest-mock-extended';
@@ -171,6 +171,49 @@ const mockObj = mockDeep<Test1>({
 expect(() => mockObj.getNumber()).toThrowError('not mocked');
 ```
 
+## Mocked Type Helpers
+If you mock objects/functions of modules and can't refer them directly due to hoist,
+mocked type helpers can be used. This is similar to the use case of Vitest's `vi.mocked`
+```ts
+// APIs
+import { mocked, mockedFn } from "vitest-mock-extended";
+import { originalObj, originalFn} from "somewhere";
+
+const mockedObj = mocked(originalObj);
+const deepMockedObj = mocked(originalObj, true);
+const mockedFunction = mockedFn(originalFn);
+```
+
+An example would be
+```ts
+// Mock a module
+// @/libs/example.mock.ts
+import { mock } from "vitest-mock-extended";
+import { ExampleClient } from "@/libs/example";
+
+vi.mock(import("@/libs/example"), async (importOriginal) => {
+  const actual = await importOriginal();
+  return {
+    ...actual,
+    // Due to vi.mock being hoisted, we have to mock here directly instead of 
+    // defining an exampleMock outside and assign it to example
+    example: mock<ExampleClient>(),
+  };
+});
+
+// Use the mocked object
+import "@/libs/example.mock"; // Mock out the actual client
+import { mocked } from "vitest-mock-extended";
+import { example } from "@/libs/example";
+
+const exampleMock = mocked(example);
+
+test("send notification", () => {
+    example.function.mockResolvedValue(xxx);
+    ...
+}
+```
+
 ## Available Matchers
 
 | Matcher                | Description                                                |
diff --git a/src/Mock.spec.ts b/src/Mock.spec.ts
index 5fbba4b3..8d6abb30 100644
--- a/src/Mock.spec.ts
+++ b/src/Mock.spec.ts
@@ -1,4 +1,4 @@
-import { mock, mockClear, mockDeep, mockReset, mockFn, VitestMockExtended } from './Mock'
+import { mock, mockClear, mockDeep, mockReset, mockFn, VitestMockExtended, mocked, mockedFn } from './Mock'
 import { anyNumber } from './Matchers'
 import { calledWithFn } from './CalledWithFn'
 import { MockProxy } from './Mock'
@@ -698,4 +698,32 @@ describe('vitest-mock-extended', () => {
       expect({ cookie: { domain: 'dummy' } }).toStrictEqual({ cookie })
     })
   })
+
+  describe('mock type utils', () => {
+    it('mocked should handle mock obj', () => {
+      const mockObj = mock<MockInt>()
+      const obj = mocked(mockObj)
+
+      expect(obj.getNumber).toHaveBeenCalledTimes(0)
+    })
+
+    it('mocked should handle mockDeep obj', () => {
+      const mockObj = mockDeep<Test6>({ funcPropSupport: true })
+      const input = new Test1(1)
+      mockObj.funcValueProp.nonDeepProp.calledWith(input).mockReturnValue(4)
+      const obj = mocked(mockObj, true)
+
+      expect(obj.funcValueProp.nonDeepProp(input)).toBe(4)
+    })
+
+    it('mockedFn should handle mockFn obj', async () => {
+      type MyFn = (x: number, y: number) => Promise<string>
+      const mockFunc = mockFn<MyFn>()
+      mockFunc.mockResolvedValue(`str`)
+
+      const obj = mockedFn(mockFunc)
+      const result: string = await obj(1, 2)
+      expect(result).toBe(`str`)
+    })
+  })
 })
diff --git a/src/Mock.ts b/src/Mock.ts
index 5af8966d..5cb08e45 100644
--- a/src/Mock.ts
+++ b/src/Mock.ts
@@ -205,6 +205,16 @@ const mockFn = <
   return calledWithFn()
 }
 
+function mocked<T>(obj: T, deep?: false): ReturnType<typeof mock<T>>
+function mocked<T>(obj: T, deep: true): ReturnType<typeof mockDeep<T>>
+function mocked<T>(obj: T, _deep?: boolean) {
+  return obj;
+}
+
+function mockedFn<T>(obj: T) {
+  return obj as ReturnType<typeof mockFn<T>>;
+}
+
 const stub = <T extends object>(): T => {
   return new Proxy<T>({} as T, {
     get: (obj, property: ProxiedProperty) => {
@@ -217,5 +227,5 @@ const stub = <T extends object>(): T => {
   })
 }
 
-export { mock, VitestMockExtended, mockClear, mockReset, mockDeep, mockFn, stub }
+export { mock, VitestMockExtended, mockClear, mockReset, mockDeep, mockFn, stub, mocked, mockedFn }
 export type { GlobalConfig, CalledWithMock, MockProxy, DeepMockProxy, MockOpts }