132 lines
4.6 KiB
TypeScript
132 lines
4.6 KiB
TypeScript
import ComfyGraph from "$lib/ComfyGraph";
|
|
import { ComfyNumberNode, ComfyComboNode } from "$lib/nodes/widgets";
|
|
import { ComfyBoxWorkflow } from "$lib/stores/workflowState";
|
|
import { LiteGraph, Subgraph } from "@litegraph-ts/core";
|
|
import { get } from "svelte/store";
|
|
import { expect } from 'vitest';
|
|
import UnitTest from "./UnitTest";
|
|
import { Watch } from "@litegraph-ts/nodes-basic";
|
|
import type { SerializedComfyWidgetNode } from "$lib/nodes/widgets/ComfyWidgetNode";
|
|
|
|
export default class ComfyGraphTests extends UnitTest {
|
|
test__onNodeAdded__updatesLayoutState() {
|
|
const [{ graph }, layoutState] = ComfyBoxWorkflow.create()
|
|
layoutState.initDefaultLayout() // adds 3 containers
|
|
|
|
const state = get(layoutState)
|
|
expect(Object.keys(state.allItems)).toHaveLength(3)
|
|
expect(Object.keys(state.allItemsByNode)).toHaveLength(0)
|
|
|
|
const widget = LiteGraph.createNode(ComfyNumberNode);
|
|
graph.add(widget)
|
|
|
|
expect(Object.keys(state.allItems)).toHaveLength(4)
|
|
expect(Object.keys(state.allItemsByNode)).toHaveLength(1)
|
|
expect(state.allItemsByNode[widget.id]).toBeTruthy();
|
|
|
|
graph.add(widget)
|
|
}
|
|
|
|
test__correctSubgraphFactory() {
|
|
const graph = new ComfyGraph();
|
|
const subgraph = LiteGraph.createNode(Subgraph);
|
|
graph.add(subgraph)
|
|
expect(subgraph.graph).toBeInstanceOf(ComfyGraph)
|
|
}
|
|
|
|
test__onNodeAdded__handlesNodesAddedInSubgraphs() {
|
|
const [{ graph }, layoutState] = ComfyBoxWorkflow.create()
|
|
layoutState.initDefaultLayout()
|
|
|
|
const subgraph = LiteGraph.createNode(Subgraph);
|
|
graph.add(subgraph)
|
|
|
|
const state = get(layoutState)
|
|
expect(Object.keys(state.allItems)).toHaveLength(3)
|
|
expect(Object.keys(state.allItemsByNode)).toHaveLength(0)
|
|
|
|
const widget = LiteGraph.createNode(ComfyNumberNode);
|
|
subgraph.subgraph.add(widget)
|
|
|
|
expect(Object.keys(state.allItems)).toHaveLength(4)
|
|
expect(Object.keys(state.allItemsByNode)).toHaveLength(1)
|
|
expect(state.allItemsByNode[widget.id]).toBeTruthy();
|
|
}
|
|
|
|
test__onNodeAdded__handlesSubgraphsWithNodes() {
|
|
const [{ graph }, layoutState] = ComfyBoxWorkflow.create()
|
|
layoutState.initDefaultLayout()
|
|
|
|
const state = get(layoutState)
|
|
expect(Object.keys(state.allItems)).toHaveLength(3)
|
|
expect(Object.keys(state.allItemsByNode)).toHaveLength(0)
|
|
|
|
const subgraph = LiteGraph.createNode(Subgraph);
|
|
const widget = LiteGraph.createNode(ComfyNumberNode);
|
|
subgraph.subgraph.add(widget)
|
|
graph.add(subgraph)
|
|
|
|
expect(Object.keys(state.allItems)).toHaveLength(4)
|
|
expect(Object.keys(state.allItemsByNode)).toHaveLength(1)
|
|
expect(state.allItemsByNode[widget.id]).toBeTruthy();
|
|
}
|
|
|
|
test__onNodeRemoved__updatesLayoutState() {
|
|
const [{ graph }, layoutState] = ComfyBoxWorkflow.create()
|
|
layoutState.initDefaultLayout()
|
|
|
|
const widget = LiteGraph.createNode(ComfyNumberNode);
|
|
graph.add(widget)
|
|
|
|
const state = get(layoutState)
|
|
expect(Object.keys(state.allItems)).toHaveLength(4)
|
|
expect(Object.keys(state.allItemsByNode)).toHaveLength(1)
|
|
expect(state.allItemsByNode[widget.id]).toBeTruthy();
|
|
|
|
graph.remove(widget)
|
|
|
|
expect(Object.keys(state.allItems)).toHaveLength(3)
|
|
expect(Object.keys(state.allItemsByNode)).toHaveLength(0)
|
|
}
|
|
|
|
test__serialize__stripsLinkData() {
|
|
const [{ graph }, layoutState] = ComfyBoxWorkflow.create()
|
|
layoutState.initDefaultLayout()
|
|
|
|
const widget = LiteGraph.createNode(ComfyNumberNode);
|
|
const watch = LiteGraph.createNode(Watch);
|
|
graph.add(widget)
|
|
graph.add(watch)
|
|
|
|
widget.connect(0, watch, 0)
|
|
const link = widget.getOutputLinks(0)[0]
|
|
widget.setOutputData(0, 42);
|
|
|
|
const result = graph.serialize();
|
|
|
|
const serNode = result.nodes.find(n => n.id === widget.id);
|
|
|
|
expect(serNode.outputs[0]._data).toBeUndefined()
|
|
}
|
|
|
|
test__serialize__savesComboData() {
|
|
const [{ graph }, layoutState] = ComfyBoxWorkflow.create()
|
|
layoutState.initDefaultLayout()
|
|
|
|
const widget = LiteGraph.createNode(ComfyComboNode);
|
|
const watch = LiteGraph.createNode(Watch);
|
|
graph.add(widget)
|
|
graph.add(watch)
|
|
|
|
widget.connect(0, watch, 0)
|
|
widget.properties.values = ["A", "B", "C"]
|
|
widget.setValue("B");
|
|
|
|
const result = graph.serialize();
|
|
|
|
const serNode = result.nodes.find(n => n.id === widget.id) as SerializedComfyWidgetNode;
|
|
|
|
expect(serNode.comfyValue).toBe("B")
|
|
}
|
|
}
|