Unit tests + widget input fixes (#1760)

* setup ui unit tests

* Refactoring, adding connections

* Few tweaks

* Fix type

* Add general test

* Refactored and extended test

* move to describe

* for groups

* Add test for converted widgets on missing nodes + fix crash

* tidy

* mores tests + refactor

* throw earlier to get less confusing error

* support outputs

* more test

* add ci action

* use lts node

* Fix?

* Prevent connecting non matching combos

* update

* accidently removed npm i

* Disable logging extension

* added step to generate object_info

* fix python

* install python

* install deps

* fix cwd?

* logging

* Fix double resolve

* create dir

* update pkg

.github/workflows/test-ui.yaml

+name: Tests CI
+
+on: [push, pull_request]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - uses: actions/setup-node@v3
+      with:
+        node-version: 18
+    - uses: actions/setup-python@v4
+      with: 
+        python-version: '3.10'
+    - name: Install requirements
+      run: |
+        python -m pip install --upgrade pip
+        pip install -r requirements.txt
+    - name: Run Tests
+      run: | 
+        npm install
+        npm run test:generate
+        npm test
+      working-directory: ./tests-ui

.gitignore

@@ -14,3 +14,4 @@ venv/
 /web/extensions/*
 !/web/extensions/logging.js.example
 !/web/extensions/core/
+/tests-ui/data/object_info.json
\ No newline at end of file

tests-ui/.gitignore

+node_modules
\ No newline at end of file

tests-ui/babel.config.json

+{
+	"presets": ["@babel/preset-env"]
+}

tests-ui/globalSetup.js

+module.exports = async function () {
+	global.ResizeObserver = class ResizeObserver {
+		observe() {}
+		unobserve() {}
+		disconnect() {}
+	};
+
+	const { nop } = require("./utils/nopProxy");
+	global.enableWebGLCanvas = nop;
+
+	HTMLCanvasElement.prototype.getContext = nop;
+
+	localStorage["Comfy.Settings.Comfy.Logging.Enabled"] = "false";
+};

tests-ui/jest.config.js

+/** @type {import('jest').Config} */
+const config = {
+	testEnvironment: "jsdom",
+	setupFiles: ["./globalSetup.js"],
+	clearMocks: true,
+	resetModules: true,
+};
+
+module.exports = config;

tests-ui/package.json

+{
+  "name": "comfui-tests",
+  "version": "1.0.0",
+  "description": "UI tests",
+  "main": "index.js",
+  "scripts": {
+    "test": "jest",
+    "test:generate": "node setup.js"
+  },
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/comfyanonymous/ComfyUI.git"
+  },
+  "keywords": [
+    "comfyui",
+    "test"
+  ],
+  "author": "comfyanonymous",
+  "license": "GPL-3.0",
+  "bugs": {
+    "url": "https://github.com/comfyanonymous/ComfyUI/issues"
+  },
+  "homepage": "https://github.com/comfyanonymous/ComfyUI#readme",
+  "devDependencies": {
+    "@babel/preset-env": "^7.22.20",
+    "@types/jest": "^29.5.5",
+    "jest": "^29.7.0",
+    "jest-environment-jsdom": "^29.7.0"
+  }
+}

tests-ui/setup.js

+const { spawn } = require("child_process");
+const { resolve } = require("path");
+const { existsSync, mkdirSync, writeFileSync } = require("fs");
+const http = require("http");
+
+async function setup() {
+	// Wait up to 30s for it to start
+	let success = false;
+	let child;
+	for (let i = 0; i < 30; i++) {
+		try {
+			await new Promise((res, rej) => {
+				http
+					.get("http://127.0.0.1:8188/object_info", (resp) => {
+						let data = "";
+						resp.on("data", (chunk) => {
+							data += chunk;
+						});
+						resp.on("end", () => {
+							// Modify the response data to add some checkpoints
+							const objectInfo = JSON.parse(data);
+							objectInfo.CheckpointLoaderSimple.input.required.ckpt_name[0] = ["model1.safetensors", "model2.ckpt"];
+
+							data = JSON.stringify(objectInfo, undefined, "\t");
+
+							const outDir = resolve("./data");
+							if (!existsSync(outDir)) {
+								mkdirSync(outDir);
+							}
+
+							const outPath = resolve(outDir, "object_info.json");
+							console.log(`Writing ${Object.keys(objectInfo).length} nodes to ${outPath}`);
+							writeFileSync(outPath, data, {
+								encoding: "utf8",
+							});
+							res();
+						});
+					})
+					.on("error", rej);
+			});
+			success = true;
+			break;
+		} catch (error) {
+			console.log(i + "/30", error);
+			if (i === 0) {
+				// Start the server on first iteration if it fails to connect
+				console.log("Starting ComfyUI server...");
+
+				let python = resolve("../../python_embeded/python.exe");
+				let args;
+				let cwd;
+				if (existsSync(python)) {
+					args = ["-s", "ComfyUI/main.py"];
+					cwd = "../..";
+				} else {
+					python = "python";
+					args = ["main.py"];
+					cwd = "..";
+				}
+				args.push("--cpu");
+				console.log(python, ...args);
+				child = spawn(python, args, { cwd });
+				child.on("error", (err) => {
+					console.log(`Server error (${err})`);
+					i = 30;
+				});
+				child.on("exit", (code) => {
+					if (!success) {
+						console.log(`Server exited (${code})`);
+						i = 30;
+					}
+				});
+			}
+			await new Promise((r) => {
+				setTimeout(r, 1000);
+			});
+		}
+	}
+
+	child?.kill();
+
+	if (!success) {
+		throw new Error("Waiting for server failed...");
+	}
+}
+
+ setup();
\ No newline at end of file

tests-ui/tests/widgetInputs.test.js

+// @ts-check
+/// <reference path="../node_modules/@types/jest/index.d.ts" />
+
+const { start, makeNodeDef, checkBeforeAndAfterReload, assertNotNullOrUndefined } = require("../utils");
+const lg = require("../utils/litegraph");
+
+/**
+ * @typedef { import("../utils/ezgraph") } Ez
+ * @typedef { ReturnType<Ez["Ez"]["graph"]>["ez"] } EzNodeFactory
+ */
+
+/**
+ * @param { EzNodeFactory } ez
+ * @param { InstanceType<Ez["EzGraph"]> } graph
+ * @param { InstanceType<Ez["EzInput"]> } input
+ * @param { string } widgetType
+ * @param { boolean } hasControlWidget
+ * @returns
+ */
+async function connectPrimitiveAndReload(ez, graph, input, widgetType, hasControlWidget) {
+	// Connect to primitive and ensure its still connected after
+	let primitive = ez.PrimitiveNode();
+	primitive.outputs[0].connectTo(input);
+
+	await checkBeforeAndAfterReload(graph, async () => {
+		primitive = graph.find(primitive);
+		let { connections } = primitive.outputs[0];
+		expect(connections).toHaveLength(1);
+		expect(connections[0].targetNode.id).toBe(input.node.node.id);
+
+		// Ensure widget is correct type
+		const valueWidget = primitive.widgets.value;
+		expect(valueWidget.widget.type).toBe(widgetType);
+
+		// Check if control_after_generate should be added
+		if (hasControlWidget) {
+			const controlWidget = primitive.widgets.control_after_generate;
+			expect(controlWidget.widget.type).toBe("combo");
+		}
+
+		// Ensure we dont have other widgets
+		expect(primitive.node.widgets).toHaveLength(1 + +!!hasControlWidget);
+	});
+
+	return primitive;
+}
+
+describe("widget inputs", () => {
+	beforeEach(() => {
+		lg.setup(global);
+	});
+
+	afterEach(() => {
+		lg.teardown(global);
+	});
+
+	[
+		{ name: "int", type: "INT", widget: "number", control: true },
+		{ name: "float", type: "FLOAT", widget: "number", control: true },
+		{ name: "text", type: "STRING" },
+		{
+			name: "customtext",
+			type: "STRING",
+			opt: { multiline: true },
+		},
+		{ name: "toggle", type: "BOOLEAN" },
+		{ name: "combo", type: ["a", "b", "c"], control: true },
+	].forEach((c) => {
+		test(`widget conversion + primitive works on ${c.name}`, async () => {
+			const { ez, graph } = await start({
+				mockNodeDefs: makeNodeDef("TestNode", { [c.name]: [c.type, c.opt ?? {}] }),
+			});
+
+			// Create test node and convert to input
+			const n = ez.TestNode();
+			const w = n.widgets[c.name];
+			w.convertToInput();
+			expect(w.isConvertedToInput).toBeTruthy();
+			const input = w.getConvertedInput();
+			expect(input).toBeTruthy();
+
+			// @ts-ignore : input is valid here
+			await connectPrimitiveAndReload(ez, graph, input, c.widget ?? c.name, c.control);
+		});
+	});
+
+	test("converted widget works after reload", async () => {
+		const { ez, graph } = await start();
+		let n = ez.CheckpointLoaderSimple();
+
+		const inputCount = n.inputs.length;
+
+		// Convert ckpt name to an input
+		n.widgets.ckpt_name.convertToInput();
+		expect(n.widgets.ckpt_name.isConvertedToInput).toBeTruthy();
+		expect(n.inputs.ckpt_name).toBeTruthy();
+		expect(n.inputs.length).toEqual(inputCount + 1);
+
+		// Convert back to widget and ensure input is removed
+		n.widgets.ckpt_name.convertToWidget();
+		expect(n.widgets.ckpt_name.isConvertedToInput).toBeFalsy();
+		expect(n.inputs.ckpt_name).toBeFalsy();
+		expect(n.inputs.length).toEqual(inputCount);
+
+		// Convert again and reload the graph to ensure it maintains state
+		n.widgets.ckpt_name.convertToInput();
+		expect(n.inputs.length).toEqual(inputCount + 1);
+
+		const primitive = await connectPrimitiveAndReload(ez, graph, n.inputs.ckpt_name, "combo", true);
+
+		// Disconnect & reconnect
+		primitive.outputs[0].connections[0].disconnect();
+		let { connections } = primitive.outputs[0];
+		expect(connections).toHaveLength(0);
+
+		primitive.outputs[0].connectTo(n.inputs.ckpt_name);
+		({ connections } = primitive.outputs[0]);
+		expect(connections).toHaveLength(1);
+		expect(connections[0].targetNode.id).toBe(n.node.id);
+
+		// Convert back to widget and ensure input is removed
+		n.widgets.ckpt_name.convertToWidget();
+		expect(n.widgets.ckpt_name.isConvertedToInput).toBeFalsy();
+		expect(n.inputs.ckpt_name).toBeFalsy();
+		expect(n.inputs.length).toEqual(inputCount);
+	});
+
+	test("converted widget works on clone", async () => {
+		const { graph, ez } = await start();
+		let n = ez.CheckpointLoaderSimple();
+
+		// Convert the widget to an input
+		n.widgets.ckpt_name.convertToInput();
+		expect(n.widgets.ckpt_name.isConvertedToInput).toBeTruthy();
+
+		// Clone the node
+		n.menu["Clone"].call();
+		expect(graph.nodes).toHaveLength(2);
+		const clone = graph.nodes[1];
+		expect(clone.id).not.toEqual(n.id);
+
+		// Ensure the clone has an input
+		expect(clone.widgets.ckpt_name.isConvertedToInput).toBeTruthy();
+		expect(clone.inputs.ckpt_name).toBeTruthy();
+
+		// Ensure primitive connects to both nodes
+		let primitive = ez.PrimitiveNode();
+		primitive.outputs[0].connectTo(n.inputs.ckpt_name);
+		primitive.outputs[0].connectTo(clone.inputs.ckpt_name);
+		expect(primitive.outputs[0].connections).toHaveLength(2);
+
+		// Convert back to widget and ensure input is removed
+		clone.widgets.ckpt_name.convertToWidget();
+		expect(clone.widgets.ckpt_name.isConvertedToInput).toBeFalsy();
+		expect(clone.inputs.ckpt_name).toBeFalsy();
+	});
+
+	test("shows missing node error on custom node with converted input", async () => {
+		const { graph } = await start();
+
+		const dialogShow = jest.spyOn(graph.app.ui.dialog, "show");
+
+		await graph.app.loadGraphData({
+			last_node_id: 3,
+			last_link_id: 4,
+			nodes: [
+				{
+					id: 1,
+					type: "TestNode",
+					pos: [41.87329101561909, 389.7381480823742],
+					size: { 0: 220, 1: 374 },
+					flags: {},
+					order: 1,
+					mode: 0,
+					inputs: [{ name: "test", type: "FLOAT", link: 4, widget: { name: "test" }, slot_index: 0 }],
+					outputs: [],
+					properties: { "Node name for S&R": "TestNode" },
+					widgets_values: [1],
+				},
+				{
+					id: 3,
+					type: "PrimitiveNode",
+					pos: [-312, 433],
+					size: { 0: 210, 1: 82 },
+					flags: {},
+					order: 0,
+					mode: 0,
+					outputs: [{ links: [4], widget: { name: "test" } }],
+					title: "test",
+					properties: {},
+				},
+			],
+			links: [[4, 3, 0, 1, 6, "FLOAT"]],
+			groups: [],
+			config: {},
+			extra: {},
+			version: 0.4,
+		});
+
+		expect(dialogShow).toBeCalledTimes(1);
+		expect(dialogShow.mock.calls[0][0]).toContain("the following node types were not found");
+		expect(dialogShow.mock.calls[0][0]).toContain("TestNode");
+	});
+
+	test("defaultInput widgets can be converted back to inputs", async () => {
+		const { graph, ez } = await start({
+			mockNodeDefs: makeNodeDef("TestNode", { example: ["INT", { defaultInput: true }] }),
+		});
+
+		// Create test node and ensure it starts as an input
+		let n = ez.TestNode();
+		let w = n.widgets.example;
+		expect(w.isConvertedToInput).toBeTruthy();
+		let input = w.getConvertedInput();
+		expect(input).toBeTruthy();
+
+		// Ensure it can be converted to
+		w.convertToWidget();
+		expect(w.isConvertedToInput).toBeFalsy();
+		expect(n.inputs.length).toEqual(0);
+		// and from
+		w.convertToInput();
+		expect(w.isConvertedToInput).toBeTruthy();
+		input = w.getConvertedInput();
+
+		// Reload and ensure it still only has 1 converted widget
+		if (!assertNotNullOrUndefined(input)) return;
+
+		await connectPrimitiveAndReload(ez, graph, input, "number", true);
+		n = graph.find(n);
+		expect(n.widgets).toHaveLength(1);
+		w = n.widgets.example;
+		expect(w.isConvertedToInput).toBeTruthy();
+
+		// Convert back to widget and ensure it is still a widget after reload
+		w.convertToWidget();
+		await graph.reload();
+		n = graph.find(n);
+		expect(n.widgets).toHaveLength(1);
+		expect(n.widgets[0].isConvertedToInput).toBeFalsy();
+		expect(n.inputs.length).toEqual(0);
+	});
+
+	test("forceInput widgets can not be converted back to inputs", async () => {
+		const { graph, ez } = await start({
+			mockNodeDefs: makeNodeDef("TestNode", { example: ["INT", { forceInput: true }] }),
+		});
+
+		// Create test node and ensure it starts as an input
+		let n = ez.TestNode();
+		let w = n.widgets.example;
+		expect(w.isConvertedToInput).toBeTruthy();
+		const input = w.getConvertedInput();
+		expect(input).toBeTruthy();
+
+		// Convert to widget should error
+		expect(() => w.convertToWidget()).toThrow();
+
+		// Reload and ensure it still only has 1 converted widget
+		if (assertNotNullOrUndefined(input)) {
+			await connectPrimitiveAndReload(ez, graph, input, "number", true);
+			n = graph.find(n);
+			expect(n.widgets).toHaveLength(1);
+			expect(n.widgets.example.isConvertedToInput).toBeTruthy();
+		}
+	});
+
+	test("primitive can connect to matching combos on converted widgets", async () => {
+		const { ez } = await start({
+			mockNodeDefs: {
+				...makeNodeDef("TestNode1", { example: [["A", "B", "C"], { forceInput: true }] }),
+				...makeNodeDef("TestNode2", { example: [["A", "B", "C"], { forceInput: true }] }),
+			},
+		});
+
+		const n1 = ez.TestNode1();
+		const n2 = ez.TestNode2();
+		const p = ez.PrimitiveNode();
+		p.outputs[0].connectTo(n1.inputs[0]);
+		p.outputs[0].connectTo(n2.inputs[0]);
+		expect(p.outputs[0].connections).toHaveLength(2);
+		const valueWidget = p.widgets.value;
+		expect(valueWidget.widget.type).toBe("combo");
+		expect(valueWidget.widget.options.values).toEqual(["A", "B", "C"]);
+	});
+
+	test("primitive can not connect to non matching combos on converted widgets", async () => {
+		const { ez } = await start({
+			mockNodeDefs: {
+				...makeNodeDef("TestNode1", { example: [["A", "B", "C"], { forceInput: true }] }),
+				...makeNodeDef("TestNode2", { example: [["A", "B"], { forceInput: true }] }),
+			},
+		});
+
+		const n1 = ez.TestNode1();
+		const n2 = ez.TestNode2();
+		const p = ez.PrimitiveNode();
+		p.outputs[0].connectTo(n1.inputs[0]);
+		expect(() => p.outputs[0].connectTo(n2.inputs[0])).toThrow();
+		expect(p.outputs[0].connections).toHaveLength(1);
+	});
+
+	test("combo output can not connect to non matching combos list input", async () => {
+		const { ez } = await start({
+			mockNodeDefs: {
+				...makeNodeDef("TestNode1", {}, [["A", "B"]]),
+				...makeNodeDef("TestNode2", { example: [["A", "B"], { forceInput: true}] }),
+				...makeNodeDef("TestNode3", { example: [["A", "B", "C"], { forceInput: true}] }),
+			},
+		});
+
+		const n1 = ez.TestNode1();
+		const n2 = ez.TestNode2();
+		const n3 = ez.TestNode3();
+
+		n1.outputs[0].connectTo(n2.inputs[0]);
+		expect(() => n1.outputs[0].connectTo(n3.inputs[0])).toThrow();
+	});
+});

tests-ui/utils/ezgraph.js

+// @ts-check
+/// <reference path="../../web/types/litegraph.d.ts" />
+
+/**
+ * @typedef { import("../../web/scripts/app")["app"] } app
+ * @typedef { import("../../web/types/litegraph") } LG
+ * @typedef { import("../../web/types/litegraph").IWidget } IWidget
+ * @typedef { import("../../web/types/litegraph").ContextMenuItem } ContextMenuItem
+ * @typedef { import("../../web/types/litegraph").INodeInputSlot } INodeInputSlot
+ * @typedef { import("../../web/types/litegraph").INodeOutputSlot } INodeOutputSlot
+ * @typedef { InstanceType<LG["LGraphNode"]> & { widgets?: Array<IWidget> } } LGNode
+ * @typedef { (...args: EzOutput[] | [...EzOutput[], Record<string, unknown>]) => EzNode } EzNodeFactory
+ */
+
+export class EzConnection {
+	/** @type { app } */
+	app;
+	/** @type { InstanceType<LG["LLink"]> } */
+	link;
+
+	get originNode() {
+		return new EzNode(this.app, this.app.graph.getNodeById(this.link.origin_id));
+	}
+
+	get originOutput() {
+		return this.originNode.outputs[this.link.origin_slot];
+	}
+
+	get targetNode() {
+		return new EzNode(this.app, this.app.graph.getNodeById(this.link.target_id));
+	}
+
+	get targetInput() {
+		return this.targetNode.inputs[this.link.target_slot];
+	}
+
+	/**
+	 * @param { app } app
+	 * @param { InstanceType<LG["LLink"]> } link
+	 */
+	constructor(app, link) {
+		this.app = app;
+		this.link = link;
+	}
+
+	disconnect() {
+		this.targetInput.disconnect();
+	}
+}
+
+export class EzSlot {
+	/** @type { EzNode } */
+	node;
+	/** @type { number } */
+	index;
+
+	/**
+	 * @param { EzNode } node
+	 * @param { number } index
+	 */
+	constructor(node, index) {
+		this.node = node;
+		this.index = index;
+	}
+}
+
+export class EzInput extends EzSlot {
+	/** @type { INodeInputSlot } */
+	input;
+
+	/**
+	 * @param { EzNode } node
+	 * @param { number } index
+	 * @param { INodeInputSlot } input
+	 */
+	constructor(node, index, input) {
+		super(node, index);
+		this.input = input;
+	}
+
+	disconnect() {
+		this.node.node.disconnectInput(this.index);
+	}
+}
+
+export class EzOutput extends EzSlot {
+	/** @type { INodeOutputSlot } */
+	output;
+
+	/**
+	 * @param { EzNode } node
+	 * @param { number } index
+	 * @param { INodeOutputSlot } output
+	 */
+	constructor(node, index, output) {
+		super(node, index);
+		this.output = output;
+	}
+
+	get connections() {
+		return (this.node.node.outputs?.[this.index]?.links ?? []).map(
+			(l) => new EzConnection(this.node.app, this.node.app.graph.links[l])
+		);
+	}
+
+	/**
+	 * @param { EzInput } input
+	 */
+	connectTo(input) {
+		if (!input) throw new Error("Invalid input");
+
+		/**
+		 * @type { LG["LLink"] | null }
+		 */
+		const link = this.node.node.connect(this.index, input.node.node, input.index);
+		if (!link) {
+			const inp = input.input;
+			const inName = inp.name || inp.label || inp.type;
+			throw new Error(
+				`Connecting from ${input.node.node.type}[${inName}#${input.index}] -> ${this.node.node.type}[${
+					this.output.name ?? this.output.type
+				}#${this.index}] failed.`
+			);
+		}
+		return link;
+	}
+}
+
+export class EzNodeMenuItem {
+	/** @type { EzNode } */
+	node;
+	/** @type { number } */
+	index;
+	/** @type { ContextMenuItem } */
+	item;
+
+	/**
+	 * @param { EzNode } node
+	 * @param { number } index
+	 * @param { ContextMenuItem } item
+	 */
+	constructor(node, index, item) {
+		this.node = node;
+		this.index = index;
+		this.item = item;
+	}
+
+	call(selectNode = true) {
+		if (!this.item?.callback) throw new Error(`Menu Item ${this.item?.content ?? "[null]"} has no callback.`);
+		if (selectNode) {
+			this.node.select();
+		}
+		this.item.callback.call(this.node.node, undefined, undefined, undefined, undefined, this.node.node);
+	}
+}
+
+export class EzWidget {
+	/** @type { EzNode } */
+	node;
+	/** @type { number } */
+	index;
+	/** @type { IWidget } */
+	widget;
+
+	/**
+	 * @param { EzNode } node
+	 * @param { number } index
+	 * @param { IWidget } widget
+	 */
+	constructor(node, index, widget) {
+		this.node = node;
+		this.index = index;
+		this.widget = widget;
+	}
+
+	get value() {
+		return this.widget.value;
+	}
+
+	set value(v) {
+		this.widget.value = v;
+	}
+
+	get isConvertedToInput() {
+		// @ts-ignore : this type is valid for converted widgets
+		return this.widget.type === "converted-widget";
+	}
+
+	getConvertedInput() {
+		if (!this.isConvertedToInput) throw new Error(`Widget ${this.widget.name} is not converted to input.`);
+
+		return this.node.inputs.find((inp) => inp.input["widget"]?.name === this.widget.name);
+	}
+
+	convertToWidget() {
+		if (!this.isConvertedToInput)
+			throw new Error(`Widget ${this.widget.name} cannot be converted as it is already a widget.`);
+		this.node.menu[`Convert ${this.widget.name} to widget`].call();
+	}
+
+	convertToInput() {
+		if (this.isConvertedToInput)
+			throw new Error(`Widget ${this.widget.name} cannot be converted as it is already an input.`);
+		this.node.menu[`Convert ${this.widget.name} to input`].call();
+	}
+}
+
+export class EzNode {
+	/** @type { app } */
+	app;
+	/** @type { LGNode } */
+	node;
+
+	/**
+	 * @param { app } app
+	 * @param { LGNode } node
+	 */
+	constructor(app, node) {
+		this.app = app;
+		this.node = node;
+	}
+
+	get id() {
+		return this.node.id;
+	}
+
+	get inputs() {
+		return this.#makeLookupArray("inputs", "name", EzInput);
+	}
+
+	get outputs() {
+		return this.#makeLookupArray("outputs", "name", EzOutput);
+	}
+
+	get widgets() {
+		return this.#makeLookupArray("widgets", "name", EzWidget);
+	}
+
+	get menu() {
+		return this.#makeLookupArray(() => this.app.canvas.getNodeMenuOptions(this.node), "content", EzNodeMenuItem);
+	}
+
+	select() {
+		this.app.canvas.selectNode(this.node);
+	}
+
+	// /**
+	//  * @template { "inputs" | "outputs" } T
+	//  * @param { T } type
+	//  * @returns { Record<string, type extends "inputs" ? EzInput : EzOutput> & (type extends "inputs" ? EzInput [] : EzOutput[]) }
+	//  */
+	// #getSlotItems(type) {
+	// 	// @ts-ignore : these items are correct
+	// 	return (this.node[type] ?? []).reduce((p, s, i) => {
+	// 		if (s.name in p) {
+	// 			throw new Error(`Unable to store input ${s.name} on array as name conflicts.`);
+	// 		}
+	// 		// @ts-ignore
+	// 		p.push((p[s.name] = new (type === "inputs" ? EzInput : EzOutput)(this, i, s)));
+	// 		return p;
+	// 	}, Object.assign([], { $: this }));
+	// }
+
+	/**
+	 * @template { { new(node: EzNode, index: number, obj: any): any } } T
+	 * @param { "inputs" | "outputs" | "widgets" | (() => Array<unknown>) } nodeProperty
+	 * @param { string } nameProperty
+	 * @param { T } ctor
+	 * @returns { Record<string, InstanceType<T>> & Array<InstanceType<T>> }
+	 */
+	#makeLookupArray(nodeProperty, nameProperty, ctor) {
+		const items = typeof nodeProperty === "function" ? nodeProperty() : this.node[nodeProperty];
+		// @ts-ignore
+		return (items ?? []).reduce((p, s, i) => {
+			if (!s) return p;
+
+			const name = s[nameProperty];
+			// @ts-ignore
+			if (!name || name in p) {
+				throw new Error(`Unable to store ${nodeProperty} ${name} on array as name conflicts.`);
+			}
+			// @ts-ignore
+			p.push((p[name] = new ctor(this, i, s)));
+			return p;
+		}, Object.assign([], { $: this }));
+	}
+}
+
+export class EzGraph {
+	/** @type { app } */
+	app;
+
+	/**
+	 * @param { app } app
+	 */
+	constructor(app) {
+		this.app = app;
+	}
+
+	get nodes() {
+		return this.app.graph._nodes.map((n) => new EzNode(this.app, n));
+	}
+
+	clear() {
+		this.app.graph.clear();
+	}
+
+	arrange() {
+		this.app.graph.arrange();
+	}
+
+	stringify() {
+		return JSON.stringify(this.app.graph.serialize(), undefined, "\t");
+	}
+
+	/**
+	 * @param { number | LGNode | EzNode } obj
+	 * @returns { EzNode }
+	 */
+	find(obj) {
+		let match;
+		let id;
+		if (typeof obj === "number") {
+			id = obj;
+		} else {
+			id = obj.id;
+		}
+
+		match = this.app.graph.getNodeById(id);
+
+		if (!match) {
+			throw new Error(`Unable to find node with ID ${id}.`);
+		}
+
+		return new EzNode(this.app, match);
+	}
+
+	/**
+	 * @returns { Promise<void> }
+	 */
+	reload() {
+		const graph = JSON.parse(JSON.stringify(this.app.graph.serialize()));
+		return new Promise((r) => {
+			this.app.graph.clear();
+			setTimeout(async () => {
+				await this.app.loadGraphData(graph);
+				r();
+			}, 10);
+		});
+	}
+}
+
+export const Ez = {
+	/**
+	 * Quickly build and interact with a ComfyUI graph
+	 * @example
+	 * const { ez, graph } = Ez.graph(app);
+	 * graph.clear();
+	 * const [model, clip, vae] = ez.CheckpointLoaderSimple();
+	 * const [pos] = ez.CLIPTextEncode(clip, { text: "positive" });
+	 * const [neg] = ez.CLIPTextEncode(clip, { text: "negative" });
+	 * const [latent] = ez.KSampler(model, pos, neg, ...ez.EmptyLatentImage());
+	 * const [image] = ez.VAEDecode(latent, vae);
+	 * const saveNode = ez.SaveImage(image).node;
+	 * console.log(saveNode);
+	 * graph.arrange();
+	 * @param { app } app
+	 * @param { LG["LiteGraph"] } LiteGraph
+	 * @param { LG["LGraphCanvas"] } LGraphCanvas
+	 * @param { boolean } clearGraph
+	 * @returns { { graph: EzGraph, ez: Record<string, EzNodeFactory> } }
+	 */
+	graph(app, LiteGraph = window["LiteGraph"], LGraphCanvas = window["LGraphCanvas"], clearGraph = true) {
+		// Always set the active canvas so things work
+		LGraphCanvas.active_canvas = app.canvas;
+
+		if (clearGraph) {
+			app.graph.clear();
+		}
+
+		// @ts-ignore : this proxy handles utility methods & node creation
+		const factory = new Proxy(
+			{},
+			{
+				get(_, p) {
+					if (typeof p !== "string") throw new Error("Invalid node");
+					const node = LiteGraph.createNode(p);
+					if (!node) throw new Error(`Unknown node "${p}"`);
+					app.graph.add(node);
+
+					/**
+					 * @param {Parameters<EzNodeFactory>} args
+					 */
+					return function (...args) {
+						const ezNode = new EzNode(app, node);
+						const inputs = ezNode.inputs;
+
+						let slot = 0;
+						for (const arg of args) {
+							if (arg instanceof EzOutput) {
+								arg.connectTo(inputs[slot++]);
+							} else {
+								for (const k in arg) {
+									ezNode.widgets[k].value = arg[k];
+								}
+							}
+						}
+
+						return ezNode;
+					};
+				},
+			}
+		);
+
+		return { graph: new EzGraph(app), ez: factory };
+	},
+};

tests-ui/utils/index.js

+const { mockApi } = require("./setup");
+const { Ez } = require("./ezgraph");
+
+/**
+ *
+ * @param { Parameters<mockApi>[0] } config
+ * @returns
+ */
+export async function start(config = undefined) {
+	mockApi(config);
+	const { app } = require("../../web/scripts/app");
+	await app.setup();
+	return Ez.graph(app, global["LiteGraph"], global["LGraphCanvas"]);
+}
+
+/**
+ * @param { ReturnType<Ez["graph"]>["graph"] } graph 
+ * @param { (hasReloaded: boolean) => (Promise<void> | void) } cb 
+ */
+export async function checkBeforeAndAfterReload(graph, cb) {
+	await cb(false);
+	await graph.reload();
+	await cb(true);
+}
+
+/**
+ * @param { string } name 
+ * @param { Record<string, string | [string | string[], any]> } input 
+ * @param { (string | string[])[] | Record<string, string | string[]> } output
+ * @returns { Record<string, import("../../web/types/comfy").ComfyObjectInfo> } 
+ */
+export function makeNodeDef(name, input, output = {}) {
+	const nodeDef = {
+		name,
+		category: "test",
+		output: [],
+		output_name: [],
+		output_is_list: [],
+		input: {
+			required: {}
+		},
+	};
+	for(const k in input) {
+		nodeDef.input.required[k] = typeof input[k] === "string" ? [input[k], {}] : [...input[k]];
+	}
+	if(output instanceof Array) {
+		output = output.reduce((p, c) => {
+			p[c] = c;
+			return p;
+		}, {})
+	}
+	for(const k in output) {
+		nodeDef.output.push(output[k]);
+		nodeDef.output_name.push(k);
+		nodeDef.output_is_list.push(false);
+	}
+
+	return { [name]: nodeDef };
+}
+
+/**
+/**
+ * @template { any } T
+ * @param { T } x
+ * @returns { x is Exclude<T, null | undefined> }
+ */
+export function assertNotNullOrUndefined(x) {
+	expect(x).not.toEqual(null);
+	expect(x).not.toEqual(undefined);
+	return true;
+}
\ No newline at end of file

tests-ui/utils/litegraph.js

+const fs = require("fs");
+const path = require("path");
+const { nop } = require("../utils/nopProxy");
+
+function forEachKey(cb) {
+	for (const k of [
+		"LiteGraph",
+		"LGraph",
+		"LLink",
+		"LGraphNode",
+		"LGraphGroup",
+		"DragAndScale",
+		"LGraphCanvas",
+		"ContextMenu",
+	]) {
+		cb(k);
+	}
+}
+
+export function setup(ctx) {
+	const lg = fs.readFileSync(path.resolve("../web/lib/litegraph.core.js"), "utf-8");
+	const globalTemp = {};
+	(function (console) {
+		eval(lg);
+	}).call(globalTemp, nop);
+
+	forEachKey((k) => (ctx[k] = globalTemp[k]));
+	require(path.resolve("../web/lib/litegraph.extensions.js"));
+}
+
+export function teardown(ctx) {
+	forEachKey((k) => delete ctx[k]);
+
+	// Clear document after each run
+	document.getElementsByTagName("html")[0].innerHTML = ""; 
+}

tests-ui/utils/nopProxy.js

+export const nop = new Proxy(function () {}, {
+	get: () => nop,
+	set: () => true,
+	apply: () => nop,
+	construct: () => nop,
+});

tests-ui/utils/setup.js

+require("../../web/scripts/api");
+
+const fs = require("fs");
+const path = require("path");
+function* walkSync(dir) {
+	const files = fs.readdirSync(dir, { withFileTypes: true });
+	for (const file of files) {
+		if (file.isDirectory()) {
+			yield* walkSync(path.join(dir, file.name));
+		} else {
+			yield path.join(dir, file.name);
+		}
+	}
+}
+
+/**
+ * @typedef { import("../../web/types/comfy").ComfyObjectInfo } ComfyObjectInfo
+ */
+
+/**
+ * @param { { mockExtensions?: string[], mockNodeDefs?: Record<string, ComfyObjectInfo> } } config
+ */
+export function mockApi({ mockExtensions, mockNodeDefs } = {}) {
+	if (!mockExtensions) {
+		mockExtensions = Array.from(walkSync(path.resolve("../web/extensions/core")))
+			.filter((x) => x.endsWith(".js"))
+			.map((x) => path.relative(path.resolve("../web"), x));
+	}
+	if (!mockNodeDefs) {
+		mockNodeDefs = JSON.parse(fs.readFileSync(path.resolve("./data/object_info.json")));
+	}
+
+	jest.mock("../../web/scripts/api", () => ({
+		get api() {
+			return {
+				addEventListener: jest.fn(),
+				getSystemStats: jest.fn(),
+				getExtensions: jest.fn(() => mockExtensions),
+				getNodeDefs: jest.fn(() => mockNodeDefs),
+				init: jest.fn(),
+				apiURL: jest.fn((x) => "../../web/" + x),
+			};
+		},
+	}));
+}

web/extensions/core/widgetInputs.js

@@ -100,6 +100,27 @@ function getWidgetType(config) {
 	return { type };
 }
 
+
+function isValidCombo(combo, obj) {
+	// New input isnt a combo
+	if (!(obj instanceof Array)) {
+		console.log(`connection rejected: tried to connect combo to ${obj}`);
+		return false;
+	}
+	// New imput combo has a different size
+	if (combo.length !== obj.length) {
+		console.log(`connection rejected: combo lists dont match`);
+		return false;
+	}
+	// New input combo has different elements
+	if (combo.find((v, i) => obj[i] !== v)) {
+		console.log(`connection rejected: combo lists dont match`);
+		return false;
+	}
+
+	return true;
+}
+
 app.registerExtension({
 	name: "Comfy.WidgetInputs",
 	async beforeRegisterNodeDef(nodeType, nodeData, app) {
@@ -256,6 +277,28 @@ app.registerExtension({
 
 			return r;
 		};
+
+		// Prevent connecting COMBO lists to converted inputs that dont match types
+		const onConnectInput = nodeType.prototype.onConnectInput;
+		nodeType.prototype.onConnectInput = function (targetSlot, type, output, originNode, originSlot) {
+			const v = onConnectInput?.(this, arguments);
+			// Not a combo, ignore
+			if (type !== "COMBO") return v;
+			// Primitive output, allow that to handle
+			if (originNode.outputs[originSlot].widget) return v;
+
+			// Ensure target is also a combo
+			const targetCombo = this.inputs[targetSlot].widget?.[GET_CONFIG]?.()?.[0];
+			if (!targetCombo || !(targetCombo instanceof Array)) return v;
+
+			// Check they match
+			const originConfig = originNode.constructor?.nodeData?.output?.[originSlot];
+			if (!originConfig || !isValidCombo(targetCombo, originConfig)) {
+				return false;
+			}
+
+			return v;
+		};
 	},
 	registerCustomNodes() {
 		class PrimitiveNode {
@@ -315,7 +358,7 @@ app.registerExtension({
 
 			onAfterGraphConfigured() {
 				if (this.outputs[0].links?.length && !this.widgets?.length) {
-					this.#onFirstConnection();
+					if (!this.#onFirstConnection()) return;
 
 					// Populate widget values from config data
 					if (this.widgets) {
@@ -386,13 +429,16 @@ app.registerExtension({
 					widget = input.widget;
 				}
 
-				const { type } = getWidgetType(widget[GET_CONFIG]());
+				const config = widget[GET_CONFIG]?.();
+				if (!config) return;
+
+				const { type } = getWidgetType(config);
 				// Update our output to restrict to the widget type
 				this.outputs[0].type = type;
 				this.outputs[0].name = type;
 				this.outputs[0].widget = widget;
 
-				this.#createWidget(widget[CONFIG] ?? widget[GET_CONFIG](), theirNode, widget.name, recreating);
+				this.#createWidget(widget[CONFIG] ?? config, theirNode, widget.name, recreating);
 			}
 
 			#createWidget(inputData, node, widgetName, recreating) {
@@ -497,21 +543,7 @@ app.registerExtension({
 				const config2 = input.widget[GET_CONFIG]();
 
 				if (config1[0] instanceof Array) {
-					// New input isnt a combo
-					if (!(config2[0] instanceof Array)) {
-						console.log(`connection rejected: tried to connect combo to ${config2[0]}`);
-						return false;
-					}
-					// New imput combo has a different size
-					if (config1[0].length !== config2[0].length) {
-						console.log(`connection rejected: combo lists dont match`);
-						return false;
-					}
-					// New input combo has different elements
-					if (config1[0].find((v, i) => config2[0][i] !== v)) {
-						console.log(`connection rejected: combo lists dont match`);
-						return false;
-					}
+					if (!isValidCombo(config1[0], config2[0])) return false;
 				} else if (config1[0] !== config2[0]) {
 					// Types dont match
 					console.log(`connection rejected: types dont match`, config1[0], config2[0]);