Refactor determine resource paths.

src/config/extra_model_config.ts

@@ -74,7 +74,7 @@ const configTemplates: Record<string, ModelPaths> = {
   },
 };
 
-export async function getModelConfigPath(): Promise<string> {
+export function getModelConfigPath(): string {
   return path.join(app.getPath('userData'), EXTRA_MODEL_CONFIG_PATH);
 }
 

src/install/resourcePaths.ts

@@ -0,0 +1,24 @@
+import { app } from 'electron';
+import { getModelConfigPath, readBasePathFromConfig } from '../config/extra_model_config';
+import path from 'path';
+
+export async function getBasePath(): Promise<string | null> {
+  const modelConfigPath = getModelConfigPath();
+  return readBasePathFromConfig(modelConfigPath);
+}
+
+export async function getPythonInstallPath(): Promise<string | null> {
+  if (!app.isPackaged) {
+    return path.join(app.getAppPath(), 'assets');
+  }
+
+  return getBasePath();
+}
+
+export async function getAppResourcesPath(): Promise<string> {
+  if (!app.isPackaged) {
+    return path.join(app.getAppPath(), 'assets');
+  }
+
+  return process.resourcesPath;
+}

src/main.ts

@@ -19,6 +19,7 @@ import dotenv from 'dotenv';
 import { buildMenu } from './menu/menu';
 import { ComfyConfigManager } from './config/comfyConfigManager';
 import { AppWindow } from './main-process/appWindow';
+import { getAppResourcesPath, getBasePath, getPythonInstallPath } from './install/resourcePaths';
 
 dotenv.config();
 
@@ -162,11 +163,11 @@ if (!gotTheLock) {
       ipcMain.on(IPC_CHANNELS.OPEN_LOGS_PATH, () => {
         shell.openPath(app.getPath('logs'));
       });
-      ipcMain.handle(IPC_CHANNELS.GET_BASE_PATH, () => {
-        return basePath;
+      ipcMain.handle(IPC_CHANNELS.GET_BASE_PATH, async () => {
+        return await getBasePath();
       });
       ipcMain.handle(IPC_CHANNELS.GET_MODEL_CONFIG_PATH, () => {
-        return modelConfigPath;
+        return getModelConfigPath();
       });
       ipcMain.on(IPC_CHANNELS.OPEN_PATH, (event, folderPath: string) => {
         log.info(`Opening path: ${folderPath}`);
@@ -179,7 +180,8 @@ if (!gotTheLock) {
         return app.isPackaged;
       });
       await handleFirstTimeSetup();
-      const { appResourcesPath, pythonInstallPath, modelConfigPath, basePath } = await determineResourcesPaths();
+      const basePath = await getBasePath();
+      const pythonInstallPath = await getPythonInstallPath();
       if (!basePath || !pythonInstallPath) {
         log.error('ERROR: Base path not found!');
         sendProgressUpdate(ProgressStatus.ERROR_INSTALL_PATH);
@@ -197,10 +199,12 @@ if (!gotTheLock) {
 
       if (!useExternalServer) {
         sendProgressUpdate(ProgressStatus.PYTHON_SETUP);
+        const appResourcesPath = await getAppResourcesPath();
         const pythonEnvironment = new PythonEnvironment(pythonInstallPath, appResourcesPath, spawnPythonAsync);
         await pythonEnvironment.setup();
 
         // TODO: Make tray setup more flexible here as not all actions depend on the python environment.
+        const modelConfigPath = getModelConfigPath();
         SetupTray(
           appWindow,
           () => {
@@ -572,8 +576,7 @@ function findAvailablePort(startPort: number, endPort: number): Promise<number>
  * This means the extra_models_config.yaml file exists in the user's data directory.
  */
 function isFirstTimeSetup(): boolean {
-  const userDataPath = app.getPath('userData');
-  const extraModelsConfigPath = path.join(userDataPath, 'extra_models_config.yaml');
+  const extraModelsConfigPath = getModelConfigPath();
   return !fs.existsSync(extraModelsConfigPath);
 }
 
@@ -594,42 +597,13 @@ async function handleFirstTimeSetup() {
     const selectedDirectory = await selectedInstallDirectory();
     const actualComfyDirectory = ComfyConfigManager.setUpComfyUI(selectedDirectory);
 
-    const modelConfigPath = await getModelConfigPath();
+    const modelConfigPath = getModelConfigPath();
     await createModelConfigFiles(modelConfigPath, actualComfyDirectory);
   } else {
     appWindow.send(IPC_CHANNELS.FIRST_TIME_SETUP_COMPLETE, null);
   }
 }
 
-export async function determineResourcesPaths(): Promise<{
-  pythonInstallPath: string | null;
-  appResourcesPath: string;
-  modelConfigPath: string;
-  basePath: string | null;
-}> {
-  const modelConfigPath = await getModelConfigPath();
-  const basePath = await readBasePathFromConfig(modelConfigPath);
-  const appResourcePath = process.resourcesPath;
-
-  if (!app.isPackaged) {
-    return {
-      // development: install python to in-tree assets dir
-      pythonInstallPath: path.join(app.getAppPath(), 'assets'),
-      appResourcesPath: path.join(app.getAppPath(), 'assets'),
-      modelConfigPath,
-      basePath,
-    };
-  }
-
-  // TODO(robinhuang): Look for extra models yaml file and use that as the userResourcesPath if it exists.
-  return {
-    pythonInstallPath: basePath, // Provide fallback
-    appResourcesPath: appResourcePath,
-    modelConfigPath,
-    basePath,
-  };
-}
-
 /**
  * Rotate old log files by adding a timestamp to the end of the file.
  * @param logDir The directory to rotate the logs in.

tests/unit/extra_model_config.test.ts

@@ -1,7 +1,7 @@
 import { app } from 'electron';
 import path from 'path';
-import { getModelConfigPath } from '../../src/config/extra_model_config';
-
+import { getModelConfigPath, readBasePathFromConfig } from '../../src/config/extra_model_config';
+import * as fsPromises from 'node:fs/promises';
 // Mock electron
 jest.mock('electron', () => ({
   app: {
@@ -20,7 +20,7 @@ describe('getModelConfigPath', () => {
       throw new Error(`Unexpected getPath key: ${key}`);
     });
 
-    const result = await getModelConfigPath();
+    const result = getModelConfigPath();
 
     // Verify the path is correctly joined
     expect(result).toBe(path.join(mockUserDataPath, 'extra_models_config.yaml'));
@@ -29,3 +29,43 @@ describe('getModelConfigPath', () => {
     expect(app.getPath).toHaveBeenCalledWith('userData');
   });
 });
+
+describe('readBasePathFromConfig', () => {
+  const testConfigPath = path.join(__dirname, 'test_config.yaml');
+
+  beforeAll(async () => {
+    // Create a test YAML file
+    const testConfig = `# Test ComfyUI config
+comfyui:
+  base_path: ~/test/comfyui
+  is_default: true
+  checkpoints: models/checkpoints/
+  loras: models/loras/`;
+
+    await fsPromises.writeFile(testConfigPath, testConfig, 'utf8');
+  });
+
+  afterAll(async () => {
+    await fsPromises.rm(testConfigPath);
+  });
+
+  it('should read base_path from valid config file', async () => {
+    const result = await readBasePathFromConfig(testConfigPath);
+    expect(result).toBe('~/test/comfyui');
+  });
+
+  it('should return null for non-existent file', async () => {
+    const result = await readBasePathFromConfig('non_existent_file.yaml');
+    expect(result).toBeNull();
+  });
+
+  it('should return null for invalid config file', async () => {
+    const invalidConfigPath = path.join(__dirname, 'invalid_config.yaml');
+    await fsPromises.writeFile(invalidConfigPath, 'invalid: yaml: content:', 'utf8');
+
+    const result = await readBasePathFromConfig(invalidConfigPath);
+    expect(result).toBeNull();
+
+    await fsPromises.rm(invalidConfigPath);
+  });
+});

tests/unit/install/resourcePath.test.ts

@@ -0,0 +1,86 @@
+import { app } from 'electron';
+import path from 'path';
+import { getBasePath, getPythonInstallPath, getAppResourcesPath } from '../../../src/install/resourcePaths';
+import { getModelConfigPath, readBasePathFromConfig } from '../../../src/config/extra_model_config';
+
+// Mock the external modules
+jest.mock('electron', () => ({
+  app: {
+    isPackaged: false,
+    getAppPath: jest.fn(),
+  },
+}));
+
+jest.mock('../../../src/config/extra_model_config', () => ({
+  getModelConfigPath: jest.fn(),
+  readBasePathFromConfig: jest.fn(),
+}));
+
+describe('resourcePaths', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  describe('getBasePath', () => {
+    it('should return the base path from config', async () => {
+      const mockConfigPath = '/mock/config/path';
+      const mockBasePath = '/mock/base/path';
+
+      (getModelConfigPath as jest.Mock).mockReturnValue(mockConfigPath);
+      (readBasePathFromConfig as jest.Mock).mockResolvedValue(mockBasePath);
+
+      const result = await getBasePath();
+
+      expect(getModelConfigPath).toHaveBeenCalled();
+      expect(readBasePathFromConfig).toHaveBeenCalledWith(mockConfigPath);
+      expect(result).toBe(mockBasePath);
+    });
+  });
+
+  describe('getPythonInstallPath', () => {
+    it('should return assets path when app is not packaged', async () => {
+      const mockAppPath = '/mock/app/path';
+      (app.getAppPath as jest.Mock).mockReturnValue(mockAppPath);
+      (app.isPackaged as boolean) = false;
+
+      const result = await getPythonInstallPath();
+
+      expect(result).toBe(path.join(mockAppPath, 'assets'));
+      expect(app.getAppPath).toHaveBeenCalled();
+    });
+
+    it('should return base path when app is packaged', async () => {
+      const mockBasePath = '/mock/base/path';
+      (app.isPackaged as boolean) = true;
+      (getModelConfigPath as jest.Mock).mockReturnValue('/mock/config');
+      (readBasePathFromConfig as jest.Mock).mockResolvedValue(mockBasePath);
+
+      const result = await getPythonInstallPath();
+
+      expect(result).toBe(mockBasePath);
+    });
+  });
+
+  describe('getAppResourcesPath', () => {
+    it('should return assets path when app is not packaged', async () => {
+      const mockAppPath = '/mock/app/path';
+      (app.getAppPath as jest.Mock).mockReturnValue(mockAppPath);
+      (app.isPackaged as boolean) = false;
+
+      const result = await getAppResourcesPath();
+
+      expect(result).toBe(path.join(mockAppPath, 'assets'));
+      expect(app.getAppPath).toHaveBeenCalled();
+    });
+
+    it('should return resources path when app is packaged', async () => {
+      (app.isPackaged as boolean) = true;
+      const mockResourcesPath = '/mock/resources/path';
+      (process as any).resourcesPath = mockResourcesPath;
+
+      const result = await getAppResourcesPath();
+
+      expect(result).toBe(mockResourcesPath);
+    });
+  });
+});