Add tests

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<void>
+  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<string, MockFunction>()
+
+      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<void>((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<string, any>
+          }
+
+          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 }
+  ]
+})

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()
+  })
+})

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'

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')
+  }
+}