Skip to content

Commit 5f7ba1f

Browse files
committed
Merge remote-tracking branch 'origin/main' into benjcooley/custom-node-e2e-review-fixes
2 parents 397034b + 32d0b6e commit 5f7ba1f

1 file changed

Lines changed: 292 additions & 0 deletions

File tree

Lines changed: 292 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,292 @@
1+
import { fromAny } from '@total-typescript/shoehorn'
2+
import { afterAll, beforeAll, describe, expect, test } from 'vitest'
3+
4+
import { LGraph, LGraphNode, LiteGraph } from '@/lib/litegraph/src/litegraph'
5+
import type { ISerialisedGraph } from '@/lib/litegraph/src/litegraph'
6+
7+
import floatingLink from './__fixtures__/assets/floatingLink.json'
8+
import linkedNodes from './__fixtures__/assets/linkedNodes.json'
9+
import reroutesComplex from './__fixtures__/assets/reroutesComplex.json'
10+
11+
/**
12+
* Loading a workflow and saving it again must not lose entities.
13+
*
14+
* The existing round-trip tests compare a second serialisation to the first,
15+
* which proves the output is a fixed point but says nothing about whether the
16+
* *input* survived. These assert the property users actually depend on: open a
17+
* workflow, save it, and everything you had is still there.
18+
*
19+
* Serialisation deliberately normalises — schema version is rewritten and
20+
* conflicting ids are reassigned — so this compares entity sets rather than
21+
* bytes. It compares them whole: a count survives an entity being renumbered,
22+
* repointed at a different slot, or replaced outright.
23+
*
24+
* The fixture node types must be registered. Without them `createNode` returns
25+
* null, every node takes the error branch, and `serialize()` echoes the input
26+
* object straight back — so the output is derived from the input by
27+
* construction and no assertion about nodes can fail. `roundTrip` asserts no
28+
* node carries `has_errors` so that can never silently return.
29+
*/
30+
31+
const FIXTURE_NODE_TYPES = {
32+
VAEDecode: {
33+
inputs: [
34+
['samples', 'LATENT'],
35+
['vae', 'VAE']
36+
],
37+
outputs: [['IMAGE', 'IMAGE']],
38+
widgets: []
39+
},
40+
SaveImage: {
41+
inputs: [['images', 'IMAGE']],
42+
outputs: [],
43+
widgets: [['filename_prefix', 'ComfyUI']]
44+
},
45+
InvertMask: {
46+
inputs: [['mask', 'MASK']],
47+
outputs: [['MASK', 'MASK']],
48+
widgets: []
49+
}
50+
} as const satisfies Record<
51+
string,
52+
{
53+
inputs: readonly (readonly [string, string])[]
54+
outputs: readonly (readonly [string, string])[]
55+
widgets: readonly (readonly [string, string])[]
56+
}
57+
>
58+
59+
const originalNodeTypes = Object.fromEntries(
60+
Object.keys(FIXTURE_NODE_TYPES).map((type) => [
61+
type,
62+
LiteGraph.registered_node_types[type]
63+
])
64+
)
65+
const originalFixtureNode = LiteGraph.Nodes.FixtureNode
66+
67+
beforeAll(() => {
68+
for (const [type, shape] of Object.entries(FIXTURE_NODE_TYPES)) {
69+
class FixtureNode extends LGraphNode {
70+
constructor(title?: string) {
71+
super(title ?? type)
72+
this.serialize_widgets = true
73+
for (const [name, slotType] of shape.inputs)
74+
this.addInput(name, slotType)
75+
for (const [name, slotType] of shape.outputs)
76+
this.addOutput(name, slotType)
77+
for (const [name, value] of shape.widgets)
78+
this.addWidget('text', name, value, () => {})
79+
}
80+
}
81+
LiteGraph.registerNodeType(type, FixtureNode)
82+
}
83+
})
84+
85+
afterAll(() => {
86+
for (const type of Object.keys(FIXTURE_NODE_TYPES)) {
87+
const originalNodeType = originalNodeTypes[type]
88+
if (originalNodeType)
89+
LiteGraph.registered_node_types[type] = originalNodeType
90+
else delete LiteGraph.registered_node_types[type]
91+
}
92+
93+
if (originalFixtureNode) LiteGraph.Nodes.FixtureNode = originalFixtureNode
94+
else delete LiteGraph.Nodes.FixtureNode
95+
})
96+
97+
interface RoundTripFixture {
98+
name: string
99+
graph: ISerialisedGraph
100+
}
101+
102+
const fixtures: RoundTripFixture[] = [
103+
{
104+
name: 'linked nodes',
105+
graph: fromAny<ISerialisedGraph, unknown>(linkedNodes)
106+
},
107+
{
108+
name: 'floating link',
109+
graph: fromAny<ISerialisedGraph, unknown>(floatingLink)
110+
},
111+
{
112+
name: 'complex reroutes',
113+
graph: fromAny<ISerialisedGraph, unknown>(reroutesComplex)
114+
}
115+
]
116+
117+
function roundTrip(source: ISerialisedGraph) {
118+
const loaded = new LGraph(structuredClone(source))
119+
expect(loaded.nodes.filter((n) => n.has_errors)).toEqual([])
120+
return loaded.serialize()
121+
}
122+
123+
/**
124+
* Nodes compared whole. By id alone, a regression that drops every input,
125+
* output, widget value or title still passes.
126+
*/
127+
function nodeKeys(graph: Pick<ISerialisedGraph, 'nodes'>) {
128+
return (graph.nodes ?? [])
129+
.map((node) =>
130+
JSON.stringify({
131+
id: node.id,
132+
type: node.type,
133+
inputs: (node.inputs ?? []).map((i) => [
134+
i.name,
135+
i.type,
136+
i.link ?? null
137+
]),
138+
outputs: (node.outputs ?? []).map((o) => [
139+
o.name,
140+
o.type,
141+
[...(o.links ?? [])].sort(ascending)
142+
]),
143+
widgets_values: node.widgets_values?.length ? node.widgets_values : null
144+
})
145+
)
146+
.sort()
147+
}
148+
149+
function ascending(a: number | string, b: number | string) {
150+
return String(a).localeCompare(String(b), undefined, { numeric: true })
151+
}
152+
153+
/**
154+
* Compared whole. Ids alone survive a regression that flattens every
155+
* `parentId`, empties every `linkIds`, or resets every `pos` — which is most of
156+
* what the reroute fixture is for.
157+
*/
158+
function rerouteKeys(graph: Pick<ISerialisedGraph, 'extra'>) {
159+
return (graph.extra?.reroutes ?? [])
160+
.map((reroute) =>
161+
JSON.stringify({
162+
id: reroute.id,
163+
parentId: reroute.parentId ?? null,
164+
pos: reroute.pos ?? null,
165+
linkIds: [...(reroute.linkIds ?? [])].sort(ascending),
166+
floating: reroute.floating ?? null
167+
})
168+
)
169+
.sort()
170+
}
171+
172+
/**
173+
* Reroute-to-link association is not on the link in schema 0.4 — `serialize()`
174+
* rebuilds it into `extra.linkExtensions`, so it needs its own assertion.
175+
*/
176+
function linkExtensionKeys(graph: Pick<ISerialisedGraph, 'extra'>) {
177+
return (graph.extra?.linkExtensions ?? [])
178+
.map((ext) =>
179+
JSON.stringify({ id: ext.id, parentId: ext.parentId ?? null })
180+
)
181+
.sort()
182+
}
183+
184+
/**
185+
* A comparison against an empty collection passes whether or not the code
186+
* works. Every assertion below runs through this so a fixture that later loses
187+
* its reroutes degrades into a failure rather than a silent no-op.
188+
*/
189+
function expectPreserved(before: string[], after: string[]) {
190+
expect(before.length).toBeGreaterThan(0)
191+
expect(after).toEqual(before)
192+
}
193+
194+
/**
195+
* Links and groups are compared whole, not counted. A count survives a link
196+
* being renumbered, repointed at a different slot, or replaced outright.
197+
*/
198+
function linkKeys(graph: Pick<ISerialisedGraph, 'links'>) {
199+
return (graph.links ?? []).map((link) => JSON.stringify(link)).sort()
200+
}
201+
202+
function floatingLinkKeys(graph: Pick<ISerialisedGraph, 'floatingLinks'>) {
203+
return (graph.floatingLinks ?? []).map((link) => JSON.stringify(link)).sort()
204+
}
205+
206+
function groupKeys(graph: Pick<ISerialisedGraph, 'groups'>) {
207+
return (graph.groups ?? [])
208+
.map(({ id, title, bounding }) => JSON.stringify({ id, title, bounding }))
209+
.sort()
210+
}
211+
212+
/**
213+
* Every fixture ships with `groups: []`, so a group assertion against them
214+
* unmodified compares nothing to nothing.
215+
*/
216+
function withGroups(graph: ISerialisedGraph): ISerialisedGraph {
217+
return {
218+
...structuredClone(graph),
219+
groups: [
220+
{ id: 1, title: 'first', bounding: [0, 0, 140, 90] },
221+
{ id: 2, title: 'second', bounding: [200, 40, 180, 120] }
222+
]
223+
}
224+
}
225+
226+
describe('LGraph round trip preserves the input', () => {
227+
for (const { name, graph } of fixtures) {
228+
describe(name, () => {
229+
test('keeps every node, whole', () => {
230+
const before = nodeKeys(graph)
231+
const after = nodeKeys(roundTrip(graph))
232+
233+
expectPreserved(before, after)
234+
})
235+
236+
test.skipIf(linkKeys(graph).length === 0)(
237+
'keeps every link, with its endpoints',
238+
() => {
239+
expectPreserved(linkKeys(graph), linkKeys(roundTrip(graph)))
240+
}
241+
)
242+
243+
test.skipIf(floatingLinkKeys(graph).length === 0)(
244+
'keeps every floating link, with its endpoints',
245+
() => {
246+
expectPreserved(
247+
floatingLinkKeys(graph),
248+
floatingLinkKeys(roundTrip(graph))
249+
)
250+
}
251+
)
252+
253+
test.skipIf(rerouteKeys(graph).length === 0)(
254+
'keeps every reroute, with its parent, position and links',
255+
() => {
256+
expectPreserved(rerouteKeys(graph), rerouteKeys(roundTrip(graph)))
257+
}
258+
)
259+
260+
test.skipIf(
261+
linkExtensionKeys(graph).length === 0 || rerouteKeys(graph).length === 0
262+
)('keeps the reroute-to-link association in extra.linkExtensions', () => {
263+
expectPreserved(
264+
linkExtensionKeys(graph),
265+
linkExtensionKeys(roundTrip(graph))
266+
)
267+
})
268+
269+
test('keeps every group, by identity and bounds', () => {
270+
const grouped = withGroups(graph)
271+
272+
expectPreserved(groupKeys(grouped), groupKeys(roundTrip(grouped)))
273+
})
274+
275+
test('does not mutate the workflow it was given', () => {
276+
const untouched = structuredClone(graph)
277+
const subject = structuredClone(graph)
278+
279+
new LGraph(subject).serialize()
280+
281+
expect(subject).toEqual(untouched)
282+
})
283+
284+
test('is stable when saved twice', () => {
285+
const once = roundTrip(graph)
286+
const twice = new LGraph(structuredClone(once)).serialize()
287+
288+
expect(twice).toEqual(once)
289+
})
290+
})
291+
}
292+
})

0 commit comments

Comments
 (0)