MajoorWaldi-ComfyUI-Majoor-AssetsManager/ui/tests/node_stream_controller.vitest.ts

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

beforeEach(() => {
    vi.resetModules();
    vi.useFakeTimers();
    globalThis.window = ({
        location: ({ href: "http://localhost/" } as any),
    } as any);
});

afterEach(() => {
    vi.useRealTimers();
    delete globalThis.HTMLCanvasElement;
});

describe("NodeStreamController", () => {
    it("streams the selected node canvas preview", async () => {
        const onOutput = vi.fn();
        const node = {
            id: 42,
            comfyClass: "PreviewImage",
            imgs: [{ src: "/view?filename=should_not_emit.png&subfolder=&type=output" }],
        };
        const app = {
            canvas: { selected_nodes: { 42: node } },
            graph: {
                getNodeById(id) {
                    return Number(id) === 42 ? node : null;
                },
            },
        };

        const mod = await import("../features/viewer/nodeStream/NodeStreamController.js");
        mod.initNodeStream({ app, onOutput });
        mod.setNodeStreamActive(true);
        await vi.advanceTimersByTimeAsync(250);

        expect(mod.getNodeStreamActive()).toBe(true);
        expect(mod.getSelectedNodeId()).toBe("42");
        expect(mod.getPinnedNodeId()).toBeNull();
        expect(onOutput).toHaveBeenCalledWith(
            expect.objectContaining({
                filename: "should_not_emit.png",
                type: "output",
                kind: "image",
                _nodeId: "42",
                _source: "canvas",
            }),
        );

        mod.teardownNodeStream(app);
    });

    it("ignores execution-driven output updates in selection-only mode", async () => {
        const onOutput = vi.fn();
        const app = {
            canvas: { selected_nodes: {} },
            graph: {
                getNodeById() {
                    return null;
                },
            },
        };

        const mod = await import("../features/viewer/nodeStream/NodeStreamController.js");
        mod.initNodeStream({ app, onOutput });
        mod.onNodeOutputs(
            new Map([["99", { images: [{ filename: "still_ignored.png", type: "output" }] }]]),
            { app },
        );
        mod.setNodeStreamActive(true);
        await vi.advanceTimersByTimeAsync(100);

        expect(mod.getNodeStreamActive()).toBe(true);
        expect(onOutput).not.toHaveBeenCalled();

        mod.teardownNodeStream(app);
    });

    it("resolves downstream previews when the graph is exposed through canvas.graph", async () => {
        const onOutput = vi.fn();
        const selectedNode = {
            id: 1,
            comfyClass: "CustomProcessor",
            outputs: [{ links: [77] }],
        };
        const previewNode = {
            id: 2,
            comfyClass: "PreviewImage",
            inputs: [{ link: 77 }],
            imgs: [{ src: "/view?filename=downstream.png&subfolder=previews&type=output" }],
        };
        const graph = {
            nodes: [selectedNode, previewNode],
            links: { 77: [77, 1, 0, 2, 0, "IMAGE"] },
            getNodeById(id) {
                return [selectedNode, previewNode].find((node) => Number(node.id) === Number(id)) || null;
            },
        };
        const app = {
            canvas: {
                graph,
                selected_nodes: { 1: selectedNode },
            },
        };

        const mod = await import("../features/viewer/nodeStream/NodeStreamController.js");
        mod.initNodeStream({ app, onOutput });
        mod.setNodeStreamActive(true);
        await vi.advanceTimersByTimeAsync(250);

        expect(onOutput).toHaveBeenCalledWith(
            expect.objectContaining({
                filename: "downstream.png",
                subfolder: "previews",
                type: "output",
                kind: "image",
                _nodeId: "1",
                _previewNodeId: "2",
                _source: "graph-downstream",
            }),
        );

        mod.teardownNodeStream(app);
    });

    it("streams the ImageOps preview canvas as a cached data URL and re-emits only on signature change", async () => {
        class MockCanvas {}
        globalThis.HTMLCanvasElement = (MockCanvas as any);

        const canvas = new MockCanvas();
        (canvas as any).width = 320;
        (canvas as any).height = 180;
        (canvas as any).toDataURL = vi.fn(() => "data:image/png;base64,preview");

        const onOutput = vi.fn();
        const node = {
            id: 7,
            comfyClass: "ImageOpsPreview",
            __imageops_state: {
                canvas,
                lastKey: "render-a",
                lastRenderTick: 1,
                nativeDirty: false,
            },
        };
        const app = {
            canvas: { selected_nodes: { 7: node } },
            graph: {
                getNodeById(id) {
                    return Number(id) === 7 ? node : null;
                },
            },
        };

        const mod = await import("../features/viewer/nodeStream/NodeStreamController.js");
        mod.initNodeStream({ app, onOutput });
        mod.setNodeStreamActive(true);
        await vi.advanceTimersByTimeAsync(10);

        expect(onOutput).toHaveBeenCalledWith(
            expect.objectContaining({
                filename: expect.stringMatching(/^imageops_7_[a-f0-9]+\.png$/),
                url: "data:image/png;base64,preview",
                kind: "image",
                width: 320,
                height: 180,
                _nodeId: "7",
                _classType: "ImageOpsPreview",
                _source: "imageops-live-preview",
            }),
        );
        expect((canvas as any).toDataURL).toHaveBeenCalledTimes(1);

        // Same signature → no re-emit and no re-encode.
        onOutput.mockClear();
        await vi.advanceTimersByTimeAsync(250);
        expect(onOutput).not.toHaveBeenCalled();
        expect((canvas as any).toDataURL).toHaveBeenCalledTimes(1);

        // Signature changes → emit + encode again.
        node.__imageops_state.lastKey = "render-b";
        await vi.advanceTimersByTimeAsync(250);
        expect(onOutput).toHaveBeenCalledTimes(1);
        expect((canvas as any).toDataURL).toHaveBeenCalledTimes(2);

        mod.teardownNodeStream(app);
    });

    it("normalizes unsupported watch mode values back to selected", async () => {
        const app = {
            canvas: { selected_nodes: {} },
            graph: {
                getNodeById() {
                    return null;
                },
            },
        };

        const mod = await import("../features/viewer/nodeStream/NodeStreamController.js");
        mod.initNodeStream({ app, onOutput: vi.fn() });

        mod.setWatchMode("all");

        expect(mod.getWatchMode()).toBe("selected");

        mod.teardownNodeStream(app);
    });

    it("emits an empty status when the selection is cleared", async () => {
        const onStatus = vi.fn();
        const node = {
            id: 42,
            comfyClass: "PreviewImage",
            imgs: [{ src: "/view?filename=selected.png&subfolder=&type=output" }],
        };
        const app = {
            canvas: { selected_nodes: { 42: node } },
            graph: {
                getNodeById(id) {
                    return Number(id) === 42 ? node : null;
                },
            },
        };

        const mod = await import("../features/viewer/nodeStream/NodeStreamController.js");
        mod.initNodeStream({ app, onOutput: vi.fn(), onStatus });
        mod.setNodeStreamActive(true);
        await vi.advanceTimersByTimeAsync(200);

        app.canvas.selected_nodes = ({} as any);
        await vi.advanceTimersByTimeAsync(200);

        expect(onStatus).toHaveBeenLastCalledWith("", "");

        mod.teardownNodeStream(app);
    });
});