diff --git a/src/components/__tests__/GenerationComposer.test.tsx b/src/components/__tests__/GenerationComposer.test.tsx
index 474ebcf8..e902e9fc 100644
--- a/src/components/__tests__/GenerationComposer.test.tsx
+++ b/src/components/__tests__/GenerationComposer.test.tsx
@@ -120,7 +120,7 @@ function imageNode(id: string, selected = false, data: Partial<NanoBananaNodeDat
   } as WorkflowNode;
 }
 
-function pricedImageNode(id: string, selected = false): WorkflowNode {
+function pricedImageNode(id: string, selected = false, data: Partial<NanoBananaNodeData> = {}): WorkflowNode {
   return imageNode(id, selected, {
     selectedModel: {
       provider: "newapiwg",
@@ -129,6 +129,7 @@ function pricedImageNode(id: string, selected = false): WorkflowNode {
       capabilities: ["text-to-image", "image-to-image"],
       pricing: { type: "per-run", amount: 2 },
     },
+    ...data,
   });
 }