ComfyUI/tests-ui/utils/index.js

115 lines
2.8 KiB
JavaScript

const { mockApi } = require("./setup");
const { Ez } = require("./ezgraph");
const lg = require("./litegraph");
/**
*
* @param { Parameters<mockApi>[0] & { resetEnv?: boolean, preSetup?(app): Promise<void> } } config
* @returns
*/
export async function start(config = {}) {
if(config.resetEnv) {
jest.resetModules();
jest.resetAllMocks();
lg.setup(global);
}
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];
}