ComfyUI-Lora-Manager/tests/frontend/core/workflowRegistry.test.js

import { beforeEach, afterEach, describe, expect, it, vi } from "vitest";

const { APP_MODULE, API_MODULE, STYLES_MODULE, REGISTRY_MODULE, appMock, apiMock, registeredExtensions } =
  vi.hoisted(() => {
    const registeredExtensions = [];
    const appMock = {
      graph: null,
      registerExtension: (ext) => registeredExtensions.push(ext),
    };
    const apiMock = {
      clientId: "client-1",
      initialClientId: null,
      addEventListener: vi.fn(),
    };
    return {
      APP_MODULE: new URL("../../../scripts/app.js", import.meta.url).pathname,
      API_MODULE: new URL("../../../scripts/api.js", import.meta.url).pathname,
      STYLES_MODULE: new URL("../../../web/comfyui/lm_styles_loader.js", import.meta.url).pathname,
      REGISTRY_MODULE: new URL("../../../web/comfyui/workflow_registry.js", import.meta.url).pathname,
      appMock,
      apiMock,
      registeredExtensions,
    };
  });

vi.mock(APP_MODULE, () => ({ app: appMock }));
vi.mock(API_MODULE, () => ({ api: apiMock }));
vi.mock(STYLES_MODULE, () => ({ ensureLmStyles: vi.fn() }));

function createTextEncodeNode({ linked = false, id = 1 } = {}) {
  const textWidget = { name: "text", type: "customtext", value: "old prompt", callback: null };
  return {
    id,
    comfyClass: "CLIPTextEncode",
    title: "CLIP Text Encode",
    mode: 0,
    properties: {},
    widgets: [textWidget, { name: "clip", type: "combo" }],
    widgets_values: ["old prompt", "clip-1"],
    inputs: [
      { name: "text", type: "STRING", widget: textWidget, link: linked ? 101 : null },
      { name: "clip", type: "CLIP", link: null },
    ],
    setDirtyCanvas: vi.fn(),
    graph: null,
  };
}

function createSubgraph({ id = "sub-1", nodes = [] } = {}) {
  const graph = {
    id,
    _nodes: nodes,
    _subgraphs: new Map(),
    getNodeById: vi.fn((nodeId) => nodes.find((n) => n.id === nodeId) ?? null),
    events: { addEventListener: vi.fn() },
  };
  for (const node of nodes) {
    node.graph = graph;
  }
  return graph;
}

function createGraph({ nodes = [], subgraphs = [] } = {}) {
  const graph = {
    id: "root",
    _nodes: nodes,
    _subgraphs: new Map(),
    getNodeById: vi.fn((nodeId) => nodes.find((n) => n.id === nodeId) ?? null),
    events: { addEventListener: vi.fn() },
  };
  for (const subgraph of subgraphs) {
    graph._subgraphs.set(subgraph.id, subgraph);
  }
  for (const node of nodes) {
    node.graph = graph;
  }
  return graph;
}

function lastRegisterPayload(fetchMock) {
  const calls = fetchMock.mock.calls.filter(
    ([url]) => url === "/api/lm/register-nodes"
  );
  expect(calls.length).toBeGreaterThan(0);
  return JSON.parse(calls[calls.length - 1][1].body);
}

describe("LoraManager.WorkflowRegistry", () => {
  let extension;
  let fetchMock;

  beforeEach(async () => {
    vi.resetModules();
    registeredExtensions.length = 0;
    appMock.graph = null;
    apiMock.addEventListener.mockClear();
    fetchMock = vi.fn().mockResolvedValue({ ok: true });
    global.fetch = fetchMock;
    await import(REGISTRY_MODULE);
    extension = registeredExtensions.find(
      (ext) => ext.name === "LoraManager.WorkflowRegistry"
    );
    expect(extension).toBeDefined();
  });

  afterEach(() => {
    delete global.fetch;
  });

  describe("refreshRegistry", () => {
    it("registers an unconnected CLIPTextEncode as a text target", async () => {
      appMock.graph = createGraph({ nodes: [createTextEncodeNode()] });

      await extension.refreshRegistry(true);

      const body = lastRegisterPayload(fetchMock);
      expect(body.nodes).toHaveLength(1);
      expect(body.nodes[0].capabilities.has_text_widget).toBe(true);
      expect(body.nodes[0].capabilities.text_widget_connected).toBe(false);
    });

    it("excludes a CLIPTextEncode whose text input is connected", async () => {
      appMock.graph = createGraph({ nodes: [createTextEncodeNode({ linked: true })] });

      await extension.refreshRegistry(true);

      const body = lastRegisterPayload(fetchMock);
      expect(body.nodes).toHaveLength(1);
      expect(body.nodes[0].capabilities.has_text_widget).toBe(false);
      expect(body.nodes[0].capabilities.text_widget_connected).toBe(true);
    });

    it("registers connection state for nodes inside subgraphs", async () => {
      const inner = createTextEncodeNode({ linked: true, id: 7 });
      const subgraph = createSubgraph({ id: "sub-1", nodes: [inner] });
      appMock.graph = createGraph({ subgraphs: [subgraph] });

      await extension.refreshRegistry(true);

      const body = lastRegisterPayload(fetchMock);
      expect(body.nodes).toHaveLength(1);
      expect(body.nodes[0].graph_id).toBe("sub-1");
      expect(body.nodes[0].node_id).toBe(7);
      expect(body.nodes[0].capabilities.text_widget_connected).toBe(true);
    });

    it("re-registers when text_widget_connected changes (fingerprint)", async () => {
      const node = createTextEncodeNode();
      appMock.graph = createGraph({ nodes: [node] });

      await extension.refreshRegistry(true);
      await extension.refreshRegistry();
      expect(
        fetchMock.mock.calls.filter(([url]) => url === "/api/lm/register-nodes")
      ).toHaveLength(1);

      node.inputs[0].link = 101;
      await extension.refreshRegistry();
      const body = lastRegisterPayload(fetchMock);
      expect(body.nodes[0].capabilities.text_widget_connected).toBe(true);
    });
  });

  describe("applyWidgetUpdate (inject_text)", () => {
    it("updates the widget value when the text input is not connected", async () => {
      const node = createTextEncodeNode();
      const callback = vi.fn();
      node.widgets[0].callback = callback;
      appMock.graph = createGraph({ nodes: [node] });
      extension.flashWidget = vi.fn();

      await extension.applyWidgetUpdate({
        node_id: 1,
        action: "inject_text",
        value: "hello",
        mode: "replace",
      });

      expect(node.widgets[0].value).toBe("hello");
      expect(node.widgets_values[0]).toBe("hello");
      expect(callback).toHaveBeenCalledWith("hello");
    });

    it("skips inject_text when the target widget is connected and self-heals the registry", async () => {
      const node = createTextEncodeNode({ linked: true });
      appMock.graph = createGraph({ nodes: [node] });
      extension.flashWidget = vi.fn();
      const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});

      await extension.applyWidgetUpdate({
        node_id: 1,
        graph_id: "root",
        action: "inject_text",
        value: "new prompt",
        mode: "replace",
      });

      expect(node.widgets[0].value).toBe("old prompt");
      expect(node.widgets_values[0]).toBe("old prompt");
      expect(warnSpy).toHaveBeenCalledWith(
        expect.stringContaining("connected to an input"),
        expect.anything(),
        expect.anything()
      );
      await vi.waitFor(() => {
        expect(
          fetchMock.mock.calls.some(([url]) => url === "/api/lm/register-nodes")
        ).toBe(true);
      });
      warnSpy.mockRestore();
    });
  });

  describe("loadWorkflowFromMessage", () => {
    beforeEach(() => {
      appMock.loadApiJson = vi.fn().mockResolvedValue(undefined);
      appMock.loadGraphData = vi.fn().mockResolvedValue(undefined);
    });

    it("warns and returns when the message carries no workflow payload", async () => {
      const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});

      await extension.loadWorkflowFromMessage({ name: "My Recipe" });

      expect(warnSpy).toHaveBeenCalledWith(
        expect.stringContaining("without a workflow payload"),
        expect.anything()
      );
      expect(appMock.loadApiJson).not.toHaveBeenCalled();
      expect(appMock.loadGraphData).not.toHaveBeenCalled();
      warnSpy.mockRestore();
    });

    it("parses a string workflow before loading", async () => {
      const workflow = { nodes: [], links: [] };

      await extension.loadWorkflowFromMessage({
        workflow: JSON.stringify(workflow),
        name: "Parsed",
      });

      expect(appMock.loadGraphData).toHaveBeenCalledWith(
        workflow,
        true,
        true,
        "Parsed",
        { openSource: "file_button" }
      );
    });

    it("warns and returns when the workflow string is not valid JSON", async () => {
      const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});

      await extension.loadWorkflowFromMessage({ workflow: "{not json" });

      expect(warnSpy).toHaveBeenCalledWith(
        expect.stringContaining("non-JSON workflow string"),
        expect.anything()
      );
      expect(appMock.loadApiJson).not.toHaveBeenCalled();
      expect(appMock.loadGraphData).not.toHaveBeenCalled();
      warnSpy.mockRestore();
    });

    it("loads API-format workflows via app.loadApiJson", async () => {
      const workflow = {
        "1": { class_type: "KSampler", inputs: {} },
        "2": { class_type: "CLIPTextEncode", inputs: {} },
      };

      await extension.loadWorkflowFromMessage({ workflow, name: "API Recipe" });

      expect(appMock.loadApiJson).toHaveBeenCalledWith(workflow, "API Recipe");
      expect(appMock.loadGraphData).not.toHaveBeenCalled();
    });

    it("loads UI-format workflows via app.loadGraphData", async () => {
      const workflow = { nodes: [{ id: 1 }], links: [] };

      await extension.loadWorkflowFromMessage({ workflow, name: "UI Recipe" });

      expect(appMock.loadGraphData).toHaveBeenCalledWith(
        workflow,
        true,
        true,
        "UI Recipe",
        { openSource: "file_button" }
      );
      expect(appMock.loadApiJson).not.toHaveBeenCalled();
    });

    it("defaults the workflow name to 'Recipe Workflow'", async () => {
      const workflow = { nodes: [], links: [] };

      await extension.loadWorkflowFromMessage({ workflow });

      expect(appMock.loadGraphData).toHaveBeenCalledWith(
        workflow,
        true,
        true,
        "Recipe Workflow",
        { openSource: "file_button" }
      );
    });

    it("logs an error when loading the workflow throws", async () => {
      const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
      const failure = new Error("load failed");
      appMock.loadGraphData.mockRejectedValue(failure);

      await extension.loadWorkflowFromMessage({
        workflow: { nodes: [], links: [] },
      });

      expect(errorSpy).toHaveBeenCalledWith(
        expect.stringContaining("failed to load workflow"),
        failure
      );
      errorSpy.mockRestore();
    });
  });

  describe("setup link-change hooks", () => {
    it("hooks root events, existing subgraphs, and future subgraphs", () => {
      const subgraph = createSubgraph({ id: "sub-1", nodes: [] });
      const graph = createGraph({ subgraphs: [subgraph] });
      appMock.graph = graph;

      extension.setup();

      expect(graph.events.addEventListener).toHaveBeenCalledWith(
        "node:slot-links:changed",
        expect.any(Function)
      );
      expect(graph.events.addEventListener).toHaveBeenCalledWith(
        "subgraph-created",
        expect.any(Function)
      );
      expect(subgraph.events.addEventListener).toHaveBeenCalledWith(
        "node:slot-links:changed",
        expect.any(Function)
      );

      const createdHandler = graph.events.addEventListener.mock.calls.find(
        ([name]) => name === "subgraph-created"
      )[1];
      const laterSubgraph = createSubgraph({ id: "sub-2", nodes: [] });
      createdHandler({ subgraph: laterSubgraph });
      expect(laterSubgraph.events.addEventListener).toHaveBeenCalledWith(
        "node:slot-links:changed",
        expect.any(Function)
      );
    });
  });
});