diff --git a/src/components/__tests__/GenerateVideoNode.test.tsx b/src/components/__tests__/GenerateVideoNode.test.tsx
index 37fb033b..494e2b23 100644
--- a/src/components/__tests__/GenerateVideoNode.test.tsx
+++ b/src/components/__tests__/GenerateVideoNode.test.tsx
@@ -177,10 +177,10 @@ describe("GenerateVideoNode", () => {
         </TestWrapper>
       );
 
-      // Should have Image and Prompt labels for inputs, Video for output
+      // Should have Image and Prompt labels for inputs, plus Video labels for input and output.
       expect(screen.getByText("Image")).toBeInTheDocument();
       expect(screen.getByText("Prompt")).toBeInTheDocument();
-      expect(screen.getByText("Video")).toBeInTheDocument();
+      expect(screen.getAllByText("Video")).toHaveLength(2);
     });
   });
 
diff --git a/src/components/__tests__/WorkflowCanvas.test.tsx b/src/components/__tests__/WorkflowCanvas.test.tsx
index 9e889a6b..55957061 100644
--- a/src/components/__tests__/WorkflowCanvas.test.tsx
+++ b/src/components/__tests__/WorkflowCanvas.test.tsx
@@ -39,6 +39,7 @@ const mockGetViewport = vi.fn(() => ({ x: 0, y: 0, zoom: 1 }));
 const mockZoomIn = vi.fn();
 const mockZoomOut = vi.fn();
 const mockSetViewport = vi.fn();
+const mockSetNodes = vi.fn();
 
 vi.mock("@xyflow/react", async () => {
   const actual = await vi.importActual("@xyflow/react");
@@ -50,6 +51,7 @@ vi.mock("@xyflow/react", async () => {
       zoomIn: mockZoomIn,
       zoomOut: mockZoomOut,
       setViewport: mockSetViewport,
+      setNodes: mockSetNodes,
     }),
   };
 });