diff --git a/openmetadata-ui/src/main/resources/ui/src/components/Lineage/Lineage.component.tsx b/openmetadata-ui/src/main/resources/ui/src/components/Lineage/Lineage.component.tsx
index b5d3d5de2ae..071065c633b 100644
--- a/openmetadata-ui/src/main/resources/ui/src/components/Lineage/Lineage.component.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/Lineage/Lineage.component.tsx
@@ -131,6 +131,7 @@ const Lineage = ({
       <div
         className="h-full relative lineage-container"
         data-testid="lineage-container"
+        id="lineage-container" // ID is required for export PNG functionality
         ref={reactFlowWrapper}>
         {entityLineage && (
           <>
diff --git a/openmetadata-ui/src/main/resources/ui/src/constants/Export.constants.ts b/openmetadata-ui/src/main/resources/ui/src/constants/Export.constants.ts
index 385f53fb0f4..8bd78181bba 100644
--- a/openmetadata-ui/src/main/resources/ui/src/constants/Export.constants.ts
+++ b/openmetadata-ui/src/main/resources/ui/src/constants/Export.constants.ts
@@ -15,3 +15,6 @@ export enum ExportTypes {
   PNG = 'PNG',
   PDF = 'PDF',
 }
+
+export const LINEAGE_EXPORT_SELECTOR =
+  '#lineage-container .react-flow__viewport';
diff --git a/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx b/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx
index 88d06171011..0dfa7235446 100644
--- a/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx
@@ -56,7 +56,10 @@ import {
 import LineageNodeRemoveButton from '../../components/Lineage/LineageNodeRemoveButton';
 import { SourceType } from '../../components/SearchedData/SearchedData.interface';
 import { ROUTES } from '../../constants/constants';
-import { ExportTypes } from '../../constants/Export.constants';
+import {
+  ExportTypes,
+  LINEAGE_EXPORT_SELECTOR,
+} from '../../constants/Export.constants';
 import {
   ELEMENT_DELETE_STATE,
   ZOOM_VALUE,
@@ -460,8 +463,8 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
               name: decodedFqn,
               exportTypes: [ExportTypes.CSV, ExportTypes.PNG],
             }),
-        documentSelector: '.react-flow__viewport',
-        viewport: getViewportForLineageExport(nodes),
+        documentSelector: LINEAGE_EXPORT_SELECTOR,
+        viewport: getViewportForLineageExport(nodes, LINEAGE_EXPORT_SELECTOR),
         onExport: exportLineageData,
       });
     }
diff --git a/openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx b/openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx
index 7cc02ecf127..fa82d8ca264 100644
--- a/openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx
@@ -1717,10 +1717,11 @@ export const getViewportForBoundsReactFlow = (
   return { x: translateX, y: translateY, zoom: scale };
 };
 
-export const getViewportForLineageExport = (nodes: Node[]): ExportViewport => {
-  const exportElement = document.querySelector(
-    '.react-flow__viewport'
-  ) as HTMLElement;
+export const getViewportForLineageExport = (
+  nodes: Node[],
+  documentSelector: string
+): ExportViewport => {
+  const exportElement = document.querySelector(documentSelector) as HTMLElement;
 
   const imageWidth = exportElement.scrollWidth;
   const imageHeight = exportElement.scrollHeight;