From 79153272e0a615fba89bc8d097ece1660107e189 Mon Sep 17 00:00:00 2001
From: Ashish Gupta <ashish@getcollate.io>
Date: Thu, 16 Jan 2025 10:45:57 +0530
Subject: [PATCH] MINOR: fix the tracing highlighter after edge clicked
 (#19346)

* fix the tracing highlighter after edge clicked

* close the drawer appeariance after switching from edge click to node

* added test to validate
---
 .../LineageProvider/LineageProvider.test.tsx  | 28 +++++++++++++++++++
 .../LineageProvider/LineageProvider.tsx       |  4 +++
 2 files changed, 32 insertions(+)

diff --git a/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.test.tsx b/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.test.tsx
index ac1ed7562f0..de1b217d872 100644
--- a/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.test.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.test.tsx
@@ -45,6 +45,7 @@ const DummyChildrenComponent = () => {
   const {
     loadChildNodesHandler,
     onEdgeClick,
+    onColumnClick,
     updateEntityType,
     onLineageEditClick,
   } = useLineageProvider();
@@ -94,6 +95,11 @@ const DummyChildrenComponent = () => {
         onClick={() => onEdgeClick(MOCK_EDGE as Edge)}>
         On Edge Click
       </button>
+      <button
+        data-testid="column-click"
+        onClick={() => onColumnClick('column')}>
+        On Column Click
+      </button>
       <button data-testid="openConfirmationModal">
         Close Confirmation Modal
       </button>
@@ -245,4 +251,26 @@ describe('LineageProvider', () => {
 
     expect(edgeDrawer).toBeInTheDocument();
   });
+
+  it('should close the drawer if open, on column click', async () => {
+    await act(async () => {
+      render(
+        <LineageProvider>
+          <DummyChildrenComponent />
+        </LineageProvider>
+      );
+    });
+
+    const edgeClick = screen.getByTestId('edge-click');
+    fireEvent.click(edgeClick);
+
+    const edgeDrawer = screen.getByText('Edge Info Drawer');
+
+    expect(edgeDrawer).toBeInTheDocument();
+
+    const columnClick = screen.getByTestId('column-click');
+    fireEvent.click(columnClick);
+
+    expect(edgeDrawer).not.toBeInTheDocument();
+  });
 });
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 b966f3fcfea..838c358be0a 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
@@ -434,6 +434,8 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
 
       setTracedColumns(connectedColumnEdges);
       setTracedNodes([]);
+      setSelectedEdge(undefined);
+      setIsDrawerOpen(false);
     },
     [nodes, edges]
   );
@@ -710,6 +712,8 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
     setActiveNode(undefined);
     setSelectedNode({} as SourceType);
     setIsDrawerOpen(true);
+    setTracedNodes([]);
+    setTracedColumns([]);
   }, []);
 
   const onLineageEditClick = useCallback(() => {