ComfyUI/tests-ui/tests/widgetInputs.test.js

/// <reference path="../node_modules/@types/jest/index.d.ts" />
// @ts-check

const { start } = require("../utils");
const lg = require("../utils/litegraph");


describe("widget inputs", () => {
	beforeEach(() => {
		lg.setup(global);
	});

	afterEach(() => {
		lg.teardown(global);
		jest.resetModules();
	});

	[
		{ 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 () => {
			/**
			 * Test node with widgets of each type
			 * @type { import("../../web/types/comfy").ComfyObjectInfo } ComfyObjectInfo
			 */
			const WidgetTestNode = {
				category: "test",
				name: "WidgetTestNode",
				output_name: [],
				input: {
					required: {
						[c.name]: [c.type, c.opt ?? {}],
					},
				},
			};

			const { ez } = await start({
				mockNodeDefs: {
					WidgetTestNode,
				},
			});

			// Create test node and convert to input
			const n = ez.WidgetTestNode();
			const w = n.widgets[c.name];
			w.convertToInput();
			expect(w.isConvertedToInput).toBeTruthy();
			const input = w.getConvertedInput();
			expect(input).toBeTruthy();

			// Connect to primitive
			const p1 = ez.PrimitiveNode();
			// @ts-ignore : input is valid
			p1.outputs[0].connectTo(input);
			expect(p1.outputs[0].connectTo).toHaveLength(1);

			// Ensure widget is correct type
			const valueWidget = p1.widgets.value;
			expect(valueWidget.widget.type).toBe(c.widget ?? c.name);

			// Check if control_after_generate should be added
			if (c.control) {
				const controlWidget = p1.widgets.control_after_generate;
				expect(controlWidget.widget.type).toBe("combo");
			}

			// Ensure we dont have other widgets
			expect(p1.node.widgets).toHaveLength(1 + +!!c.control);
		});
	});

	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).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);

		let primitive = ez.PrimitiveNode();
		primitive.outputs[0].connectTo(n.inputs.ckpt_name);

		await graph.reload();

		// Find the reloaded nodes in the graph
		n = graph.find(n);
		primitive = graph.find(primitive);

		// Ensure widget is converted
		expect(n.widgets.ckpt_name.isConvertedToInput).toBeTruthy();
		expect(n.inputs.ckpt_name).toBeTruthy();
		expect(n.inputs.length).toEqual(inputCount + 1);

		// Ensure primitive is connected
		let { connections } = primitive.outputs[0];
		expect(connections).toHaveLength(1);
		expect(connections[0].targetNode.id).toBe(n.node.id);

		// Disconnect & reconnect
		connections[0].disconnect();
		({ 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();
	});
});