setup ui unit tests

tests-ui/.gitignore

@@ -0,0 +1 @@
+node_modules

tests-ui/babel.config.json

@@ -0,0 +1,3 @@
+{
+	"presets": ["@babel/preset-env"]
+}

tests-ui/globalSetup.js

@@ -0,0 +1,12 @@
+module.exports = async function () {
+	global.ResizeObserver = class ResizeObserver {
+		observe() {}
+		unobserve() {}
+		disconnect() {}
+	};
+
+	const { nop } = require("./utils/nopProxy");
+	global.enableWebGLCanvas = nop;
+
+	HTMLCanvasElement.prototype.getContext = nop;
+};

tests-ui/jest.config.js

@@ -0,0 +1,16 @@
+const path = require("path");
+/** @type {import('jest').Config} */
+const config = {
+	testEnvironment: "jsdom",
+	// transform: {
+	// 	"^.+\\.[t|j]sx?$": "babel-jest",
+	// },
+	setupFiles: ["./globalSetup.js"],
+	// moduleDirectories: ["node_modules", path.resolve("../web/scripts")],
+	// moduleNameMapper: {
+	// 	"./api.js": path.resolve("../web/scripts/api.js"),
+	// 	"./api": path.resolve("../web/scripts/api.js"),
+	// },
+};
+
+module.exports = config;

tests-ui/package.json

@@ -0,0 +1,29 @@
+{
+  "name": "comfui-tests",
+  "version": "1.0.0",
+  "description": "UI tests",
+  "main": "index.js",
+  "scripts": {
+    "test": "jest"
+  },
+  "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/tests/widgetInputs.test.js

@@ -0,0 +1,79 @@
+/// <reference path="../node_modules/@types/jest/index.d.ts" />
+// @ts-check
+
+const { start } = require("../utils");
+const lg = require("../utils/litegraph");
+
+beforeEach(() => {
+	lg.setup(global);
+});
+
+afterEach(() => {
+	lg.teardown(global);
+	jest.resetModules();
+});
+
+test("converted widget works after reload", async () => {
+	const { graph, ez } = 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).toThrow(/Unknown input/);
+	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);
+
+	// TODO: connect primitive
+	await graph.reload();
+	// TODO: ensure primitive connected, disconnect, reconnect
+
+	// Find the reloaded node in the graph
+	n = graph.find(n);
+	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).toThrow(/Unknown input/);
+	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();
+
+	// TODO: connect primitive to clone
+
+	// 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).toThrow(/Unknown input/);
+});

tests-ui/utils/ezgraph.js

@@ -0,0 +1,318 @@
+// @ts-check
+/// <reference path="../../web/types/litegraph.d.ts" />
+
+const NODE = Symbol();
+
+/**
+ * @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 { InstanceType<LG["LGraphNode"]> & { widgets?: Array<IWidget> } } LGNode
+ * @typedef { { [k in keyof typeof Ez["util"]]: typeof Ez["util"][k] extends (app: any, ...rest: infer A) => infer R ? (...args: A) => R : never } } EzUtils
+ * @typedef { (...args: EzOutput[] | [...EzOutput[], Record<string, unknown>]) => Array<EzOutput> & { $: EzNode, node: LG["LGraphNode"]} } EzNodeFactory
+ * @typedef { ReturnType<EzNode["outputs"]>[0] } EzOutput
+ */
+
+class EzInput {
+	/** @type { EzNode } */
+	node;
+	/** @type { INodeInputSlot } */
+	input;
+	/** @type { number } */
+	index;
+
+	/**
+	 * @param { EzNode } node
+	 * @param { INodeInputSlot } input
+	 * @param { number } index
+	 */
+	constructor(node, input, index) {
+		this.node = node;
+		this.input = input;
+		this.index = index;
+	}
+}
+
+class EzNodeMenuItem {
+	/** @type { EzNode } */
+	node;
+	/** @type { ContextMenuItem } */
+	item;
+
+	/**
+	 * @param { EzNode } node
+	 * @param { ContextMenuItem } item
+	 */
+	constructor(node, item) {
+		this.node = node;
+		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);
+	}
+}
+
+class EzWidget {
+	/** @type { EzNode } */
+	node;
+	/** @type { IWidget } */
+	widget;
+
+	/**
+	 * @param { EzNode } node
+	 * @param { IWidget } widget
+	 */
+	constructor(node, widget) {
+		this.node = node;
+		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";
+	}
+
+	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();
+	}
+}
+
+class EzNode {
+	/** @type { app } */
+	app;
+	/** @type { LGNode } */
+	node;
+	/** @type { { length: number } & Record<string, EzInput> } */
+	inputs;
+	/** @type { Record<string, EzWidget> } */
+	widgets;
+	/** @type { Record<string, EzNodeMenuItem> } */
+	menu;
+
+	/**
+	 * @param { app } app
+	 * @param { LGNode } node
+	 */
+	constructor(app, node) {
+		this.app = app;
+		this.node = node;
+
+		// @ts-ignore : this proxy returns the length
+		this.inputs = new Proxy(
+			{},
+			{
+				get: (_, p) => {
+					if (typeof p !== "string") throw new Error(`Invalid widget name.`);
+					if (p === "length") return this.node.inputs?.length ?? 0;
+					const index = this.node.inputs.findIndex((i) => i.name === p);
+					if (index === -1) throw new Error(`Unknown input "${p}" on node "${this.node.type}".`);
+					return new EzInput(this, this.node.inputs[index], index);
+				},
+			}
+		);
+
+		this.widgets = new Proxy(
+			{},
+			{
+				get: (_, p) => {
+					if (typeof p !== "string") throw new Error(`Invalid widget name.`);
+					const widget = this.node.widgets?.find((w) => w.name === p);
+					if (!widget) throw new Error(`Unknown widget "${p}" on node "${this.node.type}".`);
+
+					return new EzWidget(this, widget);
+				},
+			}
+		);
+
+		this.menu = new Proxy(
+			{},
+			{
+				get: (_, p) => {
+					if (typeof p !== "string") throw new Error(`Invalid menu item name.`);
+					const options = this.menuItems();
+					const option = options.find((o) => o?.content === p);
+					if (!option) throw new Error(`Unknown menu item "${p}" on node "${this.node.type}".`);
+
+					return new EzNodeMenuItem(this, option);
+				},
+			}
+		);
+	}
+
+	get id() {
+		return this.node.id;
+	}
+
+	menuItems() {
+		return this.app.canvas.getNodeMenuOptions(this.node);
+	}
+
+	outputs() {
+		return (
+			this.node.outputs?.map((data, index) => {
+				return {
+					[NODE]: this.node,
+					index,
+					data,
+				};
+			}) ?? []
+		);
+	}
+
+	select() {
+		this.app.canvas.selectNode(this.node);
+	}
+}
+
+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();
+	}
+
+	/**
+	 * @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(() => {
+				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, 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);
+
+						// console.log("Created " + node.type, "Populating:", args);
+						let slot = 0;
+						for (let i = 0; i < args.length; i++) {
+							const arg = args[i];
+							if (arg[NODE]) {
+								arg[NODE].connect(arg.index, node, slot++);
+							} else {
+								for (const k in arg) {
+									ezNode.widgets[k].value = arg[k];
+								}
+							}
+						}
+
+						const outputs = ezNode.outputs();
+						outputs["$"] = ezNode;
+						outputs["node"] = node;
+						return outputs;
+					};
+				},
+			}
+		);
+
+		return { graph: new EzGraph(app), ez: factory };
+	},
+};

tests-ui/utils/index.js

@@ -0,0 +1,14 @@
+const { mockApi } = require("./setup");
+const { Ez } = require("./ezgraph");
+
+/**
+ *
+ * @param { Parameters<mockApi> } 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"]);
+}

tests-ui/utils/litegraph.js

@@ -0,0 +1,33 @@
+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]);
+}

tests-ui/utils/nopProxy.js

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

tests-ui/utils/setup.js

@@ -0,0 +1,45 @@
+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),
+			};
+		},
+	}));
+}