diff --git a/browser_tests/desktop/fixtures/electron.ts b/browser_tests/desktop/fixtures/electron.ts new file mode 100644 index 000000000..b175182a2 --- /dev/null +++ b/browser_tests/desktop/fixtures/electron.ts @@ -0,0 +1,124 @@ +import { test as base } from '@playwright/test' + +type ElectronFixtureOptions = { + registerDefaults?: { + downloadManager?: boolean + } +} + +type MockFunction = { + calls: unknown[][] + called: () => Promise + handler?: (args: unknown[]) => unknown +} + +export type MockElectronAPI = { + setup: (method: string, handler: (args: unknown[]) => unknown) => MockFunction +} + +export const electronFixture = base.extend<{ + electronAPI: MockElectronAPI + electronOptions: ElectronFixtureOptions +}>({ + electronOptions: [ + { + registerDefaults: { + downloadManager: true + } + }, + { option: true } + ], + + electronAPI: [ + async ({ page, electronOptions }, use) => { + const mocks = new Map() + + await page.exposeFunction( + '__handleMockCall', + async (method: string, args: unknown[]) => { + const mock = mocks.get(method) + + if (electronOptions.registerDefaults?.downloadManager) { + if (method === 'DownloadManager.getAllDownloads') { + return [] + } + } + + if (!mock) return null + mock.calls.push(args) + return mock.handler ? mock.handler(args) : null + } + ) + + const createMockFunction = ( + method: string, + handler: (args: unknown[]) => unknown + ): MockFunction => { + let resolveNextCall: (() => void) | null = null + + const mockFn: MockFunction = { + calls: [], + async called() { + if (this.calls.length > 0) return + + return new Promise((resolve) => { + resolveNextCall = resolve + }) + }, + handler: (args: unknown[]) => { + const result = handler(args) + resolveNextCall?.() + resolveNextCall = null + return result + } + } + mocks.set(method, mockFn) + + // Add the method to the window.electronAPI object + page.evaluate((methodName) => { + const w = window as typeof window & { + electronAPI: Record + } + + w.electronAPI[methodName] = async (...args: unknown[]) => { + return window['__handleMockCall'](methodName, args) + } + }, method) + + return mockFn + } + + const testAPI: MockElectronAPI = { + setup(method, handler) { + console.log('adding handler for', method) + return createMockFunction(method, handler) + } + } + + await page.addInitScript(async () => { + const getProxy = (...path: string[]) => { + return new Proxy(() => {}, { + // Handle the proxy itself being called as a function + apply: async (target, thisArg, argArray) => { + return window['__handleMockCall'](path.join('.'), argArray) + }, + // Handle property access + get: (target, prop: string) => { + return getProxy(...path, prop) + } + }) + } + + const w = window as typeof window & { + electronAPI: any + } + + w.electronAPI = getProxy() + console.log('registered electron api') + }) + + await use(testAPI) + }, + { auto: true } + ] +}) diff --git a/browser_tests/desktop/importModel.spec.ts b/browser_tests/desktop/importModel.spec.ts new file mode 100644 index 000000000..0eb7697ac --- /dev/null +++ b/browser_tests/desktop/importModel.spec.ts @@ -0,0 +1,172 @@ +import { expect, mergeTests } from '@playwright/test' + +import { ComfyPage, comfyPageFixture } from '../fixtures/ComfyPage' +import { MockElectronAPI, electronFixture } from './fixtures/electron' + +const test = mergeTests(comfyPageFixture, electronFixture) + +comfyPageFixture.describe('Import Model (web)', () => { + comfyPageFixture( + 'Import dialog does not show when electron api is not available', + async ({ comfyPage }) => { + await comfyPage.dragAndDropExternalResource({ + fileName: 'test.bin', + buffer: Buffer.from('') + }) + + // Normal unable to find workflow in file error + await expect( + comfyPage.page.locator('.p-toast-message.p-toast-message-warn') + ).toHaveCount(1) + } + ) +}) + +test.describe('Import Model (electron)', () => { + const dropFile = async ( + comfyPage: ComfyPage, + electronAPI: MockElectronAPI, + fileName: string, + metadata: string + ) => { + const getFilePathMock = electronAPI.setup('getFilePath', () => + Promise.resolve('some/file/path/' + fileName) + ) + + let buffer: Buffer | undefined + if (metadata) { + const contentBuffer = Buffer.from(metadata, 'utf-8') + + const headerSizeBuffer = Buffer.alloc(8) + headerSizeBuffer.writeBigUInt64LE(BigInt(contentBuffer.length)) + + buffer = Buffer.concat([headerSizeBuffer, contentBuffer]) + } + + await comfyPage.dragAndDropExternalResource({ + fileName, + buffer + }) + await getFilePathMock.called() + + await expect( + comfyPage.page.locator('.p-toast-message.p-toast-message-warn') + ).toHaveCount(0) + await expect(comfyPage.importModelDialog.rootEl).toBeVisible() + } + + test('Can show import file dialog by dropping file onto the app', async ({ + comfyPage, + electronAPI + }) => { + await dropFile(comfyPage, electronAPI, 'test.bin', '{}') + }) + + test('Can autodetect checkpoint model type from modelspec', async ({ + comfyPage, + electronAPI + }) => { + await dropFile( + comfyPage, + electronAPI, + 'file.safetensors', + JSON.stringify({ + __metadata__: { + 'modelspec.sai_model_spec': 'test', + 'modelspec.architecture': 'stable-diffusion-v1' + } + }) + ) + + await expect(comfyPage.importModelDialog.modelTypeInput).toHaveValue( + 'checkpoints' + ) + }) + + test('Can autodetect lora model type from modelspec', async ({ + comfyPage, + electronAPI + }) => { + await dropFile( + comfyPage, + electronAPI, + 'file.safetensors', + JSON.stringify({ + __metadata__: { + 'modelspec.sai_model_spec': 'test', + 'modelspec.architecture': 'Flux.1-AE/lora' + } + }) + ) + + await expect(comfyPage.importModelDialog.modelTypeInput).toHaveValue( + 'loras' + ) + }) + + test('Can autodetect checkpoint model type from header keys', async ({ + comfyPage, + electronAPI + }) => { + await dropFile( + comfyPage, + electronAPI, + 'file.safetensors', + JSON.stringify({ + 'model.diffusion_model.input_blocks.0.0.bias': {} + }) + ) + + await expect(comfyPage.importModelDialog.modelTypeInput).toHaveValue( + 'checkpoints' + ) + }) + + test('Can autodetect lora model type from header keys', async ({ + comfyPage, + electronAPI + }) => { + await dropFile( + comfyPage, + electronAPI, + 'file.safetensors', + JSON.stringify({ + 'lora_unet_down_blocks_0_attentions_0_proj_in.alpha': {} + }) + ) + + await expect(comfyPage.importModelDialog.modelTypeInput).toHaveValue( + 'loras' + ) + }) + + test('Can import file', async ({ comfyPage, electronAPI }) => { + await dropFile( + comfyPage, + electronAPI, + 'checkpoint_modelspec.safetensors', + '{}' + ) + + const importModelMock = electronAPI.setup( + 'importModel', + () => new Promise((resolve) => setTimeout(resolve, 100)) + ) + + // Model type is required so select one + await expect(comfyPage.importModelDialog.importButton).toBeDisabled() + await comfyPage.importModelDialog.modelTypeInput.fill('checkpoints') + await expect(comfyPage.importModelDialog.importButton).toBeEnabled() + + // Click import, ensure API is called + await comfyPage.importModelDialog.importButton.click() + await importModelMock.called() + + // Toast should be shown and dialog closes + await expect( + comfyPage.page.locator('.p-toast-message.p-toast-message-success') + ).toHaveCount(1) + + await expect(comfyPage.importModelDialog.rootEl).toBeHidden() + }) +}) diff --git a/browser_tests/fixtures/ComfyPage.ts b/browser_tests/fixtures/ComfyPage.ts index 00c9ca098..52132fc6e 100644 --- a/browser_tests/fixtures/ComfyPage.ts +++ b/browser_tests/fixtures/ComfyPage.ts @@ -13,6 +13,7 @@ import { ComfyActionbar } from '../helpers/actionbar' import { ComfyTemplates } from '../helpers/templates' import { ComfyMouse } from './ComfyMouse' import { ComfyNodeSearchBox } from './components/ComfyNodeSearchBox' +import { ImportModelDialog } from './components/ImportModelDialog' import { SettingDialog } from './components/SettingDialog' import { NodeLibrarySidebarTab, @@ -140,6 +141,7 @@ export class ComfyPage { public readonly templates: ComfyTemplates public readonly settingDialog: SettingDialog public readonly confirmDialog: ConfirmDialog + public readonly importModelDialog: ImportModelDialog /** Worker index to test user ID */ public readonly userIds: string[] = [] @@ -165,6 +167,7 @@ export class ComfyPage { this.templates = new ComfyTemplates(page) this.settingDialog = new SettingDialog(page) this.confirmDialog = new ConfirmDialog(page) + this.importModelDialog = new ImportModelDialog(page) } convertLeafToContent(structure: FolderStructure): FolderStructure { @@ -469,6 +472,7 @@ export class ComfyPage { fileName?: string url?: string dropPosition?: Position + buffer?: Buffer } = {} ) { const { dropPosition = { x: 100, y: 100 }, fileName, url } = options @@ -487,7 +491,7 @@ export class ComfyPage { // Dropping a file from the filesystem if (fileName) { const filePath = this.assetPath(fileName) - const buffer = fs.readFileSync(filePath) + const buffer = options.buffer ?? fs.readFileSync(filePath) const getFileType = (fileName: string) => { if (fileName.endsWith('.png')) return 'image/png' diff --git a/browser_tests/fixtures/components/ImportModelDialog.ts b/browser_tests/fixtures/components/ImportModelDialog.ts new file mode 100644 index 000000000..87f2dfc05 --- /dev/null +++ b/browser_tests/fixtures/components/ImportModelDialog.ts @@ -0,0 +1,17 @@ +import { Page } from '@playwright/test' + +export class ImportModelDialog { + constructor(public readonly page: Page) {} + + get rootEl() { + return this.page.locator('div[aria-labelledby="global-import-model"]') + } + + get modelTypeInput() { + return this.rootEl.locator('#model-type') + } + + get importButton() { + return this.rootEl.getByLabel('Import') + } +}