fix the cycle lineage nodes when being collapsed (#22602)

* fix the cycle lineage nodes when being collapsed

* added unit test around removal of rootNode in collapse if there is cycle lineage happen

* added playwright test and added commetns where the fix operation is perforemed

openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/Lineage.spec.ts

@@ -446,7 +446,7 @@ test('Verify function data in edge drawer', async ({ browser }) => {
   }
 });
 
-test('Verify table search with special characters as handledd', async ({
+test('Verify table search with special characters as handled', async ({
   browser,
 }) => {
   const { page } = await createNewPage(browser);
@@ -518,3 +518,154 @@ test('Verify table search with special characters as handledd', async ({
     await afterAction();
   }
 });
+
+test('Verify cycle lineage should be handled properly', async ({ browser }) => {
+  test.slow();
+
+  const { page } = await createNewPage(browser);
+  const { apiContext, afterAction } = await getApiContext(page);
+  const table = new TableClass();
+  const topic = new TopicClass();
+  const dashboard = new DashboardClass();
+
+  try {
+    await Promise.all([
+      table.create(apiContext),
+      topic.create(apiContext),
+      dashboard.create(apiContext),
+    ]);
+
+    const tableFqn = get(table, 'entityResponseData.fullyQualifiedName');
+    const topicFqn = get(topic, 'entityResponseData.fullyQualifiedName');
+    const dashboardFqn = get(
+      dashboard,
+      'entityResponseData.fullyQualifiedName'
+    );
+
+    await redirectToHomePage(page);
+    await table.visitEntityPageWithCustomSearchBox(page);
+    await visitLineageTab(page);
+    await page.getByTestId('full-screen').click();
+    await editLineage(page);
+    await performZoomOut(page);
+
+    // connect table to topic
+    await connectEdgeBetweenNodes(page, table, topic);
+    await rearrangeNodes(page);
+
+    // connect topic to dashboard
+    await connectEdgeBetweenNodes(page, topic, dashboard);
+    await rearrangeNodes(page);
+
+    // connect dashboard to table
+    await connectEdgeBetweenNodes(page, dashboard, table);
+    await rearrangeNodes(page);
+
+    await page.reload();
+    await page.waitForLoadState('networkidle');
+    await page.getByTestId('fit-screen').click();
+
+    await expect(page.getByTestId(`lineage-node-${tableFqn}`)).toBeVisible();
+    await expect(page.getByTestId(`lineage-node-${topicFqn}`)).toBeVisible();
+    await expect(
+      page.getByTestId(`lineage-node-${dashboardFqn}`)
+    ).toBeVisible();
+
+    // Collapse the cycle dashboard lineage downstreamNodeHandler
+    await page
+      .getByTestId(`lineage-node-${dashboardFqn}`)
+      .getByTestId('downstream-collapse-handle')
+      .click();
+
+    await expect(
+      page.getByTestId(`edge-${dashboardFqn}-${tableFqn}`)
+    ).not.toBeVisible();
+
+    await expect(page.getByTestId(`lineage-node-${tableFqn}`)).toBeVisible();
+    await expect(page.getByTestId(`lineage-node-${topicFqn}`)).toBeVisible();
+    await expect(
+      page.getByTestId(`lineage-node-${dashboardFqn}`)
+    ).toBeVisible();
+
+    await expect(
+      page
+        .getByTestId(`lineage-node-${tableFqn}`)
+        .getByTestId('upstream-collapse-handle')
+    ).not.toBeVisible();
+
+    await expect(
+      page.getByTestId(`lineage-node-${dashboardFqn}`).getByTestId('plus-icon')
+    ).toBeVisible();
+
+    // Reclick the plus icon to expand the cycle dashboard lineage downstreamNodeHandler
+    const downstreamResponse = page.waitForResponse(
+      `/api/v1/lineage/getLineage/Downstream?fqn=${dashboardFqn}&type=dashboard**`
+    );
+    await page
+      .getByTestId(`lineage-node-${dashboardFqn}`)
+      .getByTestId('plus-icon')
+      .click();
+
+    await downstreamResponse;
+
+    await expect(
+      page
+        .getByTestId(`lineage-node-${tableFqn}`)
+        .getByTestId('upstream-collapse-handle')
+        .getByTestId('minus-icon')
+    ).toBeVisible();
+
+    // Click the Upstream Node to expand the cycle dashboard lineage
+    await page
+      .getByTestId(`lineage-node-${dashboardFqn}`)
+      .getByTestId('upstream-collapse-handle')
+      .click();
+
+    await expect(page.getByTestId(`lineage-node-${tableFqn}`)).toBeVisible();
+    await expect(
+      page.getByTestId(`lineage-node-${dashboardFqn}`)
+    ).toBeVisible();
+    await expect(
+      page.getByTestId(`lineage-node-${topicFqn}`)
+    ).not.toBeVisible();
+
+    await expect(
+      page.getByTestId(`lineage-node-${dashboardFqn}`).getByTestId('plus-icon')
+    ).toBeVisible();
+
+    // Reclick the plus icon to expand the cycle dashboard lineage upstreamNodeHandler
+    const upStreamResponse2 = page.waitForResponse(
+      `/api/v1/lineage/getLineage/Upstream?fqn=${dashboardFqn}&type=dashboard**`
+    );
+    await page
+      .getByTestId(`lineage-node-${dashboardFqn}`)
+      .getByTestId('plus-icon')
+      .click();
+    await upStreamResponse2;
+
+    await expect(page.getByTestId(`lineage-node-${tableFqn}`)).toBeVisible();
+    await expect(
+      page.getByTestId(`lineage-node-${dashboardFqn}`)
+    ).toBeVisible();
+    await expect(page.getByTestId(`lineage-node-${topicFqn}`)).toBeVisible();
+
+    // Collapse the Node from the Parent Cycle Node
+    await page
+      .getByTestId(`lineage-node-${topicFqn}`)
+      .getByTestId('downstream-collapse-handle')
+      .click();
+
+    await expect(page.getByTestId(`lineage-node-${tableFqn}`)).toBeVisible();
+    await expect(page.getByTestId(`lineage-node-${topicFqn}`)).toBeVisible();
+    await expect(
+      page.getByTestId(`lineage-node-${dashboardFqn}`)
+    ).not.toBeVisible();
+  } finally {
+    await Promise.all([
+      table.delete(apiContext),
+      topic.delete(apiContext),
+      dashboard.delete(apiContext),
+    ]);
+    await afterAction();
+  }
+});

openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.test.tsx

@@ -302,6 +302,122 @@ describe('Test EntityLineageUtils utility', () => {
     expect(emptyResult.nodes).toEqual([]);
   });
 
+  it('getConnectedNodesEdges should filter out root nodes from child nodes', () => {
+    const selectedNode = {
+      id: '1',
+      position: { x: 0, y: 0 },
+      data: { node: { fullyQualifiedName: '1' } },
+    };
+    const nodes = [
+      {
+        id: '1',
+        position: { x: 0, y: 0 },
+        data: { node: { fullyQualifiedName: '1' } },
+      },
+      {
+        id: '2',
+        position: { x: 0, y: 0 },
+        data: {
+          node: { fullyQualifiedName: '2' },
+          isRootNode: true, // This should be filtered out
+        },
+      },
+      {
+        id: '3',
+        position: { x: 0, y: 0 },
+        data: {
+          node: { fullyQualifiedName: '3' },
+          isRootNode: false, // This should be included
+        },
+      },
+      {
+        id: '4',
+        position: { x: 0, y: 0 },
+        data: {
+          node: { fullyQualifiedName: '4' },
+          // No isRootNode property - should be included
+        },
+      },
+    ];
+    const edges = [
+      { id: '1', source: '1', target: '2' },
+      { id: '2', source: '1', target: '3' },
+      { id: '3', source: '1', target: '4' },
+    ];
+    const direction = LineageDirection.Downstream;
+
+    const result = getConnectedNodesEdges(
+      selectedNode,
+      nodes,
+      edges,
+      direction
+    );
+
+    expect(result).toHaveProperty('nodes');
+    expect(result).toHaveProperty('edges');
+    expect(result).toHaveProperty('nodeFqn');
+
+    // Should only include nodes that are not root nodes
+    expect(result.nodes).toHaveLength(2);
+    expect(result.nodes.find((node) => node.id === '2')).toBeUndefined(); // Root node should be filtered out
+    expect(result.nodes.find((node) => node.id === '3')).toBeDefined(); // Non-root node should be included
+    expect(result.nodes.find((node) => node.id === '4')).toBeDefined(); // Node without isRootNode should be included
+
+    // Verify nodeFqn contains only the filtered nodes
+    expect(result.nodeFqn).toHaveLength(2);
+    expect(result.nodeFqn).toContain('3');
+    expect(result.nodeFqn).toContain('4');
+    expect(result.nodeFqn).not.toContain('2'); // Root node FQN should not be included
+  });
+
+  it('getConnectedNodesEdges should handle nodes with undefined isRootNode property', () => {
+    const selectedNode = {
+      id: '1',
+      position: { x: 0, y: 0 },
+      data: { node: { fullyQualifiedName: '1' } },
+    };
+    const nodes = [
+      {
+        id: '1',
+        position: { x: 0, y: 0 },
+        data: { node: { fullyQualifiedName: '1' } },
+      },
+      {
+        id: '2',
+        position: { x: 0, y: 0 },
+        data: {
+          node: { fullyQualifiedName: '2' },
+          isRootNode: undefined, // Should be treated as falsy and included
+        },
+      },
+      {
+        id: '3',
+        position: { x: 0, y: 0 },
+        data: {
+          node: { fullyQualifiedName: '3' },
+          isRootNode: null, // Should be treated as falsy and included
+        },
+      },
+    ];
+    const edges = [
+      { id: '1', source: '1', target: '2' },
+      { id: '2', source: '1', target: '3' },
+    ];
+    const direction = LineageDirection.Downstream;
+
+    const result = getConnectedNodesEdges(
+      selectedNode,
+      nodes,
+      edges,
+      direction
+    );
+
+    // Should include nodes with undefined/null isRootNode
+    expect(result.nodes).toHaveLength(2);
+    expect(result.nodes.find((node) => node.id === '2')).toBeDefined();
+    expect(result.nodes.find((node) => node.id === '3')).toBeDefined();
+  });
+
   it('should call addLineage with the provided edge', async () => {
     const edge = {
       edge: { fromEntity: {}, toEntity: {} },

openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx

@@ -1133,8 +1133,14 @@ export const getConnectedNodesEdges = (
               currentNodeID
             );
 
-      stack.push(...childNodes);
-      outgoers.push(...childNodes);
+      // Removing the Root Node from the Child Nodes here, which comes when a cycle lineage is formed
+      // So while collapsing the cycle lineage, we need to prevent the Root Node not to be removed.
+      const finalChildNodeRemovingRootNode = childNodes.filter(
+        (item) => !item.data.isRootNode
+      );
+
+      stack.push(...finalChildNodeRemovingRootNode);
+      outgoers.push(...finalChildNodeRemovingRootNode);
       connectedEdges.push(...childEdges);
     }
   }