From 677d804bcd1597326bf05d72c8592f3df1605ee3 Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Wed, 6 Nov 2024 17:16:30 -0800 Subject: [PATCH 1/2] Refactor determine resource paths. --- src/config/extra_model_config.ts | 2 +- src/install/resourcePaths.ts | 24 ++++++++++++++ src/main.ts | 46 ++++++--------------------- tests/unit/extra_model_config.test.ts | 2 +- 4 files changed, 36 insertions(+), 38 deletions(-) create mode 100644 src/install/resourcePaths.ts diff --git a/src/config/extra_model_config.ts b/src/config/extra_model_config.ts index dd3ca32d..c88b5616 100644 --- a/src/config/extra_model_config.ts +++ b/src/config/extra_model_config.ts @@ -74,7 +74,7 @@ const configTemplates: Record = { }, }; -export async function getModelConfigPath(): Promise { +export function getModelConfigPath(): string { return path.join(app.getPath('userData'), EXTRA_MODEL_CONFIG_PATH); } diff --git a/src/install/resourcePaths.ts b/src/install/resourcePaths.ts new file mode 100644 index 00000000..ea9ef07a --- /dev/null +++ b/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 { + const modelConfigPath = getModelConfigPath(); + return readBasePathFromConfig(modelConfigPath); +} + +export async function getPythonInstallPath(): Promise { + if (!app.isPackaged) { + return path.join(app.getAppPath(), 'assets'); + } + + return getBasePath(); +} + +export async function getAppResourcesPath(): Promise { + if (!app.isPackaged) { + return path.join(app.getAppPath(), 'assets'); + } + + return process.resourcesPath; +} diff --git a/src/main.ts b/src/main.ts index 76085982..db1eec57 100644 --- a/src/main.ts +++ b/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 * 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. diff --git a/tests/unit/extra_model_config.test.ts b/tests/unit/extra_model_config.test.ts index b74006ae..ffc428df 100644 --- a/tests/unit/extra_model_config.test.ts +++ b/tests/unit/extra_model_config.test.ts @@ -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')); From 6fc8c248fca71492f6c28e202289aa0a40317323 Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Fri, 8 Nov 2024 14:58:25 -0800 Subject: [PATCH 2/2] Add test. --- tests/unit/extra_model_config.test.ts | 44 ++++++++++++- tests/unit/install/resourcePath.test.ts | 86 +++++++++++++++++++++++++ 2 files changed, 128 insertions(+), 2 deletions(-) create mode 100644 tests/unit/install/resourcePath.test.ts diff --git a/tests/unit/extra_model_config.test.ts b/tests/unit/extra_model_config.test.ts index ffc428df..4e12932f 100644 --- a/tests/unit/extra_model_config.test.ts +++ b/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: { @@ -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); + }); +}); diff --git a/tests/unit/install/resourcePath.test.ts b/tests/unit/install/resourcePath.test.ts new file mode 100644 index 00000000..50be268b --- /dev/null +++ b/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); + }); + }); +});