/** * Integration tests for undo/redo through the Zustand store. * * Tests that undo/redo correctly captures and restores workflow state * for all undoable actions (add/remove nodes, connect edges, etc.). */ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import { act } from "@testing-library/react"; import { useWorkflowStore } from "../workflowStore"; // Mock the Toast hook vi.mock("@/components/Toast", () => ({ useToast: { getState: () => ({ show: vi.fn(), }), }, })); // Mock the logger vi.mock("@/utils/logger", () => ({ logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), startSession: vi.fn().mockResolvedValue(undefined), endSession: vi.fn().mockResolvedValue(undefined), getCurrentSession: vi.fn().mockReturnValue(null), }, })); // Mock localStorage const mockLocalStorage: Record = {}; vi.stubGlobal("localStorage", { getItem: vi.fn((key: string) => mockLocalStorage[key] || null), setItem: vi.fn((key: string, value: string) => { mockLocalStorage[key] = value; }), removeItem: vi.fn((key: string) => { delete mockLocalStorage[key]; }), clear: vi.fn(() => { Object.keys(mockLocalStorage).forEach((key) => delete mockLocalStorage[key]); }), }); function resetStore() { const store = useWorkflowStore.getState(); store.clearWorkflow(); } describe("Undo/Redo integration", () => { beforeEach(() => { vi.useFakeTimers(); resetStore(); }); afterEach(() => { vi.useRealTimers(); }); it("starts with canUndo and canRedo both false", () => { const { canUndo, canRedo } = useWorkflowStore.getState(); expect(canUndo).toBe(false); expect(canRedo).toBe(false); }); describe("addNode + undo", () => { it("undoes adding a node", () => { let store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(0); act(() => { store.addNode("prompt", { x: 100, y: 100 }); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(1); expect(store.canUndo).toBe(true); act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(0); expect(store.canUndo).toBe(false); expect(store.canRedo).toBe(true); }); it("redoes after undoing addNode", () => { let store = useWorkflowStore.getState(); act(() => { store.addNode("prompt", { x: 100, y: 100 }); }); store = useWorkflowStore.getState(); const nodeId = store.nodes[0].id; act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(0); act(() => { store.redo(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(1); expect(store.nodes[0].id).toBe(nodeId); expect(store.canRedo).toBe(false); }); }); describe("removeNode + undo", () => { it("undoes removing a node", () => { let store = useWorkflowStore.getState(); act(() => { store.addNode("prompt", { x: 0, y: 0 }); }); store = useWorkflowStore.getState(); const nodeId = store.nodes[0].id; act(() => { store.removeNode(nodeId); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(0); act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(1); expect(store.nodes[0].id).toBe(nodeId); }); }); describe("onConnect + undo", () => { it("undoes connecting an edge", () => { let store = useWorkflowStore.getState(); act(() => { store.addNode("prompt", { x: 0, y: 0 }); store.addNode("nanoBanana", { x: 300, y: 0 }); }); store = useWorkflowStore.getState(); const promptId = store.nodes[0].id; const genId = store.nodes[1].id; const edgesBefore = store.edges.length; act(() => { store.onConnect({ source: promptId, target: genId, sourceHandle: "text", targetHandle: "text", }); }); store = useWorkflowStore.getState(); expect(store.edges.length).toBe(edgesBefore + 1); act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.edges.length).toBe(edgesBefore); }); }); describe("removeEdge + undo", () => { it("undoes removing an edge", () => { let store = useWorkflowStore.getState(); act(() => { store.addNode("prompt", { x: 0, y: 0 }); store.addNode("nanoBanana", { x: 300, y: 0 }); }); store = useWorkflowStore.getState(); const promptId = store.nodes[0].id; const genId = store.nodes[1].id; act(() => { store.onConnect({ source: promptId, target: genId, sourceHandle: "text", targetHandle: "text", }); }); store = useWorkflowStore.getState(); const edgeId = store.edges[store.edges.length - 1].id; act(() => { store.removeEdge(edgeId); }); store = useWorkflowStore.getState(); expect(store.edges.find((e) => e.id === edgeId)).toBeUndefined(); act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.edges.find((e) => e.id === edgeId)).toBeDefined(); }); }); describe("delete node via onNodesChange + onEdgesChange restores connections", () => { it("single undo restores both node and its connected edges", async () => { let store = useWorkflowStore.getState(); // Create two nodes and connect them act(() => { store.addNode("prompt", { x: 0, y: 0 }); }); store = useWorkflowStore.getState(); act(() => { store.addNode("nanoBanana", { x: 300, y: 0 }); }); store = useWorkflowStore.getState(); const promptId = store.nodes[0].id; const genId = store.nodes[1].id; act(() => { store.onConnect({ source: promptId, target: genId, sourceHandle: "text", targetHandle: "text", }); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(2); expect(store.edges.length).toBe(1); // Simulate pressing Delete: React Flow fires onNodesChange(remove) then // onEdgesChange(remove) synchronously in the same cycle act(() => { store.onNodesChange([{ type: "remove", id: promptId }]); store.onEdgesChange([{ type: "remove", id: store.edges[0].id }]); }); // Wait for microtask to clear nodeRemoveCheckpointActive await act(async () => { await Promise.resolve(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(1); expect(store.edges.length).toBe(0); // Single undo should restore both the node AND the edge act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(2); expect(store.edges.length).toBe(1); expect(store.nodes.find((n) => n.id === promptId)).toBeDefined(); // The delete should have been exactly one undo entry — undoing again // should go back to before the edge was connected, not to some // intermediate state caused by clearStaleInputImages side effects act(() => { store.undo(); }); store = useWorkflowStore.getState(); // Previous checkpoint was onConnect, so edges should be gone expect(store.edges.length).toBe(0); expect(store.nodes.length).toBe(2); }); it("single undo restores image-source node and does not create extra entries from clearStaleInputImages", async () => { let store = useWorkflowStore.getState(); // Create imageInput -> nanoBanana (image connection triggers clearStaleInputImages on delete) act(() => { store.addNode("imageInput", { x: 0, y: 0 }); }); store = useWorkflowStore.getState(); act(() => { store.addNode("nanoBanana", { x: 300, y: 0 }); }); store = useWorkflowStore.getState(); const imageInputId = store.nodes[0].id; const genId = store.nodes[1].id; act(() => { store.onConnect({ source: imageInputId, target: genId, sourceHandle: "image", targetHandle: "image", }); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(2); expect(store.edges.length).toBe(1); // Delete the imageInput node — this triggers clearStaleInputImages // which calls updateNodeData on the nanoBanana target act(() => { store.onNodesChange([{ type: "remove", id: imageInputId }]); store.onEdgesChange([{ type: "remove", id: store.edges[0].id }]); }); // Wait for microtask to clear nodeRemoveCheckpointActive await act(async () => { await Promise.resolve(); }); // Also advance past the debounced data-change timer (500ms) await act(async () => { vi.advanceTimersByTime(600); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(1); expect(store.edges.length).toBe(0); // Single undo should restore the node and edge act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(2); expect(store.edges.length).toBe(1); expect(store.nodes.find((n) => n.id === imageInputId)).toBeDefined(); // Verify no extra undo entries from clearStaleInputImages: // Next undo should go back to before onConnect (2 nodes, 0 edges) act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.edges.length).toBe(0); expect(store.nodes.length).toBe(2); }); }); describe("new action clears redo stack", () => { it("clears redo stack when a new undoable action is performed", () => { let store = useWorkflowStore.getState(); act(() => { store.addNode("prompt", { x: 0, y: 0 }); }); act(() => { store = useWorkflowStore.getState(); store.undo(); }); store = useWorkflowStore.getState(); expect(store.canRedo).toBe(true); // Perform a new action — should clear redo stack act(() => { store.addNode("imageInput", { x: 100, y: 100 }); }); store = useWorkflowStore.getState(); expect(store.canRedo).toBe(false); }); }); describe("multiple undo/redo", () => { it("supports multiple undo/redo in sequence", () => { let store = useWorkflowStore.getState(); act(() => { store.addNode("prompt", { x: 0, y: 0 }); }); act(() => { store = useWorkflowStore.getState(); store.addNode("imageInput", { x: 100, y: 0 }); }); act(() => { store = useWorkflowStore.getState(); store.addNode("nanoBanana", { x: 200, y: 0 }); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(3); // Undo 3 times act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(2); act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(1); act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(0); expect(store.canUndo).toBe(false); // Redo 3 times act(() => { store.redo(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(1); act(() => { store.redo(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(2); act(() => { store.redo(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(3); expect(store.canRedo).toBe(false); }); }); describe("undo on empty stack is no-op", () => { it("does nothing when undo stack is empty", () => { let store = useWorkflowStore.getState(); const before = store.nodes.length; act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(before); expect(store.canUndo).toBe(false); }); }); describe("redo on empty stack is no-op", () => { it("does nothing when redo stack is empty", () => { let store = useWorkflowStore.getState(); act(() => { store.addNode("prompt", { x: 0, y: 0 }); }); store = useWorkflowStore.getState(); const nodeCount = store.nodes.length; act(() => { store.redo(); }); store = useWorkflowStore.getState(); expect(store.nodes.length).toBe(nodeCount); }); }); describe("clearWorkflow resets undo history", () => { it("clears undo and redo stacks on clearWorkflow", () => { let store = useWorkflowStore.getState(); act(() => { store.addNode("prompt", { x: 0, y: 0 }); store.addNode("imageInput", { x: 100, y: 0 }); }); store = useWorkflowStore.getState(); expect(store.canUndo).toBe(true); act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.canRedo).toBe(true); act(() => { store.clearWorkflow(); }); store = useWorkflowStore.getState(); expect(store.canUndo).toBe(false); expect(store.canRedo).toBe(false); }); }); describe("setEdgeStyle + undo", () => { it("undoes edge style change", () => { let store = useWorkflowStore.getState(); const originalStyle = store.edgeStyle; act(() => { store.setEdgeStyle("angular"); }); store = useWorkflowStore.getState(); expect(store.edgeStyle).toBe("angular"); act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.edgeStyle).toBe(originalStyle); }); }); describe("group operations + undo", () => { it("undoes creating a group", () => { let store = useWorkflowStore.getState(); act(() => { store.addNode("prompt", { x: 0, y: 0 }); store.addNode("imageInput", { x: 100, y: 0 }); }); store = useWorkflowStore.getState(); const nodeIds = store.nodes.map((n) => n.id); act(() => { store.createGroup(nodeIds); }); store = useWorkflowStore.getState(); expect(Object.keys(store.groups).length).toBe(1); act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(Object.keys(store.groups).length).toBe(0); }); it("undoes deleting a group", () => { let store = useWorkflowStore.getState(); act(() => { store.addNode("prompt", { x: 0, y: 0 }); }); store = useWorkflowStore.getState(); const nodeIds = store.nodes.map((n) => n.id); let groupId: string; act(() => { groupId = store.createGroup(nodeIds); }); store = useWorkflowStore.getState(); expect(Object.keys(store.groups).length).toBe(1); act(() => { store.deleteGroup(groupId!); }); store = useWorkflowStore.getState(); expect(Object.keys(store.groups).length).toBe(0); act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(Object.keys(store.groups).length).toBe(1); }); }); describe("hasUnsavedChanges is set on undo/redo", () => { it("marks as unsaved after undo", () => { let store = useWorkflowStore.getState(); act(() => { store.addNode("prompt", { x: 0, y: 0 }); }); // Simulate that the workflow was "saved" act(() => { useWorkflowStore.setState({ hasUnsavedChanges: false }); }); store = useWorkflowStore.getState(); expect(store.hasUnsavedChanges).toBe(false); act(() => { store.undo(); }); store = useWorkflowStore.getState(); expect(store.hasUnsavedChanges).toBe(true); }); }); });