diff --git a/browser_tests/ComfyPage.ts b/browser_tests/ComfyPage.ts index fbf8e0a1ca..89d83a1105 100644 --- a/browser_tests/ComfyPage.ts +++ b/browser_tests/ComfyPage.ts @@ -4,6 +4,8 @@ import dotenv from 'dotenv' dotenv.config() import * as fs from 'fs' import { NodeBadgeMode } from '../src/types/nodeSource' +import { NodeId } from '../src/types/comfyWorkflow' +import { ManageGroupNode } from './helpers/manageGroupNode' interface Position { x: number @@ -707,6 +709,177 @@ export class ComfyPage { await this.page.getByText('Convert to Group Node').click() await this.nextFrame() } + async convertOffsetToCanvas(pos: [number, number]) { + return this.page.evaluate((pos) => { + return window['app'].canvas.ds.convertOffsetToCanvas(pos) + }, pos) + } + async getNodeRefById(id: NodeId) { + return new NodeReference(id, this) + } + async getNodeRefsByType(type: string): Promise { + return ( + await this.page.evaluate((type) => { + return window['app'].graph._nodes + .filter((n) => n.type === type) + .map((n) => n.id) + }, type) + ).map((id: NodeId) => this.getNodeRefById(id)) + } +} +class NodeSlotReference { + constructor( + readonly type: 'input' | 'output', + readonly index: number, + readonly node: NodeReference + ) {} + async getPosition() { + const pos: [number, number] = await this.node.comfyPage.page.evaluate( + ([type, id, index]) => { + const node = window['app'].graph.getNodeById(id) + if (!node) throw new Error(`Node ${id} not found.`) + return window['app'].canvas.ds.convertOffsetToCanvas( + node.getConnectionPos(type === 'input', index) + ) + }, + [this.type, this.node.id, this.index] as const + ) + return { + x: pos[0], + y: pos[1] + } + } + async getLinkCount() { + return await this.node.comfyPage.page.evaluate( + ([type, id, index]) => { + const node = window['app'].graph.getNodeById(id) + if (!node) throw new Error(`Node ${id} not found.`) + if (type === 'input') { + return node.inputs[index].link == null ? 0 : 1 + } + return node.outputs[index].links?.length ?? 0 + }, + [this.type, this.node.id, this.index] as const + ) + } + async removeLinks() { + await this.node.comfyPage.page.evaluate( + ([type, id, index]) => { + const node = window['app'].graph.getNodeById(id) + if (!node) throw new Error(`Node ${id} not found.`) + if (type === 'input') { + node.disconnectInput(index) + } else { + node.disconnectOutput(index) + } + }, + [this.type, this.node.id, this.index] as const + ) + } +} +class NodeReference { + constructor( + readonly id: NodeId, + readonly comfyPage: ComfyPage + ) {} + async exists(): Promise { + return await this.comfyPage.page.evaluate((id) => { + const node = window['app'].graph.getNodeById(id) + return !!node + }, this.id) + } + getType(): Promise { + return this.getProperty('type') + } + async getPosition(): Promise { + const pos = await this.comfyPage.convertOffsetToCanvas( + await this.getProperty<[number, number]>('pos') + ) + return { + x: pos[0], + y: pos[1] + } + } + async getSize(): Promise { + const size = await this.getProperty<[number, number]>('size') + return { + width: size[0], + height: size[1] + } + } + async getProperty(prop: string): Promise { + return await this.comfyPage.page.evaluate( + ([id, prop]) => { + const node = window['app'].graph.getNodeById(id) + if (!node) throw new Error('Node not found') + return node[prop] + }, + [this.id, prop] as const + ) + } + async getOutput(index: number) { + return new NodeSlotReference('output', index, this) + } + async getInput(index: number) { + return new NodeSlotReference('input', index, this) + } + async click(position: 'title', options?: Parameters[1]) { + const nodePos = await this.getPosition() + const nodeSize = await this.getSize() + let clickPos: Position + switch (position) { + case 'title': + clickPos = { x: nodePos.x + nodeSize.width / 2, y: nodePos.y + 15 } + break + default: + throw new Error(`Invalid click position ${position}`) + } + await this.comfyPage.canvas.click({ + ...options, + position: clickPos + }) + await this.comfyPage.nextFrame() + } + async connectOutput( + originSlotIndex: number, + targetNode: NodeReference, + targetSlotIndex: number + ) { + const originSlot = await this.getOutput(originSlotIndex) + const targetSlot = await targetNode.getInput(targetSlotIndex) + await this.comfyPage.dragAndDrop( + await originSlot.getPosition(), + await targetSlot.getPosition() + ) + return originSlot + } + async clickContextMenuOption(optionText: string) { + await this.click('title', { button: 'right' }) + const ctx = this.comfyPage.page.locator('.litecontextmenu') + await ctx.getByText(optionText).click() + } + async convertToGroupNode(groupNodeName: string = 'GroupNode') { + this.comfyPage.page.once('dialog', async (dialog) => { + await dialog.accept(groupNodeName) + }) + await this.clickContextMenuOption('Convert to Group Node') + await this.comfyPage.nextFrame() + const nodes = await this.comfyPage.getNodeRefsByType( + `workflow/${groupNodeName}` + ) + if (nodes.length !== 1) { + throw new Error(`Did not find single group node (found=${nodes.length})`) + } + return nodes[0] + } + async manageGroupNode() { + await this.clickContextMenuOption('Manage Group Node') + await this.comfyPage.nextFrame() + return new ManageGroupNode( + this.comfyPage.page, + this.comfyPage.page.locator('.comfy-group-manage') + ) + } } export const comfyPageFixture = base.extend<{ comfyPage: ComfyPage }>({ diff --git a/browser_tests/groupNode.spec.ts b/browser_tests/groupNode.spec.ts index 8e5d69ed27..97bdecd873 100644 --- a/browser_tests/groupNode.spec.ts +++ b/browser_tests/groupNode.spec.ts @@ -53,4 +53,40 @@ test.describe('Group Node', () => { await comfyPage.page.waitForTimeout(tooltipTimeout + 16) await expect(comfyPage.page.locator('.node-tooltip')).toBeVisible() }) + + test('Reconnects inputs after configuration changed via manage dialog save', async ({ + comfyPage + }) => { + const expectSingleNode = async (type: string) => { + const nodes = await comfyPage.getNodeRefsByType(type) + expect(nodes).toHaveLength(1) + return nodes[0] + } + const latent = await expectSingleNode('EmptyLatentImage') + const sampler = await expectSingleNode('KSampler') + // Remove existing link + const samplerInput = await sampler.getInput(0) + await samplerInput.removeLinks() + // Group latent + sampler + await latent.click('title', { + modifiers: ['Shift'] + }) + await sampler.click('title', { + modifiers: ['Shift'] + }) + const groupNode = await sampler.convertToGroupNode() + // Connect node to group + const ckpt = await expectSingleNode('CheckpointLoaderSimple') + const input = await ckpt.connectOutput(0, groupNode, 0) + expect(await input.getLinkCount()).toBe(1) + // Modify the group node via manage dialog + const manage = await groupNode.manageGroupNode() + await manage.selectNode('KSampler') + await manage.changeTab('Inputs') + await manage.setLabel('model', 'test') + await manage.save() + await manage.close() + // Ensure the link is still present + expect(await input.getLinkCount()).toBe(1) + }) }) diff --git a/browser_tests/helpers/manageGroupNode.ts b/browser_tests/helpers/manageGroupNode.ts new file mode 100644 index 0000000000..27bc98346c --- /dev/null +++ b/browser_tests/helpers/manageGroupNode.ts @@ -0,0 +1,37 @@ +import { Locator, Page } from '@playwright/test' +export class ManageGroupNode { + footer: Locator + + constructor( + readonly page: Page, + readonly root: Locator + ) { + this.footer = root.locator('footer') + } + + async setLabel(name: string, label: string) { + const active = this.root.locator('.comfy-group-manage-node-page.active') + const input = active.getByPlaceholder(name) + await input.fill(label) + } + + async save() { + await this.footer.getByText('Save').click() + } + + async close() { + await this.footer.getByText('Close').click() + } + + async selectNode(name: string) { + const list = this.root.locator('.comfy-group-manage-list-items') + const item = list.getByText(name) + await item.click() + } + + async changeTab(name: 'Inputs' | 'Widgets' | 'Outputs') { + const header = this.root.locator('.comfy-group-manage-node header') + const tab = header.getByText(name) + await tab.click() + } +} diff --git a/src/extensions/core/groupNode.ts b/src/extensions/core/groupNode.ts index 822df7b7d1..34e743dd4a 100644 --- a/src/extensions/core/groupNode.ts +++ b/src/extensions/core/groupNode.ts @@ -828,6 +828,9 @@ export class GroupNodeHandler { ] // Remove all converted nodes and relink them + const builder = new GroupNodeBuilder(nodes) + const nodeData = builder.getNodeData() + groupNode[GROUP].groupData.nodeData.links = nodeData.links groupNode[GROUP].replaceNodes(nodes) return groupNode }