From 01697a0dcf7d49788abb3328ecf19c10d9e934f0 Mon Sep 17 00:00:00 2001 From: bymyself Date: Sat, 8 Feb 2025 22:03:47 -0700 Subject: [PATCH] Make comfy settings instance accessible from anywhere (#863) * Make comfy settings a singleton * update tests * remove unnecessary mock * load settings in fromConfig * suppress test directive * change to service locator pattern * update tests --- src/config/comfySettings.ts | 41 +++++- src/desktopApp.ts | 5 +- src/install/installWizard.ts | 29 ++-- src/install/installationManager.ts | 6 +- src/main-process/comfyDesktopApp.ts | 10 +- src/main-process/comfyInstallation.ts | 21 ++- src/services/telemetry.ts | 14 +- tests/unit/config/comfySettings.test.ts | 118 ++++++++-------- tests/unit/desktopApp.test.ts | 36 +++-- tests/unit/install/installWizard.test.ts | 128 +++++++++--------- .../unit/install/installationManager.test.ts | 72 ++++++---- .../unit/main-process/comfyDesktopApp.test.ts | 14 ++ tests/unit/services/telemetry.test.ts | 83 +++++++++--- 13 files changed, 351 insertions(+), 226 deletions(-) diff --git a/src/config/comfySettings.ts b/src/config/comfySettings.ts index e007dc11..8dee65cc 100644 --- a/src/config/comfySettings.ts +++ b/src/config/comfySettings.ts @@ -29,6 +29,15 @@ export interface ComfySettingsData { [key: string]: unknown; } +/** Backing ref for the singleton settings instance. */ +let current: ComfySettings; + +/** Service locator for settings. ComfySettings.load() must be called before access. */ +export function useComfySettings() { + if (!current) throw new Error('Cannot access ComfySettings before initialization.'); + return current; +} + /** * A read-only interface to an in-memory cache of frontend settings. * @see {@link ComfySettings} concrete implementation @@ -75,12 +84,12 @@ export interface IComfySettings extends FrontendSettingsCache { * @see {@link IComfySettings} read-write interface */ export class ComfySettings implements IComfySettings { - public readonly filePath: string; private settings: ComfySettingsData = structuredClone(DEFAULT_SETTINGS); private static writeLocked = false; + readonly #basePath: string; - constructor(basePath: string) { - this.filePath = path.join(basePath, 'user', 'default', 'comfy.settings.json'); + private constructor(basePath: string) { + this.#basePath = basePath; } /** @@ -91,7 +100,11 @@ export class ComfySettings implements IComfySettings { ComfySettings.writeLocked = true; } - public async loadSettings() { + get filePath(): string { + return path.join(this.#basePath, 'user', 'default', 'comfy.settings.json'); + } + + private async loadSettings() { try { await fs.access(this.filePath); } catch { @@ -102,9 +115,13 @@ export class ComfySettings implements IComfySettings { const fileContent = await fs.readFile(this.filePath, 'utf8'); // TODO: Reimplement with validation and error reporting. // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment - this.settings = JSON.parse(fileContent); + this.settings = { ...this.settings, ...JSON.parse(fileContent) }; } catch (error) { - log.error(`Settings file cannot be loaded.`, error); + if (error instanceof SyntaxError) { + log.error(`Settings file contains invalid JSON:`, error); + } else { + log.error(`Settings file cannot be loaded.`, error); + } } } @@ -135,4 +152,16 @@ export class ComfySettings implements IComfySettings { get(key: K): ComfySettingsData[K] { return this.settings[key] ?? DEFAULT_SETTINGS[key]; } + + /** + * Static factory method. Loads the settings from disk. + * @param basePath The base path where ComfyUI is installed + * @returns The newly created instance + */ + static async load(basePath: string): Promise { + const instance = new ComfySettings(basePath); + await instance.loadSettings(); + current = instance; + return instance; + } } diff --git a/src/desktopApp.ts b/src/desktopApp.ts index e6011bc6..b6bdac4c 100644 --- a/src/desktopApp.ts +++ b/src/desktopApp.ts @@ -44,13 +44,10 @@ export class DesktopApp implements HasTelemetry { } private async initializeTelemetry(installation: ComfyInstallation): Promise { - const { comfySettings } = installation; - await SentryLogging.setSentryGpuContext(); - SentryLogging.shouldSendStatistics = () => comfySettings.get('Comfy-Desktop.SendStatistics'); SentryLogging.getBasePath = () => installation.basePath; - const allowMetrics = await promptMetricsConsent(this.config, this.appWindow, comfySettings); + const allowMetrics = await promptMetricsConsent(this.config, this.appWindow); this.telemetry.hasConsent = allowMetrics; if (allowMetrics) this.telemetry.flush(); } diff --git a/src/install/installWizard.ts b/src/install/installWizard.ts index beebf65a..ed26008a 100644 --- a/src/install/installWizard.ts +++ b/src/install/installWizard.ts @@ -4,7 +4,7 @@ import path from 'node:path'; import { ComfyConfigManager } from '../config/comfyConfigManager'; import { ComfyServerConfig, ModelPaths } from '../config/comfyServerConfig'; -import { type ComfySettingsData, DEFAULT_SETTINGS } from '../config/comfySettings'; +import { ComfySettings, type ComfySettingsData } from '../config/comfySettings'; import { InstallOptions } from '../preload'; import { HasTelemetry, ITelemetry, trackEvent } from '../services/telemetry'; @@ -31,7 +31,7 @@ export class InstallWizard implements HasTelemetry { // Setup the ComfyUI folder structure. ComfyConfigManager.createComfyDirectories(this.basePath); this.initializeUserFiles(); - this.initializeSettings(); + await this.initializeSettings(); await this.initializeModelPaths(); } @@ -52,16 +52,12 @@ export class InstallWizard implements HasTelemetry { /** * Setup comfy.settings.json file */ - public initializeSettings() { - const settingsPath = path.join(this.basePath, 'user', 'default', 'comfy.settings.json'); - // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment - const existingSettings: ComfySettingsData = fs.existsSync(settingsPath) - ? JSON.parse(fs.readFileSync(settingsPath, 'utf8')) - : {}; - - const settings = { - ...DEFAULT_SETTINGS, - ...existingSettings, + public async initializeSettings() { + // Load any existing settings if they exist + const existingSettings = await ComfySettings.load(this.basePath); + + // Add install options to settings + const settings: Partial = { 'Comfy-Desktop.AutoUpdate': this.installOptions.autoUpdate, 'Comfy-Desktop.SendStatistics': this.installOptions.allowMetrics, 'Comfy-Desktop.UV.PythonInstallMirror': this.installOptions.pythonMirror, @@ -74,9 +70,12 @@ export class InstallWizard implements HasTelemetry { settings['Comfy.Server.LaunchArgs']['cpu'] = ''; } - const settingsJson = JSON.stringify(settings, null, 2); - fs.writeFileSync(settingsPath, settingsJson); - log.info(`Wrote settings to ${settingsPath}: ${settingsJson}`); + for (const [key, value] of Object.entries(settings)) { + existingSettings.set(key, value); + } + + await existingSettings.saveSettings(); + log.info(`Wrote install options to comfy settings file.`); } /** diff --git a/src/install/installationManager.ts b/src/install/installationManager.ts index 096c0938..30cfcc46 100644 --- a/src/install/installationManager.ts +++ b/src/install/installationManager.ts @@ -1,8 +1,6 @@ import { Notification, app, dialog, ipcMain, shell } from 'electron'; import log from 'electron-log/main'; -import { ComfySettings } from '@/config/comfySettings'; - import { IPC_CHANNELS, ProgressStatus } from '../constants'; import type { AppWindow } from '../main-process/appWindow'; import { ComfyInstallation } from '../main-process/comfyInstallation'; @@ -238,9 +236,7 @@ export class InstallationManager implements HasTelemetry { useDesktopConfig().set('migrateCustomNodesFrom', installWizard.migrationSource); } - const comfySettings = new ComfySettings(installWizard.basePath); - await comfySettings.loadSettings(); - const installation = new ComfyInstallation('started', installWizard.basePath, this.telemetry, comfySettings); + const installation = new ComfyInstallation('started', installWizard.basePath, this.telemetry); InstallationManager.setReinstallHandler(installation); const { virtualEnvironment } = installation; diff --git a/src/main-process/comfyDesktopApp.ts b/src/main-process/comfyDesktopApp.ts index 21c323ad..e46deb22 100644 --- a/src/main-process/comfyDesktopApp.ts +++ b/src/main-process/comfyDesktopApp.ts @@ -2,6 +2,8 @@ import todesktop from '@todesktop/runtime'; import { app, ipcMain } from 'electron'; import log from 'electron-log/main'; +import { useComfySettings } from '@/config/comfySettings'; + import { DEFAULT_SERVER_ARGS, IPC_CHANNELS, ProgressStatus, ServerArgs } from '../constants'; import { DownloadManager } from '../models/DownloadManager'; import { HasTelemetry, ITelemetry } from '../services/telemetry'; @@ -22,10 +24,6 @@ export class ComfyDesktopApp implements HasTelemetry { readonly telemetry: ITelemetry ) {} - get comfySettings() { - return this.installation.comfySettings; - } - get basePath() { return this.installation.basePath; } @@ -47,7 +45,7 @@ export class ComfyDesktopApp implements HasTelemetry { const serverArgs: ServerArgs = { listen: DEFAULT_SERVER_ARGS.listen, port: DEFAULT_SERVER_ARGS.port, - ...this.comfySettings.get('Comfy.Server.LaunchArgs'), + ...useComfySettings().get('Comfy.Server.LaunchArgs'), }; if (COMFY_HOST) serverArgs.listen = COMFY_HOST; @@ -69,7 +67,7 @@ export class ComfyDesktopApp implements HasTelemetry { autoCheckInterval: 60 * 60 * 1000, // every hour customLogger: log, updateReadyAction: { showInstallAndRestartPrompt: 'always', showNotification: 'always' }, - autoUpdater: this.comfySettings.get('Comfy-Desktop.AutoUpdate'), + autoUpdater: useComfySettings().get('Comfy-Desktop.AutoUpdate'), }); todesktop.autoUpdater?.setFeedURL('https://updater.comfy.org'); } diff --git a/src/main-process/comfyInstallation.ts b/src/main-process/comfyInstallation.ts index 4ad466cb..d7db6cad 100644 --- a/src/main-process/comfyInstallation.ts +++ b/src/main-process/comfyInstallation.ts @@ -2,7 +2,7 @@ import log from 'electron-log/main'; import { rm } from 'node:fs/promises'; import { ComfyServerConfig } from '../config/comfyServerConfig'; -import { ComfySettings } from '../config/comfySettings'; +import { ComfySettings, useComfySettings } from '../config/comfySettings'; import type { DesktopInstallState } from '../main_types'; import type { InstallValidation } from '../preload'; import { type ITelemetry, getTelemetry } from '../services/telemetry'; @@ -52,8 +52,7 @@ export class ComfyInstallation { /** The base path of the desktop app. Models, nodes, and configuration are saved here by default. */ basePath: string, /** The device type to use for the installation. */ - public readonly telemetry: ITelemetry, - public comfySettings: ComfySettings + public readonly telemetry: ITelemetry ) { this._basePath = basePath; this._virtualEnvironment = this.createVirtualEnvironment(basePath); @@ -63,9 +62,9 @@ export class ComfyInstallation { return new VirtualEnvironment(basePath, { telemetry: this.telemetry, selectedDevice: useDesktopConfig().get('selectedDevice'), - pythonMirror: this.comfySettings.get('Comfy-Desktop.UV.PythonInstallMirror'), - pypiMirror: this.comfySettings.get('Comfy-Desktop.UV.PypiInstallMirror'), - torchMirror: this.comfySettings.get('Comfy-Desktop.UV.TorchInstallMirror'), + pythonMirror: useComfySettings().get('Comfy-Desktop.UV.PythonInstallMirror'), + pypiMirror: useComfySettings().get('Comfy-Desktop.UV.PypiInstallMirror'), + torchMirror: useComfySettings().get('Comfy-Desktop.UV.TorchInstallMirror'), }); } @@ -79,9 +78,8 @@ export class ComfyInstallation { const state = config.get('installState'); const basePath = config.get('basePath'); if (state && basePath) { - const comfySettings = new ComfySettings(basePath); - await comfySettings.loadSettings(); - return new ComfyInstallation(state, basePath, getTelemetry(), comfySettings); + await ComfySettings.load(basePath); + return new ComfyInstallation(state, basePath, getTelemetry()); } } @@ -215,9 +213,10 @@ export class ComfyInstallation { this._basePath = basePath; this._virtualEnvironment = this.createVirtualEnvironment(basePath); - this.comfySettings = new ComfySettings(basePath); - await this.comfySettings.loadSettings(); useDesktopConfig().set('basePath', basePath); + + // If settings file exists at new location, load it + await ComfySettings.load(basePath); } /** diff --git a/src/services/telemetry.ts b/src/services/telemetry.ts index e2e040f1..f469776f 100644 --- a/src/services/telemetry.ts +++ b/src/services/telemetry.ts @@ -7,7 +7,7 @@ import os from 'node:os'; import path from 'node:path'; import si from 'systeminformation'; -import type { IComfySettings } from '@/config/comfySettings'; +import { useComfySettings } from '@/config/comfySettings'; import { IPC_CHANNELS } from '../constants'; import { AppWindow } from '../main-process/appWindow'; @@ -241,12 +241,8 @@ export function trackEvent(eventName: string) { } /** @returns Whether the user has consented to sending metrics. */ -export async function promptMetricsConsent( - store: DesktopConfig, - appWindow: AppWindow, - comfySettings: IComfySettings -): Promise { - const consent = comfySettings.get('Comfy-Desktop.SendStatistics') ?? false; +export async function promptMetricsConsent(store: DesktopConfig, appWindow: AppWindow): Promise { + const consent = useComfySettings().get('Comfy-Desktop.SendStatistics') ?? false; const consentedOn = store.get('versionConsentedMetrics'); const isOutdated = !consentedOn || compareVersions(consentedOn, '0.4.12') < 0; if (!isOutdated) return consent; @@ -260,8 +256,8 @@ export async function promptMetricsConsent( await appWindow.loadPage('metrics-consent'); const newConsent = await consentPromise; if (newConsent !== consent) { - comfySettings.set('Comfy-Desktop.SendStatistics', newConsent); - await comfySettings.saveSettings(); + useComfySettings().set('Comfy-Desktop.SendStatistics', newConsent); + await useComfySettings().saveSettings(); } return newConsent; diff --git a/tests/unit/config/comfySettings.test.ts b/tests/unit/config/comfySettings.test.ts index e30f3502..3bfc9891 100644 --- a/tests/unit/config/comfySettings.test.ts +++ b/tests/unit/config/comfySettings.test.ts @@ -1,8 +1,9 @@ -import fs from 'node:fs/promises'; +import log from 'electron-log/main'; +import fsPromises from 'node:fs/promises'; import path from 'node:path'; import { beforeEach, describe, expect, it, vi } from 'vitest'; -import { ComfySettings, type ComfySettingsData, DEFAULT_SETTINGS } from '@/config/comfySettings'; +import { ComfySettings, type ComfySettingsData, DEFAULT_SETTINGS, useComfySettings } from '@/config/comfySettings'; vi.mock('electron-log/main', () => ({ default: { @@ -17,33 +18,42 @@ vi.mock('node:fs/promises', () => ({ readFile: vi.fn(), writeFile: vi.fn(), }, + access: vi.fn(), + readFile: vi.fn(), + writeFile: vi.fn(), })); -async function expectLogError() { - const log = await import('electron-log/main'); - expect(vi.mocked(log.default.error)).toHaveBeenCalled(); -} - describe('ComfySettings', () => { + const basePath = path.join('test', 'base', 'path'); + const expectedFilePath = path.join(basePath, 'user', 'default', 'comfy.settings.json'); let settings: ComfySettings; - const testBasePath = '/test/path'; - beforeEach(() => { - settings = new ComfySettings(testBasePath); - ComfySettings['writeLocked'] = false; + beforeEach(async () => { + vi.resetModules(); vi.clearAllMocks(); + + // Reset writeLocked state + // @ts-expect-error accessing private static + ComfySettings.writeLocked = false; + + // Reset fs mocks with default behaviors + vi.mocked(fsPromises.access).mockResolvedValue(undefined); + vi.mocked(fsPromises.readFile).mockResolvedValue('{}'); + vi.mocked(fsPromises.writeFile).mockResolvedValue(undefined); + + settings = await ComfySettings.load(basePath); }); describe('write locking', () => { it('should allow writes before being locked', async () => { await settings.saveSettings(); - expect(vi.mocked(fs).writeFile).toHaveBeenCalled(); + expect(fsPromises.writeFile).toHaveBeenCalledWith(expectedFilePath, JSON.stringify(DEFAULT_SETTINGS, null, 2)); }); it('should prevent writes after being locked', async () => { ComfySettings.lockWrites(); await expect(settings.saveSettings()).rejects.toThrow('Settings are locked'); - expect(fs.writeFile).not.toHaveBeenCalled(); + expect(fsPromises.writeFile).not.toHaveBeenCalled(); }); it('should prevent modifications after being locked', () => { @@ -56,9 +66,9 @@ describe('ComfySettings', () => { expect(() => settings.get('Comfy-Desktop.AutoUpdate')).not.toThrow(); }); - it('should share lock state across instances', () => { - const settings1 = new ComfySettings('/path1'); - const settings2 = new ComfySettings('/path2'); + it('should share lock state across references', async () => { + const settings1 = settings; + const settings2 = await ComfySettings.load(basePath); ComfySettings.lockWrites(); @@ -66,16 +76,20 @@ describe('ComfySettings', () => { expect(() => settings2.set('Comfy-Desktop.AutoUpdate', false)).toThrow('Settings are locked'); }); - it('should log error when saving locked settings', async () => { + it('should throw error when saving locked settings', async () => { ComfySettings.lockWrites(); await expect(settings.saveSettings()).rejects.toThrow('Settings are locked'); - await expectLogError(); }); }); describe('file operations', () => { - it('should use correct file path', () => { - expect(settings.filePath).toBe(path.join(testBasePath, 'user', 'default', 'comfy.settings.json')); + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('should use correct file path', async () => { + await settings.saveSettings(); + expect(fsPromises.writeFile).toHaveBeenCalledWith(expectedFilePath, JSON.stringify(DEFAULT_SETTINGS, null, 2)); }); it('should load settings from file when available', async () => { @@ -92,41 +106,40 @@ describe('ComfySettings', () => { 'Comfy-Desktop.UV.TorchInstallMirror': '', }; - vi.mocked(fs.access).mockResolvedValue(); - vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(mockSettings)); + vi.mocked(fsPromises.access).mockResolvedValue(undefined); + vi.mocked(fsPromises.readFile).mockResolvedValue(JSON.stringify(mockSettings)); - await settings.loadSettings(); + settings = await ComfySettings.load(basePath); expect(settings.get('Comfy-Desktop.AutoUpdate')).toBe(false); expect(settings.get('Comfy.Server.LaunchArgs')).toEqual({ test: 'value' }); expect(settings.get('Comfy-Desktop.SendStatistics')).toBe(false); }); it('should use default settings when file does not exist', async () => { - vi.mocked(fs.access).mockRejectedValue(new Error('ENOENT')); - - await settings.loadSettings(); - expect(settings.get('Comfy-Desktop.AutoUpdate')).toBe(true); - expect(settings.get('Comfy-Desktop.SendStatistics')).toBe(true); + vi.mocked(fsPromises.access).mockRejectedValue(new Error('ENOENT')); + settings = await ComfySettings.load(basePath); + expect(settings.get('Comfy-Desktop.AutoUpdate')).toBe(DEFAULT_SETTINGS['Comfy-Desktop.AutoUpdate']); }); it('should save settings to correct path with proper formatting', async () => { settings.set('Comfy-Desktop.AutoUpdate', false); await settings.saveSettings(); - const writeCall = vi.mocked(fs).writeFile.mock.calls[0]; + const writeCall = vi.mocked(fsPromises.writeFile).mock.calls.at(-1); + if (!writeCall) throw new Error('No write calls recorded'); const savedJson = JSON.parse(writeCall[1] as string); - expect(writeCall[0]).toBe(settings.filePath); + expect(writeCall[0]).toBe(expectedFilePath); expect(savedJson['Comfy-Desktop.AutoUpdate']).toBe(false); }); it('should fall back to defaults on file read error', async () => { - vi.mocked(fs.access).mockResolvedValue(); - vi.mocked(fs.readFile).mockRejectedValue(new Error('Permission denied')); + vi.mocked(fsPromises.access).mockResolvedValue(undefined); + vi.mocked(fsPromises.readFile).mockRejectedValue(new Error('Permission denied')); - await settings.loadSettings(); - await expectLogError(); + settings = await ComfySettings.load(basePath); expect(settings.get('Comfy-Desktop.AutoUpdate')).toBe(DEFAULT_SETTINGS['Comfy-Desktop.AutoUpdate']); + expect(log.error).toHaveBeenCalled(); }); }); @@ -138,8 +151,8 @@ describe('ComfySettings', () => { }); it('should preserve primitive and object types when getting/setting values', () => { - settings.set('Comfy-Desktop.AutoUpdate', false); - expect(typeof settings.get('Comfy-Desktop.AutoUpdate')).toBe('boolean'); + settings.set('Comfy-Desktop.SendStatistics', false); + expect(typeof settings.get('Comfy-Desktop.SendStatistics')).toBe('boolean'); const serverArgs = { test: 'value' }; settings.set('Comfy.Server.LaunchArgs', serverArgs); @@ -152,37 +165,32 @@ describe('ComfySettings', () => { 'Comfy.Server.LaunchArgs': null, }; - vi.mocked(fs.access).mockResolvedValue(); - vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(invalidSettings)); + vi.mocked(fsPromises.access).mockResolvedValue(undefined); + vi.mocked(fsPromises.readFile).mockResolvedValue(JSON.stringify(invalidSettings)); - await settings.loadSettings(); + settings = await ComfySettings.load(basePath); expect(settings.get('Comfy-Desktop.AutoUpdate')).toBe(DEFAULT_SETTINGS['Comfy-Desktop.AutoUpdate']); expect(settings.get('Comfy.Server.LaunchArgs')).toEqual(DEFAULT_SETTINGS['Comfy.Server.LaunchArgs']); }); it('should fall back to defaults when settings file contains invalid JSON', async () => { - vi.mocked(fs.access).mockResolvedValue(); - vi.mocked(fs.readFile).mockResolvedValue('{ invalid json }'); + vi.mocked(fsPromises.access).mockResolvedValue(undefined); + vi.mocked(fsPromises.readFile).mockRejectedValue(new Error('Invalid JSON')); - await settings.loadSettings(); - await expectLogError(); + settings = await ComfySettings.load(basePath); expect(settings.get('Comfy-Desktop.AutoUpdate')).toBe(DEFAULT_SETTINGS['Comfy-Desktop.AutoUpdate']); }); - it('should handle attempts to save null settings', async () => { - const saveSettingsSpy = vi.spyOn(settings, 'saveSettings'); - // @ts-expect-error: explicitly setting settings to null - settings['settings'] = null; - await settings.saveSettings(); - - expect(saveSettingsSpy).toHaveReturned(); - expect(fs.writeFile).not.toHaveBeenCalled(); + it('should throw error on write error during saveSettings', async () => { + vi.mocked(fsPromises.writeFile).mockRejectedValue(new Error('Permission denied')); + await expect(settings.saveSettings()).rejects.toThrow('Permission denied'); }); + }); - it('should log and throw error on write error during saveSettings', async () => { - vi.mocked(fs.writeFile).mockRejectedValue(new Error('Permission denied')); - await expect(settings.saveSettings()).rejects.toThrow('Permission denied'); - await expectLogError(); + describe('useComfySettings', () => { + it('should return the current instance after initialization', async () => { + settings = await ComfySettings.load(basePath); + expect(useComfySettings()).toBe(settings); }); }); }); diff --git a/tests/unit/desktopApp.test.ts b/tests/unit/desktopApp.test.ts index 1690f9fb..319c8214 100644 --- a/tests/unit/desktopApp.test.ts +++ b/tests/unit/desktopApp.test.ts @@ -2,7 +2,7 @@ import { app, dialog } from 'electron'; import log from 'electron-log/main'; import { beforeEach, describe, expect, it, vi } from 'vitest'; -import { ComfySettings } from '@/config/comfySettings'; +import { useComfySettings } from '@/config/comfySettings'; import { ProgressStatus } from '@/constants'; import { IPC_CHANNELS } from '@/constants'; import { DesktopApp } from '@/desktopApp'; @@ -57,16 +57,32 @@ vi.mock('@/main-process/appWindow', () => ({ AppWindow: vi.fn().mockImplementation(() => mockAppWindow), })); -const mockComfySettings = new ComfySettings('/mock/path'); -mockComfySettings.get = vi.fn().mockReturnValue('true'); -mockComfySettings.set = vi.fn(); -mockComfySettings.saveSettings = vi.fn(); +vi.mock('@/config/comfySettings', () => ({ + ComfySettings: { + load: vi.fn().mockResolvedValue({ + get: vi.fn().mockReturnValue('true'), + set: vi.fn(), + saveSettings: vi.fn(), + }), + }, + useComfySettings: vi.fn().mockReturnValue({ + get: vi.fn().mockReturnValue('true'), + set: vi.fn(), + saveSettings: vi.fn(), + }), +})); + +vi.mock('@/store/desktopConfig', () => ({ + useDesktopConfig: vi.fn().mockReturnValue({ + get: vi.fn().mockReturnValue('/mock/path'), + set: vi.fn(), + }), +})); const mockInstallation: Partial = { basePath: '/mock/path', virtualEnvironment: {} as any, validation: {} as any, - comfySettings: mockComfySettings, hasIssues: false, isValid: true, state: 'installed', @@ -223,10 +239,11 @@ describe('DesktopApp', () => { it('should initialize telemetry with user consent', async () => { vi.mocked(promptMetricsConsent).mockResolvedValueOnce(true); vi.mocked(mockConfig.get).mockReturnValue('true'); + vi.mocked(useComfySettings().get).mockReturnValue('true'); await desktopApp['initializeTelemetry'](testInstallation); - expect(promptMetricsConsent).toHaveBeenCalledWith(mockConfig, mockAppWindow, testInstallation.comfySettings); + expect(promptMetricsConsent).toHaveBeenCalledWith(mockConfig, mockAppWindow); expect(SentryLogging.setSentryGpuContext).toHaveBeenCalled(); expect(desktopApp.telemetry.hasConsent).toBe(true); expect(desktopApp.telemetry.flush).toHaveBeenCalled(); @@ -235,11 +252,12 @@ describe('DesktopApp', () => { it('should respect user rejection of telemetry', async () => { vi.mocked(promptMetricsConsent).mockResolvedValueOnce(false); vi.mocked(mockConfig.get).mockReturnValue('false'); - vi.mocked(testInstallation.comfySettings.get).mockReturnValue('false'); + vi.mocked(useComfySettings().get).mockReturnValue('false'); await desktopApp['initializeTelemetry'](testInstallation); - expect(promptMetricsConsent).toHaveBeenCalledWith(mockConfig, mockAppWindow, testInstallation.comfySettings); + expect(promptMetricsConsent).toHaveBeenCalledWith(mockConfig, mockAppWindow); + expect(SentryLogging.setSentryGpuContext).toHaveBeenCalled(); expect(desktopApp.telemetry.hasConsent).toBe(false); expect(desktopApp.telemetry.flush).not.toHaveBeenCalled(); }); diff --git a/tests/unit/install/installWizard.test.ts b/tests/unit/install/installWizard.test.ts index 7a24dff9..aff942e1 100644 --- a/tests/unit/install/installWizard.test.ts +++ b/tests/unit/install/installWizard.test.ts @@ -1,15 +1,35 @@ import fs from 'node:fs'; +import fsPromises from 'node:fs/promises'; import path from 'node:path'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { ComfyConfigManager } from '../../../src/config/comfyConfigManager'; import { ComfyServerConfig, ModelPaths } from '../../../src/config/comfyServerConfig'; -import { DEFAULT_SETTINGS } from '../../../src/config/comfySettings'; +import { ComfySettings } from '../../../src/config/comfySettings'; import { InstallWizard } from '../../../src/install/installWizard'; import { InstallOptions } from '../../../src/preload'; import { ITelemetry } from '../../../src/services/telemetry'; -vi.mock('node:fs'); +vi.mock('node:fs', () => ({ + default: { + cpSync: vi.fn(), + existsSync: vi.fn(), + }, + cpSync: vi.fn(), + existsSync: vi.fn(), +})); + +vi.mock('node:fs/promises', () => ({ + default: { + access: vi.fn(), + readFile: vi.fn(), + writeFile: vi.fn(), + }, + access: vi.fn(), + readFile: vi.fn(), + writeFile: vi.fn(), +})); + vi.mock('electron-log/main'); vi.mock('../../../src/config/comfyConfigManager'); vi.mock('../../../src/config/comfyServerConfig'); @@ -65,8 +85,9 @@ describe('InstallWizard', () => { torchMirror: 'default', }; - beforeEach(() => { + beforeEach(async () => { vi.clearAllMocks(); + await ComfySettings.load('/test/path'); installWizard = new InstallWizard(defaultInstallOptions, mockTelemetry); }); @@ -111,54 +132,48 @@ describe('InstallWizard', () => { }); describe('initializeSettings', () => { - it('should create settings file with default values when no existing settings', () => { - vi.spyOn(fs, 'existsSync').mockReturnValue(false); - const writeFileSpy = vi.spyOn(fs, 'writeFileSync'); - - installWizard.initializeSettings(); - - const expectedSettings = { - ...DEFAULT_SETTINGS, - 'Comfy-Desktop.AutoUpdate': true, - 'Comfy-Desktop.SendStatistics': true, - 'Comfy-Desktop.UV.PythonInstallMirror': 'default', - 'Comfy-Desktop.UV.PypiInstallMirror': 'default', - 'Comfy-Desktop.UV.TorchInstallMirror': 'default', - }; - - expect(writeFileSpy).toHaveBeenCalledWith( - path.join('/test/path', 'user', 'default', 'comfy.settings.json'), - JSON.stringify(expectedSettings, null, 2) - ); + it('should create settings file with default values when no existing settings', async () => { + // Mock fs to simulate no existing settings + vi.mocked(fs.existsSync).mockReturnValue(false); + vi.mocked(fsPromises.access).mockRejectedValue(new Error('ENOENT')); + vi.mocked(fsPromises.readFile).mockResolvedValue('{}'); + + await installWizard.initializeSettings(); + + // Verify settings were saved + expect(fsPromises.writeFile).toHaveBeenCalled(); + const savedSettings = JSON.parse(vi.mocked(fsPromises.writeFile).mock.calls[0][1] as string); + expect(savedSettings['Comfy-Desktop.AutoUpdate']).toBe(true); + expect(savedSettings['Comfy-Desktop.SendStatistics']).toBe(true); + expect(savedSettings['Comfy-Desktop.UV.PythonInstallMirror']).toBe('default'); + expect(savedSettings['Comfy-Desktop.UV.PypiInstallMirror']).toBe('default'); + expect(savedSettings['Comfy-Desktop.UV.TorchInstallMirror']).toBe('default'); }); - it('should merge with existing settings when settings file exists', () => { + it('should merge with existing settings when settings file exists', async () => { + // Mock fs to simulate existing settings const existingSettings = { 'Existing.Setting': 'value', + 'Comfy.ColorPalette': 'light', + 'Comfy.Server.LaunchArgs': { existingArg: true }, }; - vi.spyOn(fs, 'existsSync').mockReturnValue(true); - vi.spyOn(fs, 'readFileSync').mockReturnValue(JSON.stringify(existingSettings)); - const writeFileSpy = vi.spyOn(fs, 'writeFileSync'); - - installWizard.initializeSettings(); - - const expectedSettings = { - ...DEFAULT_SETTINGS, - ...existingSettings, - 'Comfy-Desktop.AutoUpdate': true, - 'Comfy-Desktop.SendStatistics': true, - 'Comfy-Desktop.UV.PythonInstallMirror': 'default', - 'Comfy-Desktop.UV.PypiInstallMirror': 'default', - 'Comfy-Desktop.UV.TorchInstallMirror': 'default', - }; - - expect(writeFileSpy).toHaveBeenCalledWith( - path.join('/test/path', 'user', 'default', 'comfy.settings.json'), - JSON.stringify(expectedSettings, null, 2) - ); + vi.mocked(fs.existsSync).mockReturnValue(true); + vi.mocked(fsPromises.access).mockResolvedValue(undefined); + vi.mocked(fsPromises.readFile).mockResolvedValue(JSON.stringify(existingSettings)); + + await installWizard.initializeSettings(); + + // Verify settings were merged and saved + expect(fsPromises.writeFile).toHaveBeenCalled(); + const savedSettings = JSON.parse(vi.mocked(fsPromises.writeFile).mock.calls[0][1] as string); + expect(savedSettings['Existing.Setting']).toBe('value'); + expect(savedSettings['Comfy.ColorPalette']).toBe('light'); + expect(savedSettings['Comfy.Server.LaunchArgs']).toEqual({ existingArg: true }); + expect(savedSettings['Comfy-Desktop.AutoUpdate']).toBe(true); + expect(savedSettings['Comfy-Desktop.SendStatistics']).toBe(true); }); - it('should add CPU launch args when device is cpu', () => { + it('should add CPU launch args when device is cpu', async () => { const wizardWithCpu = new InstallWizard( { ...defaultInstallOptions, @@ -167,27 +182,12 @@ describe('InstallWizard', () => { mockTelemetry ); - vi.spyOn(fs, 'existsSync').mockReturnValue(false); - const writeFileSpy = vi.spyOn(fs, 'writeFileSync'); - - wizardWithCpu.initializeSettings(); + await wizardWithCpu.initializeSettings(); - const expectedSettings = { - ...DEFAULT_SETTINGS, - 'Comfy-Desktop.AutoUpdate': true, - 'Comfy-Desktop.SendStatistics': true, - 'Comfy-Desktop.UV.PythonInstallMirror': 'default', - 'Comfy-Desktop.UV.PypiInstallMirror': 'default', - 'Comfy-Desktop.UV.TorchInstallMirror': 'default', - 'Comfy.Server.LaunchArgs': { - cpu: '', - }, - }; - - expect(writeFileSpy).toHaveBeenCalledWith( - path.join('/test/path', 'user', 'default', 'comfy.settings.json'), - JSON.stringify(expectedSettings, null, 2) - ); + // Verify CPU settings were saved + expect(fsPromises.writeFile).toHaveBeenCalled(); + const savedSettings = JSON.parse(vi.mocked(fsPromises.writeFile).mock.calls[0][1] as string); + expect(savedSettings['Comfy.Server.LaunchArgs']).toEqual({ cpu: '' }); }); }); diff --git a/tests/unit/install/installationManager.test.ts b/tests/unit/install/installationManager.test.ts index 094f57d8..07b26a78 100644 --- a/tests/unit/install/installationManager.test.ts +++ b/tests/unit/install/installationManager.test.ts @@ -1,3 +1,4 @@ +import fsPromises from 'node:fs/promises'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { ComfyServerConfig } from '@/config/comfyServerConfig'; @@ -15,11 +16,20 @@ vi.mock('electron', () => ({ ipcMain: { handle: vi.fn(), removeHandler: vi.fn(), + on: vi.fn(), + }, + app: { + getPath: vi.fn().mockReturnValue('valid/path'), }, })); vi.mock('node:fs/promises', () => ({ - rm: vi.fn(), + default: { + access: vi.fn(), + readFile: vi.fn().mockResolvedValue('{}'), + }, + access: vi.fn(), + readFile: vi.fn().mockResolvedValue('{}'), })); vi.mock('@/store/desktopConfig', () => ({ @@ -77,6 +87,14 @@ vi.mock('@/virtualEnvironment', () => { }; }); +// Mock Telemetry +vi.mock('@/services/telemetry', () => ({ + getTelemetry: vi.fn().mockReturnValue({ + track: vi.fn(), + }), + trackEvent: () => (target: any, propertyKey: string, descriptor: PropertyDescriptor) => descriptor, +})); + const createMockAppWindow = () => { const mock = { send: vi.fn(), @@ -87,22 +105,29 @@ const createMockAppWindow = () => { return mock as unknown as AppWindow; }; -const createMockTelemetry = () => { - const mock = { - track: vi.fn(), - }; - return mock as unknown as ITelemetry; -}; +const createMockTelemetry = (): ITelemetry => ({ + track: vi.fn(), + hasConsent: true, + flush: vi.fn(), + registerHandlers: vi.fn(), + queueSentryEvent: vi.fn(), + popSentryEvent: vi.fn(), + hasPendingSentryEvents: vi.fn().mockReturnValue(false), + clearSentryQueue: vi.fn(), +}); describe('InstallationManager', () => { let manager: InstallationManager; let mockAppWindow: ReturnType; let validationUpdates: InstallValidation[]; - beforeEach(() => { + beforeEach(async () => { vi.clearAllMocks(); validationUpdates = []; + // Reset fs mocks with default behaviors - only the ones we need + vi.mocked(fsPromises.access).mockResolvedValue(undefined); + mockAppWindow = createMockAppWindow(); manager = new InstallationManager(mockAppWindow, createMockTelemetry()); @@ -111,27 +136,32 @@ describe('InstallationManager', () => { path: 'valid/base', }); + // Initialize ComfySettings before creating ComfyInstallation + await ComfySettings.load('valid/base'); + // Capture validation updates vi.spyOn(mockAppWindow, 'send').mockImplementation((channel: string, data: unknown) => { if (channel === IPC_CHANNELS.VALIDATION_UPDATE) { validationUpdates.push({ ...(data as InstallValidation) }); } }); + + // Wait for any pending promises + await Promise.resolve(); }); describe('ensureInstalled', () => { - it('returns existing valid installation', async () => { - const installation = new ComfyInstallation( - 'installed', - 'valid/base', - createMockTelemetry(), - new ComfySettings('valid/base') - ); - vi.spyOn(ComfyInstallation, 'fromConfig').mockResolvedValue(installation); + beforeEach(() => { + // eslint-disable-next-line @typescript-eslint/require-await + vi.spyOn(ComfyInstallation, 'fromConfig').mockImplementation(async () => { + return new ComfyInstallation('installed', 'valid/base', createMockTelemetry()); + }); + }); + it('returns existing valid installation', async () => { const result = await manager.ensureInstalled(); - expect(result).toBe(installation); + expect(result).toBeDefined(); expect(result.hasIssues).toBe(false); expect(result.isValid).toBe(true); expect(mockAppWindow.loadPage).not.toHaveBeenCalledWith('maintenance'); @@ -172,14 +202,6 @@ describe('InstallationManager', () => { ])('$scenario', async ({ mockSetup, expectedErrors }) => { const cleanup = mockSetup?.() as (() => void) | undefined; - const installation = new ComfyInstallation( - 'installed', - 'valid/base', - createMockTelemetry(), - new ComfySettings('valid/base') - ); - vi.spyOn(ComfyInstallation, 'fromConfig').mockResolvedValue(installation); - vi.spyOn( manager as unknown as { resolveIssues: (installation: ComfyInstallation) => Promise }, 'resolveIssues' diff --git a/tests/unit/main-process/comfyDesktopApp.test.ts b/tests/unit/main-process/comfyDesktopApp.test.ts index 56f58fcd..ee0f175d 100644 --- a/tests/unit/main-process/comfyDesktopApp.test.ts +++ b/tests/unit/main-process/comfyDesktopApp.test.ts @@ -13,6 +13,20 @@ import { Terminal } from '@/shell/terminal'; import { findAvailablePort, getModelsDirectory } from '@/utils'; // Mock dependencies +vi.mock('@/config/comfySettings', () => { + const mockSettings = { + get: vi.fn().mockReturnValue(true), + set: vi.fn(), + saveSettings: vi.fn(), + }; + return { + ComfySettings: { + load: vi.fn().mockResolvedValue(mockSettings), + }, + useComfySettings: vi.fn().mockReturnValue(mockSettings), + }; +}); + vi.mock('electron', () => ({ app: { on: vi.fn(), diff --git a/tests/unit/services/telemetry.test.ts b/tests/unit/services/telemetry.test.ts index c62b7382..beb7bd15 100644 --- a/tests/unit/services/telemetry.test.ts +++ b/tests/unit/services/telemetry.test.ts @@ -3,7 +3,7 @@ import fs from 'node:fs'; import path from 'node:path'; import { beforeEach, describe, expect, it, vi } from 'vitest'; -import type { ComfySettings } from '@/config/comfySettings'; +import { ComfySettings, useComfySettings } from '@/config/comfySettings'; import { IPC_CHANNELS } from '@/constants'; import type { AppWindow } from '@/main-process/appWindow'; import { MixpanelTelemetry, promptMetricsConsent } from '@/services/telemetry'; @@ -13,6 +13,7 @@ vi.mock('electron', () => ({ app: { getPath: vi.fn().mockReturnValue('/mock/user/data'), isPackaged: true, + getVersion: vi.fn().mockReturnValue('1.0.0'), }, ipcMain: { on: vi.fn(), @@ -22,7 +23,38 @@ vi.mock('electron', () => ({ }, })); -vi.mock('fs'); +vi.mock('node:fs', () => { + const mockFs = { + existsSync: vi.fn(), + readFileSync: vi.fn(), + writeFileSync: vi.fn(), + }; + return { + default: mockFs, + ...mockFs, + }; +}); + +vi.mock('node:fs/promises', () => ({ + access: vi.fn().mockResolvedValue(undefined), + readFile: vi.fn().mockResolvedValue('{"Comfy-Desktop.SendStatistics": true}'), + writeFile: vi.fn().mockResolvedValue(undefined), +})); + +vi.mock('@/config/comfySettings', () => { + const mockSettings = { + get: vi.fn(), + set: vi.fn(), + saveSettings: vi.fn(), + }; + return { + ComfySettings: { + load: vi.fn().mockResolvedValue(mockSettings), + }, + useComfySettings: vi.fn().mockReturnValue(mockSettings), + }; +}); + vi.mock('mixpanel', () => ({ default: { init: vi.fn(), @@ -33,6 +65,13 @@ vi.mock('mixpanel', () => ({ }, })); +vi.mock('@/store/desktopConfig', () => ({ + useDesktopConfig: vi.fn().mockReturnValue({ + get: vi.fn().mockReturnValue('/mock/path'), + set: vi.fn(), + }), +})); + describe('MixpanelTelemetry', () => { let telemetry: MixpanelTelemetry; const mockInitializedMixpanelClient = { @@ -46,8 +85,10 @@ describe('MixpanelTelemetry', () => { init: vi.fn().mockReturnValue(mockInitializedMixpanelClient), }; - beforeEach(() => { + beforeEach(async () => { vi.clearAllMocks(); + // Initialize settings before each test + await ComfySettings.load('/mock/path'); }); describe('distinct ID management', () => { @@ -197,16 +238,16 @@ describe('MixpanelTelemetry', () => { describe('promptMetricsConsent', () => { let store: Pick; let appWindow: Pick; - let comfySettings: Pick; const versionBeforeUpdate = '0.4.1'; const versionAfterUpdate = '1.0.1'; - beforeEach(() => { + beforeEach(async () => { vi.clearAllMocks(); store = { get: vi.fn(), set: vi.fn() }; appWindow = { loadPage: vi.fn() }; - comfySettings = { get: vi.fn(), set: vi.fn(), saveSettings: vi.fn() }; + // Initialize settings before each test + await ComfySettings.load('/mock/path'); }); const runTest = async ({ @@ -223,20 +264,30 @@ describe('promptMetricsConsent', () => { promptUser?: boolean; }) => { vi.mocked(store.get).mockReturnValue(storeValue); - vi.mocked(comfySettings.get).mockReturnValue(settingsValue); - - if (promptUser) { - vi.mocked(ipcMain.handleOnce).mockImplementationOnce((channel, handler) => { - if (channel === IPC_CHANNELS.SET_METRICS_CONSENT) { - handler(null!, mockConsent); + vi.mocked(useComfySettings().get).mockReturnValue(settingsValue); + + if (mockConsent !== undefined) { + vi.mocked(ipcMain.handleOnce).mockImplementation( + (channel: string, handler: (event: IpcMainEvent, ...args: any[]) => any) => { + if (channel === IPC_CHANNELS.SET_METRICS_CONSENT) { + // Call the handler with the mock consent value and return its result + return handler(null!, mockConsent); + } } - }); + ); } - // @ts-expect-error - store is a mock and doesn't implement all of DesktopConfig - const result = await promptMetricsConsent(store, appWindow, comfySettings); + // @ts-expect-error - appWindow is a mock and doesn't implement all of AppWindow + const result = await promptMetricsConsent(store as DesktopConfig, appWindow); + expect(result).toBe(expectedResult); + if (promptUser) { + expect(store.set).toHaveBeenCalled(); + expect(appWindow.loadPage).toHaveBeenCalledWith('metrics-consent'); + expect(ipcMain.handleOnce).toHaveBeenCalledWith(IPC_CHANNELS.SET_METRICS_CONSENT, expect.any(Function)); + } + if (promptUser) ipcMain.removeHandler(IPC_CHANNELS.SET_METRICS_CONSENT); }; @@ -281,8 +332,6 @@ describe('promptMetricsConsent', () => { expectedResult: false, }); expect(store.set).not.toHaveBeenCalled(); - expect(appWindow.loadPage).not.toHaveBeenCalled(); - expect(ipcMain.handleOnce).not.toHaveBeenCalled(); }); it('should return false if consent is out-of-date and metrics are disabled', async () => {