diff --git a/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/CustomEdge.component.test.tsx b/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/CustomEdge.component.test.tsx
new file mode 100644
index 00000000000..c86a9bc504f
--- /dev/null
+++ b/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/CustomEdge.component.test.tsx
@@ -0,0 +1,82 @@
+/*
+ *  Copyright 2021 Collate
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+import { findByTestId, queryByTestId, render } from '@testing-library/react';
+import React from 'react';
+import { ArrowHeadType, EdgeProps, Position } from 'react-flow-renderer';
+import { MemoryRouter } from 'react-router-dom';
+import { CustomEdge } from './CustomEdge.component';
+
+jest.mock('../../constants/Lineage.constants', () => ({
+  foreignObjectSize: 40,
+}));
+
+const mockCustomEdgeProp = {
+  id: 'id1',
+  sourceX: 20,
+  sourceY: 20,
+  targetX: 20,
+  targetY: 20,
+  sourcePosition: Position.Left,
+  targetPosition: Position.Right,
+  style: {},
+  arrowHeadType: ArrowHeadType.ArrowClosed,
+  markerEndId: '',
+  data: {
+    source: 'node1',
+    target: 'node2',
+    onEdgeClick: jest.fn(),
+    selectedNode: {
+      id: 'node1',
+    },
+  },
+} as EdgeProps;
+
+describe('Test CustomEdge Component', () => {
+  it('Check if CustomeEdge has all child elements', async () => {
+    const { container } = render(<CustomEdge {...mockCustomEdgeProp} />, {
+      wrapper: MemoryRouter,
+    });
+
+    const deleteButton = await findByTestId(container, 'delete-button');
+    const edgePathElement = await findByTestId(
+      container,
+      'react-flow-edge-path'
+    );
+
+    expect(deleteButton).toBeInTheDocument();
+    expect(edgePathElement).toBeInTheDocument();
+  });
+
+  it('Check if CustomeEdge has selectedNode as empty object', async () => {
+    const { container } = render(
+      <CustomEdge
+        {...mockCustomEdgeProp}
+        data={{ ...mockCustomEdgeProp.data, selectedNode: {} }}
+      />,
+      {
+        wrapper: MemoryRouter,
+      }
+    );
+
+    const edgePathElement = await findByTestId(
+      container,
+      'react-flow-edge-path'
+    );
+
+    const deleteButton = queryByTestId(container, 'delete-button');
+
+    expect(deleteButton).not.toBeInTheDocument();
+    expect(edgePathElement).toBeInTheDocument();
+  });
+});
diff --git a/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/CustomEdge.component.tsx b/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/CustomEdge.component.tsx
index a734cd132f4..99401e4aea1 100644
--- a/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/CustomEdge.component.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/CustomEdge.component.tsx
@@ -35,7 +35,8 @@ export const CustomEdge = ({
   markerEndId,
   data,
 }: EdgeProps) => {
-  const { onEdgeClick, ...rest } = data;
+  const { onEdgeClick, selectedNode, ...rest } = data;
+
   const edgePath = getBezierPath({
     sourceX,
     sourceY,
@@ -57,23 +58,31 @@ export const CustomEdge = ({
       <path
         className="react-flow__edge-path"
         d={edgePath}
+        data-testid="react-flow-edge-path"
         id={id}
         markerEnd={markerEnd}
         style={style}
       />
-      <foreignObject
-        className="tw-opacity-0 hover:tw-opacity-100"
-        height={foreignObjectSize}
-        requiredExtensions="http://www.w3.org/1999/xhtml"
-        width={foreignObjectSize}
-        x={edgeCenterX - foreignObjectSize / 4}
-        y={edgeCenterY - foreignObjectSize / 4}>
-        <button
-          className="tw-cursor-pointer tw-flex tw-z-9999"
-          onClick={(event) => onEdgeClick?.(event, rest as CustomEdgeData)}>
-          <SVGIcons alt="times-circle" icon="icon-times-circle" width="14px" />
-        </button>
-      </foreignObject>
+      {(rest as CustomEdgeData)?.source?.includes(selectedNode?.id) ||
+      (rest as CustomEdgeData)?.target?.includes(selectedNode?.id) ? (
+        <foreignObject
+          data-testid="delete-button"
+          height={foreignObjectSize}
+          requiredExtensions="http://www.w3.org/1999/xhtml"
+          width={foreignObjectSize}
+          x={edgeCenterX - foreignObjectSize / 4}
+          y={edgeCenterY - foreignObjectSize / 4}>
+          <button
+            className="tw-cursor-pointer tw-flex tw-z-9999"
+            onClick={(event) => onEdgeClick?.(event, rest as CustomEdgeData)}>
+            <SVGIcons
+              alt="times-circle"
+              icon="icon-times-circle"
+              width="14px"
+            />
+          </button>
+        </foreignObject>
+      ) : null}
     </Fragment>
   );
 };
diff --git a/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/EntityLineage.component.tsx b/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/EntityLineage.component.tsx
index 356137c8d88..f8b565226c4 100644
--- a/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/EntityLineage.component.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/EntityLineage.component.tsx
@@ -215,6 +215,9 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
       // eslint-disable-next-line @typescript-eslint/no-use-before-define
       setElements((es) => es.filter((e) => e.id !== data.id));
 
+      /**
+       * Get new downstreamEdges
+       */
       const newDownStreamEdges = lineageData.downstreamEdges?.filter(
         (dn) =>
           !lineageData.downstreamEdges?.find(
@@ -222,6 +225,10 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
               edgeData.fromId === dn.fromEntity && edgeData.toId === dn.toEntity
           )
       );
+
+      /**
+       * Get new upstreamEdges
+       */
       const newUpStreamEdges = lineageData.upstreamEdges?.filter(
         (up) =>
           !lineageData.upstreamEdges?.find(
@@ -230,10 +237,20 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
           )
       );
 
+      /**
+       * Get new nodes that have either downstreamEdge or upstreamEdge
+       */
+      const newNodes = lineageData.nodes?.filter(
+        (n) =>
+          !isUndefined(newDownStreamEdges?.find((d) => d.toEntity === n.id)) ||
+          !isUndefined(newUpStreamEdges?.find((u) => u.fromEntity === n.id))
+      );
+
       setNewAddedNode({} as FlowElement);
       setSelectedEntity({} as EntityReference);
       entityLineageHandler({
         ...lineageData,
+        nodes: newNodes,
         downstreamEdges: newDownStreamEdges,
         upstreamEdges: newUpStreamEdges,
       });
@@ -698,6 +715,9 @@ const Entitylineage: FunctionComponent<EntityLineageProp> = ({
     if (!isEmpty(selectedNode)) {
       setExpandNode(undefined);
     }
+    setElements((pre) => {
+      return pre.map((el) => ({ ...el, data: { ...el.data, selectedNode } }));
+    });
   }, [selectedNode]);
 
   useEffect(() => {