first commit

tests-ui/utils/ezgraph.js

@@ -0,0 +1,452 @@
+// @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;
+	}
+
+	get connection() {
+		const link = this.node.node.inputs?.[this.index]?.link;
+		if (link == null) {
+			return null;
+		}
+		return new EzConnection(this.node.app, this.node.app.graph.links[link]);
+	}
+
+	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}#${input.node.id}[${inName}#${input.index}] -> ${this.node.node.type}#${this.node.id}[${
+					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();
+		}
+		return 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;
+		this.widget.callback?.call?.(this.widget, 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.`);
+		var menu = this.node.menu["Convert Input to Widget"].item.submenu.options;
+		var index = menu.findIndex(a => a.content == `Convert ${this.widget.name} to widget`);
+		menu[index].callback.call();
+	}
+
+	convertToInput() {
+		if (this.isConvertedToInput)
+			throw new Error(`Widget ${this.widget.name} cannot be converted as it is already an input.`);
+		var menu = this.node.menu["Convert Widget to Input"].item.submenu.options;
+		var index = menu.findIndex(a => a.content == `Convert ${this.widget.name} to input`);
+		menu[index].callback.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);
+	}
+
+	get isRemoved() {
+		return !this.app.graph.getNodeById(this.id);
+	}
+
+	select(addToSelection = false) {
+		this.app.canvas.selectNode(this.node, addToSelection);
+	}
+
+	// /**
+	//  * @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];
+			const item = new ctor(this, i, s);
+			// @ts-ignore
+			p.push(item);
+			if (name) {
+				// @ts-ignore
+				if (name in p) {
+					throw new Error(`Unable to store ${nodeProperty} ${name} on array as name conflicts.`);
+				}
+			}
+			// @ts-ignore
+			p[name] = item;
+			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);
+	}
+
+	/**
+	 * @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);
+		});
+	}
+
+	/**
+	 * @returns { Promise<{
+	 * 	workflow: {},
+	 * 	output: Record<string, {
+	 * 		class_name: string,
+	 * 		inputs: Record<string, [string, number] | unknown>
+	 * }>}> }
+	 */
+	toPrompt() {
+		// @ts-ignore
+		return this.app.graphToPrompt();
+	}
+}
+
+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().outputs;
+	 * const [pos] = ez.CLIPTextEncode(clip, { text: "positive" }).outputs;
+	 * const [neg] = ez.CLIPTextEncode(clip, { text: "negative" }).outputs;
+	 * const [latent] = ez.KSampler(model, pos, neg, ...ez.EmptyLatentImage().outputs).outputs;
+	 * const [image] = ez.VAEDecode(latent, vae).outputs;
+	 * const saveNode = ez.SaveImage(image);
+	 * 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

@@ -0,0 +1,129 @@
+const { mockApi } = require("./setup");
+const { Ez } = require("./ezgraph");
+const lg = require("./litegraph");
+const fs = require("fs");
+const path = require("path");
+
+const html = fs.readFileSync(path.resolve(__dirname, "../../web/index.html"))
+
+/**
+ *
+ * @param { Parameters<typeof mockApi>[0] & { 
+ * 	resetEnv?: boolean, 
+ * 	preSetup?(app): Promise<void>,
+ *  localStorage?: Record<string, string> 
+ * } } config
+ * @returns
+ */
+export async function start(config = {}) {
+	if(config.resetEnv) {
+		jest.resetModules();
+		jest.resetAllMocks();
+        lg.setup(global);
+		localStorage.clear();
+		sessionStorage.clear();
+	}
+
+	Object.assign(localStorage, config.localStorage ?? {});
+	document.body.innerHTML = html;
+
+	mockApi(config);
+	const { app } = require("../../web/scripts/app");
+	config.preSetup?.(app);
+	await app.setup();
+
+	return { ...Ez.graph(app, global["LiteGraph"], global["LGraphCanvas"]), app };
+}
+
+/**
+ * @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;
+}
+
+/**
+ *
+ * @param { ReturnType<Ez["graph"]>["ez"] } ez
+ * @param { ReturnType<Ez["graph"]>["graph"] } graph
+ */
+export function createDefaultWorkflow(ez, graph) {
+	graph.clear();
+	const ckpt = ez.CheckpointLoaderSimple();
+
+	const pos = ez.CLIPTextEncode(ckpt.outputs.CLIP, { text: "positive" });
+	const neg = ez.CLIPTextEncode(ckpt.outputs.CLIP, { text: "negative" });
+
+	const empty = ez.EmptyLatentImage();
+	const sampler = ez.KSampler(
+		ckpt.outputs.MODEL,
+		pos.outputs.CONDITIONING,
+		neg.outputs.CONDITIONING,
+		empty.outputs.LATENT
+	);
+
+	const decode = ez.VAEDecode(sampler.outputs.LATENT, ckpt.outputs.VAE);
+	const save = ez.SaveImage(decode.outputs.IMAGE);
+	graph.arrange();
+
+	return { ckpt, pos, neg, empty, sampler, decode, save };
+}
+
+export async function getNodeDefs() {
+	const { api } = require("../../web/scripts/api");
+	return api.getNodeDefs();
+}
+
+export async function getNodeDef(nodeId) {
+	return (await getNodeDefs())[nodeId];
+}

tests-ui/utils/litegraph.js

@@ -0,0 +1,36 @@
+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

@@ -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,82 @@
+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>,
+* 	settings?: Record<string, string>
+* 	userConfig?: {storage: "server" | "browser", users?: Record<string, any>, migrated?: boolean },
+* 	userData?: Record<string, any>
+ * }} config
+ */
+export function mockApi(config = {}) {
+	let { mockExtensions, mockNodeDefs, userConfig, settings, userData } = {
+		userConfig,
+		settings: {},
+		userData: {},
+		...config,
+	};
+	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")));
+	}
+
+	const events = new EventTarget();
+	const mockApi = {
+		addEventListener: events.addEventListener.bind(events),
+		removeEventListener: events.removeEventListener.bind(events),
+		dispatchEvent: events.dispatchEvent.bind(events),
+		getSystemStats: jest.fn(),
+		getExtensions: jest.fn(() => mockExtensions),
+		getNodeDefs: jest.fn(() => mockNodeDefs),
+		init: jest.fn(),
+		apiURL: jest.fn((x) => "../../web/" + x),
+		createUser: jest.fn((username) => {
+			if(username in userConfig.users) {
+				return { status: 400, json: () => "Duplicate" }
+			}
+			userConfig.users[username + "!"] = username;
+			return { status: 200, json: () => username + "!" }
+		}),
+		getUserConfig: jest.fn(() => userConfig ?? { storage: "browser", migrated: false }),
+		getSettings: jest.fn(() => settings),
+		storeSettings: jest.fn((v) => Object.assign(settings, v)),
+		getUserData: jest.fn((f) => {
+			if (f in userData) {
+				return { status: 200, json: () => userData[f] };
+			} else {
+				return { status: 404 };
+			}
+		}),
+		storeUserData: jest.fn((file, data) => {
+			userData[file] = data;
+		}),
+		listUserData: jest.fn(() => [])
+	};
+	jest.mock("../../web/scripts/api", () => ({
+		get api() {
+			return mockApi;
+		},
+	}));
+}