Lineage node column pagination

openmetadata-ui/src/main/resources/ui/playwright/utils/lineage.ts

@@ -162,7 +162,7 @@ export const dragAndDropNode = async (
   await page.mouse.down();
   const box = (await destinationElement.boundingBox()) as DOMRect;
   const x = box.x + 250;
-  const y = box.y + box.height / 2;
+  const y = box.y + box.height / 2 + 100;
   await page.mouse.move(x, y, { steps: 20 });
   await page.mouse.up();
 };
@@ -475,6 +475,8 @@ export const addColumnLineage = async (
   );
   await lineageRes;
 
+  await page.getByTestId(`column-${toColumnNode}`).click();
+
   if (exitEditMode) {
     await editLineageClick(page);
   }

openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomEdge.component.test.tsx

@@ -11,9 +11,13 @@
  *  limitations under the License.
  */
 
+import { createTheme, Theme, ThemeProvider } from '@mui/material/styles';
+import { ThemeColors } from '@openmetadata/ui-core-components';
 import { render, screen } from '@testing-library/react';
+import React from 'react';
 import { MemoryRouter } from 'react-router-dom';
 import { EdgeProps, Position } from 'reactflow';
+import { useLineageProvider } from '../../../context/LineageProvider/LineageProvider';
 import { EntityType } from '../../../enums/entity.enum';
 import { LineageLayer } from '../../../generated/settings/settings';
 import { CustomEdge } from './CustomEdge.component';
@@ -61,6 +65,7 @@ const mockCustomEdgeProp = {
       },
     },
     isEditMode: true,
+    dataTestId: 'edge-id1',
   },
   selected: true,
 } as EdgeProps;
@@ -72,13 +77,50 @@ jest.mock('../../../context/LineageProvider/LineageProvider', () => ({
     pipelineStatus: {},
     activeLayer: [LineageLayer.ColumnLevelLineage],
     fetchPipelineStatus: jest.fn(),
+    columnsInCurrentPages: [],
   })),
 }));
 
+const mockThemeColors: ThemeColors = {
+  indigo: {
+    600: '#4F46E5',
+  },
+  error: {
+    600: '#D92D20',
+  },
+} as ThemeColors;
+
+const theme: Theme = createTheme({
+  palette: {
+    primary: {
+      main: '#0958D9',
+    },
+    allShades: mockThemeColors,
+  },
+});
+
+const Wrapper = ({ children }: { children: React.ReactNode }) => (
+  <ThemeProvider theme={theme}>
+    <MemoryRouter>{children}</MemoryRouter>
+  </ThemeProvider>
+);
+
+const mockLineageProvider = (overrides = {}) => {
+  (useLineageProvider as jest.Mock).mockImplementation(() => ({
+    tracedNodes: [],
+    tracedColumns: [],
+    pipelineStatus: {},
+    activeLayer: [LineageLayer.ColumnLevelLineage],
+    fetchPipelineStatus: jest.fn(),
+    columnsInCurrentPages: [],
+    ...overrides,
+  }));
+};
+
 describe('Test CustomEdge Component', () => {
   it('Check if CustomEdge has selected as false', async () => {
     render(<CustomEdge {...mockCustomEdgeProp} selected={false} />, {
-      wrapper: MemoryRouter,
+      wrapper: Wrapper,
     });
 
     const deleteButton = screen.queryByTestId('delete-button');
@@ -101,7 +143,7 @@ describe('Test CustomEdge Component', () => {
         }}
       />,
       {
-        wrapper: MemoryRouter,
+        wrapper: Wrapper,
       }
     );
 
@@ -111,4 +153,143 @@ describe('Test CustomEdge Component', () => {
 
     expect(pipelineLabelAsEdge).toBeInTheDocument();
   });
+
+  it('should display edge when no nodes or columns are traced', () => {
+    render(<CustomEdge {...mockCustomEdgeProp} />, {
+      wrapper: Wrapper,
+    });
+
+    const edgePath = screen.getByTestId('edge-id1');
+
+    expect(edgePath).toBeInTheDocument();
+    expect(edgePath).toHaveStyle({ display: 'block' });
+  });
+
+  it('should hide edge when nodes are traced but current edge nodes are not in traced list', () => {
+    mockLineageProvider({ tracedNodes: ['other-node-1', 'other-node-2'] });
+
+    render(<CustomEdge {...mockCustomEdgeProp} />, {
+      wrapper: Wrapper,
+    });
+
+    const edgePath = screen.getByTestId('edge-id1');
+
+    expect(edgePath).toBeInTheDocument();
+    expect(edgePath).toHaveStyle({ display: 'none' });
+  });
+
+  it('should display edge when both edge nodes are in traced nodes list', () => {
+    mockLineageProvider({ tracedNodes: ['1', '2'] });
+
+    render(<CustomEdge {...mockCustomEdgeProp} />, {
+      wrapper: Wrapper,
+    });
+
+    const edgePath = screen.getByTestId('edge-id1');
+
+    expect(edgePath).toBeInTheDocument();
+    expect(edgePath).toHaveStyle({ display: 'block' });
+  });
+
+  it('should display column lineage edge when columns are highlighted', () => {
+    const encodedSourceHandle = btoa(encodeURIComponent('column1'));
+    const encodedTargetHandle = btoa(encodeURIComponent('column2'));
+
+    mockLineageProvider({ tracedColumns: ['column1', 'column2'] });
+
+    render(
+      <CustomEdge
+        {...mockCustomEdgeProp}
+        data={{
+          ...mockCustomEdgeProp.data,
+          isColumnLineage: true,
+          sourceHandle: encodedSourceHandle,
+          targetHandle: encodedTargetHandle,
+        }}
+      />,
+      {
+        wrapper: Wrapper,
+      }
+    );
+
+    const edgePath = screen.getByTestId('edge-id1');
+
+    expect(edgePath).toBeInTheDocument();
+    expect(edgePath).toHaveStyle({ display: 'block' });
+  });
+
+  it('should hide column lineage edge when columns are traced but not highlighted', () => {
+    mockLineageProvider({ tracedColumns: ['other-column'] });
+
+    render(
+      <CustomEdge
+        {...mockCustomEdgeProp}
+        data={{
+          ...mockCustomEdgeProp.data,
+          isColumnLineage: true,
+          sourceHandle: 'column1',
+          targetHandle: 'column2',
+        }}
+      />,
+      {
+        wrapper: Wrapper,
+      }
+    );
+
+    const edgePath = screen.getByTestId('edge-id1');
+
+    expect(edgePath).toBeInTheDocument();
+    expect(edgePath).toHaveStyle({ display: 'none' });
+  });
+
+  it('should hide column lineage edge when columns are not in current page', () => {
+    mockLineageProvider({ columnsInCurrentPages: ['other-column'] });
+
+    render(
+      <CustomEdge
+        {...mockCustomEdgeProp}
+        data={{
+          ...mockCustomEdgeProp.data,
+          isColumnLineage: true,
+          sourceHandle: 'column1',
+          targetHandle: 'column2',
+        }}
+      />,
+      {
+        wrapper: Wrapper,
+      }
+    );
+
+    const edgePath = screen.getByTestId('edge-id1');
+
+    expect(edgePath).toBeInTheDocument();
+    expect(edgePath).toHaveStyle({ display: 'none' });
+  });
+
+  it('should display column lineage edge when both columns are in current page', () => {
+    const encodedSourceHandle = btoa(encodeURIComponent('column1'));
+    const encodedTargetHandle = btoa(encodeURIComponent('column2'));
+
+    mockLineageProvider({ columnsInCurrentPages: ['column1', 'column2'] });
+
+    render(
+      <CustomEdge
+        {...mockCustomEdgeProp}
+        data={{
+          ...mockCustomEdgeProp.data,
+          isColumnLineage: true,
+          sourceHandle: encodedSourceHandle,
+          targetHandle: encodedTargetHandle,
+        }}
+      />,
+      {
+        wrapper: Wrapper,
+      }
+    );
+
+    const edgePath = screen.getByTestId('edge-id1');
+
+    expect(edgePath).toBeInTheDocument();
+    expect(edgePath).toHaveStyle({ display: 'block' });
+  });
 });