Refactor creating ComfyUI directories + unit tests. (#185)

src/__tests__/unit/comfyConfigManager.test.ts

@@ -0,0 +1,165 @@
+import fs from 'fs';
+import { ComfyConfigManager, DirectoryStructure } from '../../config/comfyConfigManager';
+
+// Mock the fs module
+jest.mock('fs');
+jest.mock('electron-log/main', () => ({
+  info: jest.fn(),
+  error: jest.fn(),
+  warn: jest.fn(),
+}));
+
+describe('ComfyConfigManager', () => {
+  // Reset all mocks before each test
+  beforeEach(() => {
+    jest.clearAllMocks();
+    (fs.existsSync as jest.Mock).mockReset();
+    (fs.mkdirSync as jest.Mock).mockReset();
+    (fs.writeFileSync as jest.Mock).mockReset();
+    (fs.renameSync as jest.Mock).mockReset();
+  });
+
+  describe('setUpComfyUI', () => {
+    it('should use existing directory when it contains ComfyUI structure', () => {
+      // Mock isComfyUIDirectory to return true for the input path
+      (fs.existsSync as jest.Mock).mockImplementation((path: string) => {
+        const requiredDirs = [
+          '/existing/ComfyUI/models',
+          '/existing/ComfyUI/input',
+          '/existing/ComfyUI/user',
+          '/existing/ComfyUI/output',
+          '/existing/ComfyUI/custom_nodes',
+        ];
+        return requiredDirs.includes(path);
+      });
+
+      const result = ComfyConfigManager.setUpComfyUI('/existing/ComfyUI');
+
+      expect(result).toBe('/existing/ComfyUI');
+    });
+
+    it('should create ComfyUI subdirectory when it is missing', () => {
+      (fs.existsSync as jest.Mock).mockImplementationOnce((path: string) => {
+        if (path === '/some/base/path/ComfyUI') {
+          return false;
+        }
+        return true;
+      });
+
+      const result = ComfyConfigManager.setUpComfyUI('/some/base/path');
+
+      expect(result).toBe('/some/base/path/ComfyUI');
+    });
+  });
+
+  describe('isComfyUIDirectory', () => {
+    it('should return true when all required directories exist', () => {
+      (fs.existsSync as jest.Mock).mockImplementation((path: string) => {
+        const requiredDirs = [
+          '/fake/path/models',
+          '/fake/path/input',
+          '/fake/path/user',
+          '/fake/path/output',
+          '/fake/path/custom_nodes',
+        ];
+        return requiredDirs.includes(path);
+      });
+
+      const result = ComfyConfigManager.isComfyUIDirectory('/fake/path');
+
+      expect(result).toBe(true);
+      expect(fs.existsSync).toHaveBeenCalledTimes(5);
+    });
+
+    it('should return false when some required directories are missing', () => {
+      (fs.existsSync as jest.Mock)
+        .mockReturnValueOnce(true) // models exists
+        .mockReturnValueOnce(true) // input exists
+        .mockReturnValueOnce(false) // user missing
+        .mockReturnValueOnce(true) // output exists
+        .mockReturnValueOnce(true); // custom_nodes exists
+
+      const result = ComfyConfigManager.isComfyUIDirectory('/fake/path');
+
+      expect(result).toBe(false);
+    });
+  });
+
+  describe('createComfyDirectories', () => {
+    it('should create all necessary directories when none exist', () => {
+      (fs.existsSync as jest.Mock).mockReturnValue(false);
+
+      ComfyConfigManager.createComfyDirectories('/fake/path/ComfyUI');
+
+      // Verify each required directory was created
+      expect(fs.mkdirSync).toHaveBeenCalledWith('/fake/path/ComfyUI/models', { recursive: true });
+      expect(fs.mkdirSync).toHaveBeenCalledWith('/fake/path/ComfyUI/input', { recursive: true });
+      expect(fs.mkdirSync).toHaveBeenCalledWith('/fake/path/ComfyUI/user', { recursive: true });
+      expect(fs.mkdirSync).toHaveBeenCalledWith('/fake/path/ComfyUI/output', { recursive: true });
+      expect(fs.mkdirSync).toHaveBeenCalledWith('/fake/path/ComfyUI/custom_nodes', { recursive: true });
+    });
+  });
+
+  describe('createComfyConfigFile', () => {
+    it('should create new config file when none exists', () => {
+      (fs.existsSync as jest.Mock).mockReturnValue(false);
+
+      ComfyConfigManager.createComfyConfigFile('/fake/path', false);
+
+      expect(fs.writeFileSync).toHaveBeenCalledTimes(1);
+      expect(fs.renameSync).not.toHaveBeenCalled();
+    });
+
+    it('should backup existing config file when overwrite is true', () => {
+      (fs.existsSync as jest.Mock).mockImplementation((path: string) => {
+        return path === '/user/default/comfy.settings.json';
+      });
+
+      ComfyConfigManager.createComfyConfigFile('/user/default', true);
+
+      expect(fs.renameSync).toHaveBeenCalledTimes(1);
+      expect(fs.writeFileSync).toHaveBeenCalledTimes(1);
+    });
+
+    it('should handle backup failure gracefully', () => {
+      (fs.existsSync as jest.Mock).mockReturnValue(true);
+      (fs.renameSync as jest.Mock).mockImplementation(() => {
+        throw new Error('Backup failed');
+      });
+
+      ComfyConfigManager.createComfyConfigFile('/fake/path', true);
+
+      expect(fs.writeFileSync).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('createNestedDirectories', () => {
+    it('should create nested directory structure correctly', () => {
+      (fs.existsSync as jest.Mock).mockReturnValue(false);
+
+      const structure = ['dir1', ['dir2', ['subdir1', 'subdir2']], ['dir3', [['subdir3', ['subsubdir1']]]]];
+
+      ComfyConfigManager['createNestedDirectories']('/fake/path', structure);
+
+      // Verify the correct paths were created
+      expect(fs.mkdirSync).toHaveBeenCalledWith(expect.stringContaining('dir1'), expect.any(Object));
+      expect(fs.mkdirSync).toHaveBeenCalledWith(expect.stringContaining('dir2'), expect.any(Object));
+      expect(fs.mkdirSync).toHaveBeenCalledWith(expect.stringContaining('subdir1'), expect.any(Object));
+      expect(fs.mkdirSync).toHaveBeenCalledWith(expect.stringContaining('subsubdir1'), expect.any(Object));
+    });
+
+    it('should handle invalid directory structure items', () => {
+      const invalidStructure = [
+        'dir1',
+        ['dir2'], // Invalid: array with only one item
+        [123, ['subdir1']], // Invalid: non-string directory name
+      ];
+
+      ComfyConfigManager['createNestedDirectories']('/fake/path', invalidStructure as DirectoryStructure);
+
+      // Verify only valid directories were created
+      expect(fs.mkdirSync).toHaveBeenCalledWith(expect.stringContaining('dir1'), expect.any(Object));
+      expect(fs.mkdirSync).not.toHaveBeenCalledWith(expect.stringContaining('subdir1'), expect.any(Object));
+    });
+  });
+});

src/config/comfyConfigManager.ts

@@ -0,0 +1,142 @@
+import fs from 'fs';
+import path from 'path';
+import log from 'electron-log/main';
+
+export type DirectoryStructure = (string | DirectoryStructure)[];
+
+export class ComfyConfigManager {
+  private static readonly DEFAULT_DIRECTORIES: DirectoryStructure = [
+    'custom_nodes',
+    'input',
+    'output',
+    ['user', ['default']],
+    [
+      'models',
+      [
+        'checkpoints',
+        'clip',
+        'clip_vision',
+        'configs',
+        'controlnet',
+        'diffusers',
+        'diffusion_models',
+        'embeddings',
+        'gligen',
+        'hypernetworks',
+        'loras',
+        'photomaker',
+        'style_models',
+        'unet',
+        'upscale_models',
+        'vae',
+        'vae_approx',
+
+        // TODO(robinhuang): Remove when we have a better way to specify base model paths.
+        'animatediff_models',
+        'animatediff_motion_lora',
+        'animatediff_video_formats',
+        'liveportrait',
+        ['insightface', ['buffalo_1']],
+        ['blip', ['checkpoints']],
+        'CogVideo',
+        ['xlabs', ['loras', 'controlnets']],
+        'layerstyle',
+        'LLM',
+        'Joy_caption',
+      ],
+    ],
+  ];
+
+  private static readonly DEFAULT_CONFIG = {
+    'Comfy.ColorPalette': 'dark',
+    'Comfy.UseNewMenu': 'Top',
+    'Comfy.Workflow.WorkflowTabsPosition': 'Topbar',
+    'Comfy.Workflow.ShowMissingModelsWarning': true,
+  };
+
+  public static setUpComfyUI(localComfyDirectory: string): string {
+    if (!this.isComfyUIDirectory(localComfyDirectory)) {
+      log.info(
+        `Selected directory ${localComfyDirectory} is not a ComfyUI directory. Appending ComfyUI to install path.`
+      );
+      localComfyDirectory = path.join(localComfyDirectory, 'ComfyUI');
+    }
+
+    this.createComfyDirectories(localComfyDirectory);
+    const userSettingsPath = path.join(localComfyDirectory, 'user', 'default');
+    this.createComfyConfigFile(userSettingsPath, true);
+    return localComfyDirectory;
+  }
+
+  public static createComfyConfigFile(userSettingsPath: string, overwrite: boolean = false): void {
+    const configFilePath = path.join(userSettingsPath, 'comfy.settings.json');
+
+    if (fs.existsSync(configFilePath) && overwrite) {
+      const backupFilePath = path.join(userSettingsPath, 'old_comfy.settings.json');
+      try {
+        fs.renameSync(configFilePath, backupFilePath);
+        log.info(`Renaming existing user settings file to: ${backupFilePath}`);
+      } catch (error) {
+        log.error(`Failed to backup existing user settings file: ${error}`);
+        return;
+      }
+    }
+
+    try {
+      fs.writeFileSync(configFilePath, JSON.stringify(this.DEFAULT_CONFIG, null, 2));
+      log.info(`Created new ComfyUI config file at: ${configFilePath}`);
+    } catch (error) {
+      log.error(`Failed to create new ComfyUI config file: ${error}`);
+    }
+  }
+
+  public static isComfyUIDirectory(directory: string): boolean {
+    const requiredSubdirs = ['models', 'input', 'user', 'output', 'custom_nodes'];
+    return requiredSubdirs.every((subdir) => fs.existsSync(path.join(directory, subdir)));
+  }
+
+  static createComfyDirectories(localComfyDirectory: string): void {
+    log.info(`Creating ComfyUI directories in ${localComfyDirectory}`);
+
+    try {
+      this.createNestedDirectories(localComfyDirectory, this.DEFAULT_DIRECTORIES);
+    } catch (error) {
+      log.error(`Failed to create ComfyUI directories: ${error}`);
+    }
+  }
+
+  static createNestedDirectories(basePath: string, structure: DirectoryStructure): void {
+    structure.forEach((item) => {
+      if (typeof item === 'string') {
+        const dirPath = path.join(basePath, item);
+        this.createDirIfNotExists(dirPath);
+      } else if (Array.isArray(item) && item.length === 2) {
+        const [dirName, subDirs] = item;
+        if (typeof dirName === 'string') {
+          const newBasePath = path.join(basePath, dirName);
+          this.createDirIfNotExists(newBasePath);
+          if (Array.isArray(subDirs)) {
+            this.createNestedDirectories(newBasePath, subDirs);
+          }
+        } else {
+          log.warn(`Invalid directory structure item: ${JSON.stringify(item)}`);
+        }
+      } else {
+        log.warn(`Invalid directory structure item: ${JSON.stringify(item)}`);
+      }
+    });
+  }
+
+  /**
+   * Create a directory if not exists
+   * @param dirPath
+   */
+  static createDirIfNotExists(dirPath: string): void {
+    if (!fs.existsSync(dirPath)) {
+      fs.mkdirSync(dirPath, { recursive: true });
+      log.info(`Created directory: ${dirPath}`);
+    } else {
+      log.info(`Directory already exists: ${dirPath}`);
+    }
+  }
+}