[Test] Centralise common mocks

.cursor/rules/vitest.mdc

@@ -10,6 +10,7 @@ globs:
   - Directory structure of `tests/unit/` should match that of the tested file in `src/`
 - Tests should be cross-platform compatible; able to run on Windows, macOS, and linux
   - e.g. the use of `path.normalize`, or `path.join` and `path.sep` to ensure that tests work the same on all platforms
+- vitest automatically runs the setup file [setup.ts](mdc:tests/unit/setup.ts) before running each test file
 - Tests should be mocked properly
   - Mocks should be cleanly written and easy to understand
   - Mocks should be re-usable where possible

tests/unit/config/comfyServerConfig.test.ts

@@ -9,12 +9,6 @@ import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi }
 
 import { ComfyServerConfig } from '@/config/comfyServerConfig';
 
-vi.mock('electron', () => ({
-  app: {
-    getPath: vi.fn(),
-  },
-}));
-
 vi.mock('@/install/resourcePaths', () => ({
   getAppResourcesPath: vi.fn(() => '/mocked/app_resources'),
 }));

tests/unit/desktopApp.test.ts

@@ -18,29 +18,6 @@ import { promptMetricsConsent } from '@/services/telemetry';
 import type { DesktopConfig } from '@/store/desktopConfig';
 
 // Mock dependencies
-vi.mock('electron', () => ({
-  app: {
-    quit: vi.fn(() => {
-      throw new Error('Test exited via app.quit()');
-    }),
-    exit: vi.fn(() => {
-      throw new Error('Test exited via app.exit()');
-    }),
-    getPath: vi.fn(() => '/mock/app/path'),
-    getAppPath: vi.fn(() => '/mock/app/path'),
-  },
-  dialog: {
-    showErrorBox: vi.fn(),
-  },
-  ipcMain: {
-    on: vi.fn(),
-    once: vi.fn(),
-    handle: vi.fn(),
-    handleOnce: vi.fn(),
-    removeHandler: vi.fn(),
-  },
-}));
-
 const mockAppWindow = {
   loadPage: vi.fn(),
   send: vi.fn(),

tests/unit/handlers/AppHandlers.test.ts

@@ -1,31 +1,17 @@
 import { app, ipcMain } from 'electron';
-import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest';
 
 import { IPC_CHANNELS } from '@/constants';
 import { registerAppHandlers } from '@/handlers/AppHandlers';
 
-vi.mock('electron', () => ({
-  app: {
-    quit: vi.fn(),
-  },
-  ipcMain: {
-    handle: vi.fn(),
-  },
-}));
-
-interface TestCase {
-  channel: string;
-  expected: any;
-  args?: any[];
-}
+import { quitMessage } from '../setup';
+
 const getHandler = (channel: string) => {
-  const [, handlerFn] = (ipcMain.handle as Mock).mock.calls.find(([ch]) => ch === channel) || [];
+  const [, handlerFn] = vi.mocked(ipcMain.handle).mock.calls.find(([ch]) => ch === channel) || [];
   return handlerFn;
 };
 
 describe('AppHandlers', () => {
-  const testCases: TestCase[] = [{ channel: IPC_CHANNELS.QUIT, expected: undefined }];
-
   beforeEach(() => {
     registerAppHandlers();
   });
@@ -35,24 +21,24 @@ describe('AppHandlers', () => {
   });
 
   describe('registerHandlers', () => {
-    it.each(testCases)('should register handler for $channel', ({ channel }) => {
+    const channels = [IPC_CHANNELS.QUIT, IPC_CHANNELS.RESTART_APP];
+
+    test.each(channels)('should register handler for $channel', (channel) => {
       expect(ipcMain.handle).toHaveBeenCalledWith(channel, expect.any(Function));
     });
+  });
 
-    it.each(testCases)(
-      '$channel handler should return mock value ($expected)',
-      async ({ channel, expected, args = [] }) => {
-        const handlerFn = getHandler(channel);
-        const result = await handlerFn(...args);
+  test('restart handler should call app.relaunch', async () => {
+    expect(ipcMain.handle).toHaveBeenCalledWith(IPC_CHANNELS.RESTART_APP, expect.any(Function));
 
-        expect(result).toEqual(expected);
-      }
-    );
+    const handlerFn = getHandler(IPC_CHANNELS.RESTART_APP);
+    await expect(handlerFn).rejects.toThrow(/^Cannot destructure property 'customMessage' of/);
+    await expect(handlerFn?.(null!, [{}])).rejects.toThrow(quitMessage);
+    expect(app.relaunch).toHaveBeenCalledTimes(1);
   });
 
-  it('quit handler should call app.quit', async () => {
+  test('quit handler should call app.quit', () => {
     const handlerFn = getHandler(IPC_CHANNELS.QUIT);
-    await handlerFn();
-    expect(app.quit).toHaveBeenCalled();
+    expect(handlerFn).toThrow(quitMessage);
   });
 });

tests/unit/handlers/appinfoHandlers.test.ts

@@ -8,18 +8,6 @@ const MOCK_WINDOW_STYLE = 'default';
 const MOCK_GPU_NAME = 'mock-gpu';
 const MOCK_BASE_PATH = '/set/user/changed/base/path';
 
-vi.mock('electron', () => ({
-  app: {
-    isPackaged: false,
-    getPath: vi.fn(),
-    getVersion: vi.fn(() => '1.0.0'),
-  },
-  ipcMain: {
-    on: vi.fn(),
-    handle: vi.fn(),
-  },
-}));
-
 vi.mock('@/store/desktopConfig', () => ({
   useDesktopConfig: vi.fn(() => ({
     get: vi.fn((key: string) => {
@@ -53,7 +41,7 @@ const getHandler = (channel: string) => {
 
 describe('AppInfoHandlers', () => {
   const testCases: TestCase[] = [
-    { channel: IPC_CHANNELS.IS_PACKAGED, expected: false },
+    { channel: IPC_CHANNELS.IS_PACKAGED, expected: true },
     { channel: IPC_CHANNELS.GET_ELECTRON_VERSION, expected: '1.0.0' },
     { channel: IPC_CHANNELS.GET_BASE_PATH, expected: MOCK_BASE_PATH },
     { channel: IPC_CHANNELS.GET_GPU, expected: MOCK_GPU_NAME },

tests/unit/handlers/pathHandler.test.ts

@@ -12,6 +12,8 @@ import { IPC_CHANNELS } from '@/constants';
 import { REQUIRED_SPACE, registerPathHandlers } from '@/handlers/pathHandlers';
 import type { SystemPaths } from '@/preload';
 
+import { electronMock } from '../setup';
+
 const DEFAULT_FREE_SPACE = 20 * 1024 * 1024 * 1024; // 20GB
 const LOW_FREE_SPACE = 5 * 1024 * 1024 * 1024; // 5GB
 
@@ -23,38 +25,23 @@ const MOCK_PATHS = {
   appPath: '/mock/app/path',
 } as const;
 
-vi.mock('electron', () => {
-  return {
-    ipcMain: {
-      on: vi.fn(),
-      handle: vi.fn(),
-    },
-    app: {
-      getPath: vi.fn((name: string): string => {
-        switch (name) {
-          case 'userData':
-            return '/mock/user/data';
-          case 'logs':
-            return '/mock/logs/path';
-          case 'documents':
-            return '/mock/documents';
-          case 'appData':
-            return '/mock/appData';
-          default:
-            return `/mock/${name}`;
-        }
-      }),
-      getAppPath: vi.fn().mockReturnValue('/mock/app/path'),
-    },
-    shell: {
-      openPath: vi.fn(),
-    },
-    dialog: {
-      showOpenDialog: vi.fn(),
-    },
-  };
+electronMock.app.getPath = vi.fn((name: string) => {
+  switch (name) {
+    case 'userData':
+      return '/mock/user/data';
+    case 'logs':
+      return '/mock/logs/path';
+    case 'documents':
+      return '/mock/documents';
+    case 'appData':
+      return '/mock/appData';
+    default:
+      return `/mock/${name}`;
+  }
 });
 
+electronMock.shell = { openPath: vi.fn() };
+
 vi.mock('systeminformation');
 vi.mock('node:fs');
 vi.mock('@/config/comfyServerConfig', () => ({

tests/unit/install/installWizard.test.ts

@@ -9,6 +9,7 @@ import { ComfySettings } from '../../../src/config/comfySettings';
 import { InstallWizard } from '../../../src/install/installWizard';
 import { InstallOptions } from '../../../src/preload';
 import { getTelemetry } from '../../../src/services/telemetry';
+import { electronMock } from '../setup';
 
 vi.mock('node:fs', () => ({
   default: {
@@ -32,23 +33,19 @@ vi.mock('node:fs/promises', () => ({
 
 vi.mock('../../../src/config/comfyConfigManager');
 vi.mock('../../../src/config/comfyServerConfig');
-vi.mock('electron', () => ({
-  app: {
-    isPackaged: true,
-    getPath: vi.fn((name: string) => {
-      switch (name) {
-        case 'userData':
-          return '/test/user/data';
-        case 'appData':
-          return '/test/app/data';
-        case 'temp':
-          return '/test/temp';
-        default:
-          return '/test/default';
-      }
-    }),
-  },
-}));
+
+electronMock.app.getPath = vi.fn((name: string) => {
+  switch (name) {
+    case 'userData':
+      return '/test/user/data';
+    case 'appData':
+      return '/test/app/data';
+    case 'temp':
+      return '/test/temp';
+    default:
+      return '/test/default';
+  }
+});
 
 // Mock process.resourcesPath since app.isPackaged is true
 vi.stubGlobal('process', {

tests/unit/install/installationManager.test.ts

@@ -12,17 +12,6 @@ import type { ITelemetry } from '@/services/telemetry';
 import { useDesktopConfig } from '@/store/desktopConfig';
 import * as utils from '@/utils';
 
-vi.mock('electron', () => ({
-  ipcMain: {
-    handle: vi.fn(),
-    removeHandler: vi.fn(),
-    on: vi.fn(),
-  },
-  app: {
-    getPath: vi.fn(() => 'valid/path'),
-  },
-}));
-
 vi.mock('node:fs/promises', () => ({
   default: {
     access: vi.fn(),

tests/unit/main-process/appState.test.ts

@@ -6,13 +6,6 @@ import type { Page } from '@/infrastructure/interfaces';
 // Clear global mock
 vi.unmock('@/main-process/appState');
 
-// Mock electron app
-vi.mock('electron', () => ({
-  app: {
-    once: vi.fn(),
-  },
-}));
-
 type AppStateModule = typeof import('@/main-process/appState');
 
 const test = baseTest.extend<AppStateModule & { imported: AppStateModule }>({

tests/unit/main-process/appWindow.test.ts

@@ -1,19 +1,12 @@
-import { BrowserWindow } from 'electron';
+import { BrowserWindow, type Tray } from 'electron';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { AppWindow } from '@/main-process/appWindow';
 
-vi.mock('electron', () => ({
-  BrowserWindow: vi.fn(),
-  app: {
-    // Required if process.resourcesPath is not mocked
-    isPackaged: false,
-    on: vi.fn(),
-  },
-  ipcMain: {
-    on: vi.fn(),
-    handle: vi.fn(),
-  },
+import { type PartialMock, electronMock } from '../setup';
+
+const additionalMocks: PartialMock<typeof Electron> = {
+  BrowserWindow: vi.fn() as PartialMock<BrowserWindow>,
   nativeTheme: {
     shouldUseDarkColors: true,
   },
@@ -26,13 +19,15 @@ vi.mock('electron', () => ({
     setPressedImage: vi.fn(),
     setToolTip: vi.fn(),
     on: vi.fn(),
-  })),
+  })) as PartialMock<Tray>,
   screen: {
     getPrimaryDisplay: vi.fn(() => ({
       workAreaSize: { width: 1024, height: 768 },
     })),
   },
-}));
+};
+
+Object.assign(electronMock, additionalMocks);
 
 vi.mock('electron-store', () => ({
   default: vi.fn(() => ({
@@ -60,6 +55,11 @@ describe('AppWindow.isOnPage', () => {
       setWindowOpenHandler: vi.fn(),
     };
 
+    vi.stubGlobal('process', {
+      ...process,
+      resourcesPath: '/mock/app/path/assets',
+    });
+
     vi.mocked(BrowserWindow).mockImplementation(
       () =>
         ({

tests/unit/main-process/comfyDesktopApp.test.ts

@@ -27,16 +27,6 @@ vi.mock('@/config/comfySettings', () => {
   };
 });
 
-vi.mock('electron', () => ({
-  app: {
-    on: vi.fn(),
-  },
-  ipcMain: {
-    handle: vi.fn(),
-    removeHandler: vi.fn(),
-  },
-}));
-
 vi.mock('@todesktop/runtime', () => ({
   default: {
     init: vi.fn(),

tests/unit/main-process/comfyServer.test.ts

@@ -12,11 +12,7 @@ import { ComfyServer } from '@/main-process/comfyServer';
 import { type ITelemetry, getTelemetry } from '@/services/telemetry';
 import type { VirtualEnvironment } from '@/virtualEnvironment';
 
-const basePath = '/mock/base/path';
-
-vi.mock('electron', () => ({
-  app: { getPath: vi.fn(() => basePath) },
-}));
+const basePath = '/mock/app/path';
 
 vi.mock('@/install/resourcePaths', () => ({
   getAppResourcesPath: vi.fn(() => '/mocked/app_resources'),