ComfyUI_frontend/tests-ui/utils/index.ts

import { APIConfig, mockApi } from "./setup";
import { Ez, EzGraph, EzNameSpace } from "./ezgraph";
import lg from "./litegraph";
import fs from "fs";
import path from "path";

const html = fs.readFileSync(path.resolve(__dirname, "../../index.html"))

interface StartConfig extends APIConfig {
  resetEnv?: boolean;
  preSetup?(app): Promise<void>;
  localStorage?: Record<string, string>;
}

interface StartResult {
  app: any;
  graph: EzGraph;
  ez: EzNameSpace;
}

/**
 *
 * @param { Parameters<typeof mockApi>[0] & { 
 *   resetEnv?: boolean, 
 *   preSetup?(app): Promise<void>,
 *  localStorage?: Record<string, string> 
 * } } config
 * @returns
 */
export async function start(config: StartConfig = {}): Promise<StartResult> {
  if(config.resetEnv) {
    jest.resetModules();
    jest.resetAllMocks();
        lg.setup(global);
    localStorage.clear();
    sessionStorage.clear();
  }

  Object.assign(localStorage, config.localStorage ?? {});
  document.body.innerHTML = html.toString();

	mockApi(config);
	const { app } = await import("../../src/scripts/app");
	const { LiteGraph, LGraphCanvas } = await import("@comfyorg/litegraph");
	config.preSetup?.(app);
	await app.setup();

	// @ts-ignore
	return { ...Ez.graph(app, LiteGraph, 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("./src/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) {
    // @ts-ignore
    nodeDef.output.push(output[k]);
    // @ts-ignore
    nodeDef.output_name.push(k);
    // @ts-ignore
    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 } = await import("../../src/scripts/api");
  return api.getNodeDefs();
}

export async function getNodeDef(nodeId) {
  return (await getNodeDefs())[nodeId];
}