From 3abf40b56cb797c349aa27c82b9b44f2bef1f688 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adam=20Zapa=C5=9Bnik?=
 <9281806+adamzapasnik@users.noreply.github.com>
Date: Wed, 9 Jun 2021 09:30:02 +0300
Subject: [PATCH] Unmount React Components when removing elements from Dom

---
 __tests__/utilities/graph-util.test.js | 3 +++
 src/utilities/graph-util.js            | 2 ++
 2 files changed, 5 insertions(+)

diff --git a/__tests__/utilities/graph-util.test.js b/__tests__/utilities/graph-util.test.js
index c18de118..4e1dda7f 100644
--- a/__tests__/utilities/graph-util.test.js
+++ b/__tests__/utilities/graph-util.test.js
@@ -1,5 +1,6 @@
 // @flow
 import * as React from 'react';
+import ReactDOM from 'react-dom';
 import { shallow } from 'enzyme';
 import GraphView from '../../src/components/graph-view';
 
@@ -128,6 +129,7 @@ describe('GraphUtils class', () => {
 
   describe('removeElementFromDom method', () => {
     it('removes an element using an id', () => {
+      ReactDOM.unmountComponentAtNode = jest.fn();
       const fakeElement = {
         parentNode: {
           removeChild: jest.fn(),
@@ -137,6 +139,7 @@ describe('GraphUtils class', () => {
       jest.spyOn(document, 'querySelector').mockReturnValue(fakeElement);
       const result = GraphUtils.removeElementFromDom('fake');
 
+      expect(ReactDOM.unmountComponentAtNode).toHaveBeenCalledWith(fakeElement);
       expect(fakeElement.parentNode.removeChild).toHaveBeenCalledWith(
         fakeElement
       );
diff --git a/src/utilities/graph-util.js b/src/utilities/graph-util.js
index d06c3033..070cbba3 100644
--- a/src/utilities/graph-util.js
+++ b/src/utilities/graph-util.js
@@ -19,6 +19,7 @@ import { type IEdge } from '../components/edge';
 import { type INode } from '../components/node';
 import { type IPoint } from '../components/graph-view-props';
 import fastDeepEqual from 'fast-deep-equal';
+import ReactDOM from 'react-dom';
 
 export type INodeMapNode = {
   node: INode,
@@ -106,6 +107,7 @@ class GraphUtils {
     const container = searchElement.querySelector(`[id='${id}']`);
 
     if (container && container.parentNode) {
+      ReactDOM.unmountComponentAtNode(container);
       container.parentNode.removeChild(container);
 
       return true;