127 lines
4.9 KiB
TypeScript
127 lines
4.9 KiB
TypeScript
// SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
// Copyright (C) 2026 Artificial Sweetener and contributors
|
|
// SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
import { describe, expect, it, vi } from "vitest";
|
|
|
|
import { registerSamplerSocketOrder } from "../../src/samplerSocketOrder";
|
|
import { createFakeComfyApp } from "../support/testUtils";
|
|
|
|
const correctOrder = [
|
|
"model", "positive", "negative", "latent_image", "options", "segs", "region_masks"
|
|
];
|
|
|
|
/** Represent the host's required-first socket layout and saved connection slots. */
|
|
function fakeSampler() {
|
|
return {
|
|
id: 10,
|
|
constructor: { comfyClass: "SimpleSyrup.KSampler" },
|
|
inputs: [
|
|
{ name: "model", link: 1 },
|
|
{ name: "positive", link: 2 },
|
|
{ name: "latent_image", link: 3 },
|
|
{ name: "negative", link: 4 },
|
|
{ name: "options", link: null },
|
|
{ name: "segs", link: 5 },
|
|
{ name: "region_masks", link: 6 }
|
|
],
|
|
widgets: [{ name: "seed", value: 17 }, { name: "denoise", value: 0.5 }],
|
|
graph: {
|
|
links: {
|
|
1: { target_id: 10, target_slot: 0 },
|
|
2: { target_id: 10, target_slot: 1 },
|
|
3: { target_id: 10, target_slot: 2 },
|
|
4: { target_id: 10, target_slot: 3 },
|
|
5: { target_id: 10, target_slot: 5 },
|
|
6: { target_id: 10, target_slot: 6 }
|
|
},
|
|
setDirtyCanvas: vi.fn()
|
|
},
|
|
onGraphConfigured: vi.fn(() => "preserved-result")
|
|
};
|
|
}
|
|
|
|
describe("sampler socket ordering", () => {
|
|
it("orders a new sampler without changing its inputs or widget values", async () => {
|
|
const app = createFakeComfyApp();
|
|
registerSamplerSocketOrder(app);
|
|
const node = fakeSampler();
|
|
const originalInputs = [...node.inputs];
|
|
const originalArray = node.inputs;
|
|
const widgets = [...node.widgets];
|
|
await app.extensions[0]?.nodeCreated?.(node);
|
|
expect(node.inputs.map((input) => input.name)).toEqual(correctOrder);
|
|
expect(node.inputs).toBe(originalArray);
|
|
expect(node.inputs).toEqual([
|
|
originalInputs[0], originalInputs[1], originalInputs[3],
|
|
originalInputs[2], originalInputs[4], originalInputs[5], originalInputs[6]
|
|
]);
|
|
expect(node.widgets).toEqual(widgets);
|
|
expect(node.graph.links[3].target_slot).toBe(3);
|
|
expect(node.graph.links[4].target_slot).toBe(2);
|
|
expect(node.inputs[5]?.link).toBe(5);
|
|
expect(node.graph.links[5].target_slot).toBe(5);
|
|
expect(node.inputs[6]?.link).toBe(6);
|
|
expect(node.graph.links[6].target_slot).toBe(6);
|
|
});
|
|
|
|
it("repairs old workflow link indices by socket name after configuration", async () => {
|
|
const app = createFakeComfyApp();
|
|
registerSamplerSocketOrder(app);
|
|
const node = fakeSampler();
|
|
const originalCallback = node.onGraphConfigured;
|
|
await app.extensions[0]?.nodeCreated?.(node);
|
|
// Comfy restores the old saved sockets and link indices when loading a workflow.
|
|
const restored = fakeSampler().inputs;
|
|
node.inputs = restored;
|
|
node.graph.links[3].target_slot = 2;
|
|
node.graph.links[4].target_slot = 3;
|
|
expect(node.onGraphConfigured()).toBe("preserved-result");
|
|
expect(originalCallback).toHaveBeenCalledOnce();
|
|
expect(node.inputs.map((input) => input.name)).toEqual(correctOrder);
|
|
expect(node.inputs[2]?.link).toBe(4);
|
|
expect(node.inputs[3]?.link).toBe(3);
|
|
expect(node.graph.links[3].target_slot).toBe(3);
|
|
expect(node.graph.links[4].target_slot).toBe(2);
|
|
// Already ordered serialized sockets still need their graph link slots repaired.
|
|
node.graph.links[3].target_slot = 2;
|
|
node.onGraphConfigured();
|
|
expect(node.graph.links[3].target_slot).toBe(3);
|
|
});
|
|
|
|
it("preserves added widget sockets and works without links or an existing callback", async () => {
|
|
const app = createFakeComfyApp();
|
|
registerSamplerSocketOrder(app);
|
|
const node = {
|
|
constructor: { comfyClass: "SimpleSyrup.KSampler" },
|
|
inputs: [
|
|
{ name: "model", link: null }, { name: "seed", link: null },
|
|
{ name: "positive", link: null }, { name: "latent_image", link: null },
|
|
{ name: "negative", link: null }, { name: "options", link: null },
|
|
{ name: "segs", link: null },
|
|
{ name: "region_masks", link: null },
|
|
{ name: "denoise", link: null }
|
|
],
|
|
onGraphConfigured: undefined as (() => unknown) | undefined
|
|
};
|
|
await app.extensions[0]?.nodeCreated?.(node);
|
|
expect(node.inputs.map((input) => input.name)).toEqual([
|
|
...correctOrder, "seed", "denoise"
|
|
]);
|
|
node.onGraphConfigured?.();
|
|
});
|
|
|
|
it.each<{ node: unknown }>([
|
|
{ node: null }, { node: {} },
|
|
{ node: { constructor: { comfyClass: "Other.KSampler" }, inputs: [] } }
|
|
])(
|
|
"does not alter unrelated or malformed nodes: $node", async ({ node }) => {
|
|
const app = createFakeComfyApp();
|
|
registerSamplerSocketOrder(app);
|
|
const snapshot = JSON.stringify(node);
|
|
await app.extensions[0]?.nodeCreated?.(node);
|
|
expect(JSON.stringify(node)).toBe(snapshot);
|
|
}
|
|
);
|
|
});
|