diff --git a/src/main-process/comfyInstallation.ts b/src/main-process/comfyInstallation.ts index 3859f59aa..3a14d29ca 100644 --- a/src/main-process/comfyInstallation.ts +++ b/src/main-process/comfyInstallation.ts @@ -1,6 +1,7 @@ import log from 'electron-log/main'; import { rm } from 'node:fs/promises'; +import { ComfyConfigManager } from '../config/comfyConfigManager'; import { ComfyServerConfig } from '../config/comfyServerConfig'; import { ComfySettings, useComfySettings } from '../config/comfySettings'; import { evaluatePathRestrictions } from '../handlers/pathHandlers'; @@ -135,7 +136,8 @@ export class ComfyInstallation { } else { await this.updateBasePathAndVenv(basePath); - validation.basePath = 'OK'; + const hasDirectories = this.ensureBaseDirectoryStructure(basePath); + validation.basePath = hasDirectories ? 'OK' : 'error'; } this.onUpdate?.(validation); @@ -260,4 +262,23 @@ export class ComfyInstallation { } await useDesktopConfig().permanentlyDeleteConfigFile(); } + + /** + * Ensure the base path contains the expected ComfyUI directory structure. + * Creates missing directories; returns false if the structure is still invalid. + */ + private ensureBaseDirectoryStructure(basePath: string): boolean { + try { + ComfyConfigManager.createComfyDirectories(basePath); + } catch (error) { + log.error('Failed to create ComfyUI directories at base path.', error); + return false; + } + + const isValid = ComfyConfigManager.isComfyUIDirectory(basePath); + if (!isValid) { + log.error(`Base path is missing required ComfyUI directories after creation attempt. [${basePath}]`); + } + return isValid; + } } diff --git a/tests/unit/main-process/comfyInstallation.test.ts b/tests/unit/main-process/comfyInstallation.test.ts new file mode 100644 index 000000000..93a54b77f --- /dev/null +++ b/tests/unit/main-process/comfyInstallation.test.ts @@ -0,0 +1,97 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { ComfyInstallation } from '@/main-process/comfyInstallation'; +import type { ITelemetry } from '@/services/telemetry'; + +vi.mock('electron-log/main', () => ({ + default: { + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + verbose: vi.fn(), + }, +})); + +vi.mock('@/virtualEnvironment', () => ({ + VirtualEnvironment: vi.fn(() => ({})), +})); + +const createDirectories = vi.fn(); +const isComfyDirectory = vi.fn(); + +vi.mock('@/config/comfyConfigManager', () => ({ + ComfyConfigManager: { + createComfyDirectories: createDirectories, + isComfyUIDirectory: isComfyDirectory, + }, +})); + +vi.mock('@/store/desktopConfig', () => ({ + useDesktopConfig: vi.fn(() => ({ + get: vi.fn((key: string) => { + if (key === 'selectedDevice') return 'cpu'; + if (key === 'basePath') return '/base/path'; + return undefined; + }), + set: vi.fn(), + })), +})); + +vi.mock('@/config/comfySettings', () => ({ + useComfySettings: vi.fn(() => ({ + get: vi.fn(), + set: vi.fn(), + })), + ComfySettings: { + load: vi.fn(), + }, +})); + +const telemetry: ITelemetry = { + track: vi.fn(), + hasConsent: true, + flush: vi.fn(), + registerHandlers: vi.fn(), + loadGenerationCount: vi.fn(), +}; + +describe('ComfyInstallation.ensureBaseDirectoryStructure', () => { + const basePath = '/base/path'; + let installation: ComfyInstallation; + + beforeEach(() => { + vi.clearAllMocks(); + installation = new ComfyInstallation('installed', basePath, telemetry); + }); + + it('creates and validates directory structure', () => { + isComfyDirectory.mockReturnValue(true); + + const result = (installation as any).ensureBaseDirectoryStructure(basePath); + + expect(result).toBe(true); + expect(createDirectories).toHaveBeenCalledWith(basePath); + expect(isComfyDirectory).toHaveBeenCalledWith(basePath); + }); + + it('returns false when directory creation throws', () => { + createDirectories.mockImplementation(() => { + throw new Error('fail'); + }); + + const result = (installation as any).ensureBaseDirectoryStructure(basePath); + + expect(result).toBe(false); + expect(isComfyDirectory).not.toHaveBeenCalled(); + }); + + it('returns false when directory validation fails', () => { + isComfyDirectory.mockReturnValue(false); + + const result = (installation as any).ensureBaseDirectoryStructure(basePath); + + expect(result).toBe(false); + expect(createDirectories).toHaveBeenCalledWith(basePath); + expect(isComfyDirectory).toHaveBeenCalledWith(basePath); + }); +});