diff --git a/web/common/src/components/Lineage/LineageLayout.tsx b/web/common/src/components/Lineage/LineageLayout.tsx
index e01e8ae9e9..2ab4a34879 100644
--- a/web/common/src/components/Lineage/LineageLayout.tsx
+++ b/web/common/src/components/Lineage/LineageLayout.tsx
@@ -66,13 +66,15 @@ export function LineageLayout<
 }) {
   return (
     <ReactFlowProvider>
-      <LineageLayoutContainer isBuildingLayout={isBuildingLayout}>
+      <LineageLayoutContainer
+        isBuildingLayout={isBuildingLayout}
+        className={className}
+      >
         <LineageLayoutBase
           nodeTypes={nodeTypes}
           edgeTypes={edgeTypes}
           nodesDraggable={nodesDraggable}
           nodesConnectable={nodesConnectable}
-          className={className}
           controls={controls}
           useLineage={useLineage}
           onNodeClick={onNodeClick}
diff --git a/web/common/src/components/Lineage/LineageLayoutContainer.tsx b/web/common/src/components/Lineage/LineageLayoutContainer.tsx
index e3385a3294..4bd0d42402 100644
--- a/web/common/src/components/Lineage/LineageLayoutContainer.tsx
+++ b/web/common/src/components/Lineage/LineageLayoutContainer.tsx
@@ -19,7 +19,7 @@ export function LineageLayoutContainer({
 }) {
   return (
     <VerticalContainer
-      data-component="LineageLayout"
+      data-component="LineageLayoutContainer"
       className={cn(
         'border-2 border-lineage-border bg-lineage-background relative h-full',
         className,
diff --git a/web/common/src/components/Lineage/stories/ModelLineage.tsx b/web/common/src/components/Lineage/stories/ModelLineage.tsx
index 800f292c4b..ddcaf47663 100644
--- a/web/common/src/components/Lineage/stories/ModelLineage.tsx
+++ b/web/common/src/components/Lineage/stories/ModelLineage.tsx
@@ -383,6 +383,7 @@ export const ModelLineage = ({
         ModelEdgeId,
         ModelColumnID
       >
+        isBuildingLayout={isBuildingLayout}
         useLineage={useModelLineage}
         nodeTypes={nodeTypes}
         edgeTypes={edgeTypes}