fix(tree@v1): ensures initialDepth is reapplied if dataset changes (#47)

Author: bkrem

src/Tree/index.js

@@ -19,10 +19,10 @@ class Tree extends React.Component {
     d3: Tree.calculateD3Geometry(this.props),
     rd3tSvgClassName: `_${uuid.v4()}`,
     rd3tGClassName: `_${uuid.v4()}`,
+    isInitialRenderForDataset: true,
   };
 
   internalState = {
-    initialRender: true,
     targetNode: null,
     isTransitioning: false,
   };
@@ -36,6 +36,7 @@ class Tree extends React.Component {
         // eslint-disable-next-line react/no-unused-state
         dataRef: nextProps.data,
         data: Tree.assignInternalProperties(clone(nextProps.data)),
+        isInitialRenderForDataset: true,
       };
     }
 
@@ -50,10 +51,17 @@ class Tree extends React.Component {
 
   componentDidMount() {
     this.bindZoomListener(this.props);
-    this.internalState.initialRender = false;
+    // eslint-disable-next-line
+    this.setState({ isInitialRenderForDataset: false });
   }
 
   componentDidUpdate(prevProps) {
+    if (this.props.data !== prevProps.data) {
+      // If last `render` was due to change in dataset -> mark the initial render as done.
+      // eslint-disable-next-line
+      this.setState({ isInitialRenderForDataset: false });
+    }
+
     // If zoom-specific props change -> rebind listener with new values
     // Or: rebind zoom listeners to new DOM nodes in case NodeWrapper switched <TransitionGroup> <-> <g>
     if (
@@ -423,24 +431,29 @@ class Tree extends React.Component {
       nodeSize,
       orientation,
     } = this.props;
+    const { isInitialRenderForDataset } = this.state;
 
     const tree = layout
       .tree()
       .nodeSize(orientation === 'horizontal' ? [nodeSize.y, nodeSize.x] : [nodeSize.x, nodeSize.y])
       .separation(
         (a, b) => (a.parent.id === b.parent.id ? separation.siblings : separation.nonSiblings),
       )
-      .children(d => (d._collapsed ? null : d._children));
+      .children(d => {
+        if (initialDepth !== undefined && isInitialRenderForDataset) {
+          // If `initialDepth` is defined, a node's depth determines
+          // whether we append `children` on the first render.
+          return d.depth >= initialDepth ? null : d._children;
+        }
+        // Node's `collapsed` property determines appending of `children` for subsequent renders.
+        return d._collapsed ? null : d._children;
+      });
 
     const rootNode = this.state.data[0];
     let nodes = tree.nodes(rootNode);
 
     // set `initialDepth` on first render if specified
-    if (
-      useCollapseData === false &&
-      initialDepth !== undefined &&
-      this.internalState.initialRender
-    ) {
+    if (useCollapseData === false && initialDepth !== undefined && isInitialRenderForDataset) {
       this.setInitialTreeDepth(nodes, initialDepth);
       nodes = tree.nodes(rootNode);
     }

src/Tree/tests/index.test.js

@@ -1,3 +1,5 @@
+/* eslint-disable camelcase */
+
 import React from 'react';
 import { shallow, mount } from 'enzyme';
 import { render } from 'react-dom';
@@ -6,7 +8,7 @@ import NodeWrapper from '../NodeWrapper';
 import Node from '../../Node';
 import Link from '../../Link';
 import Tree from '../index';
-import { mockData, mockData2, mockData3, mockData4 } from './mockData';
+import { mockData, mockData2, mockData3, mockData4, mockTree_D1N2_D2N2 } from './mockData';
 
 describe('<Tree />', () => {
   jest.spyOn(Tree.prototype, 'generateTree');
@@ -270,14 +272,19 @@ describe('<Tree />', () => {
   });
 
   describe('initialDepth', () => {
-    it('sets tree depth to `props.initialDepth` if specified', () => {
-      mount(<Tree data={mockData} initialDepth={1} />);
-      expect(Tree.prototype.setInitialTreeDepth).toHaveBeenCalled();
+    it('expands tree to full depth by default', () => {
+      const renderedComponent = shallow(<Tree data={mockTree_D1N2_D2N2} />);
+      expect(renderedComponent.find(Node).length).toBe(5);
+    });
+
+    it('expands tree to `props.initialDepth` if specified', () => {
+      const renderedComponent = shallow(<Tree data={mockTree_D1N2_D2N2} initialDepth={1} />);
+      expect(renderedComponent.find(Node).length).toBe(3);
     });
 
-    it('does not set an initialDepth if `props.useCollapseData` is true', () => {
-      mount(<Tree data={mockData} initialDepth={1} useCollapseData />);
-      expect(Tree.prototype.setInitialTreeDepth).not.toHaveBeenCalled();
+    it('renders only the root node if `initialDepth === 0`', () => {
+      const renderedComponent = shallow(<Tree data={mockTree_D1N2_D2N2} initialDepth={0} />);
+      expect(renderedComponent.find(Node).length).toBe(1);
     });
   });
 
@@ -746,7 +753,6 @@ describe('<Tree />', () => {
         .first()
         .simulate('click'); // collapse
 
-      expect(onUpdateSpy).toHaveBeenCalledTimes(1);
       expect(onUpdateSpy).toHaveBeenCalledWith({
         node: expect.any(Object),
         zoom: 1,
@@ -764,7 +770,6 @@ describe('<Tree />', () => {
       );
       const scrollableComponent = document.querySelector('.rd3t-tree-container > svg');
       scrollableComponent.dispatchEvent(new Event('wheel'));
-      expect(onUpdateSpy).toHaveBeenCalledTimes(1);
       expect(onUpdateSpy).toHaveBeenCalledWith({
         node: null,
         translate: { x: expect.any(Number), y: expect.any(Number) },
@@ -798,7 +803,6 @@ describe('<Tree />', () => {
         .first()
         .simulate('click');
 
-      expect(onUpdateSpy).toHaveBeenCalledTimes(1);
       expect(onUpdateSpy).toHaveBeenCalledWith({
         node: expect.any(Object),
         translate,

src/Tree/tests/mockData.js

@@ -1,3 +1,47 @@
+/* eslint-disable camelcase */
+
+export const mockTree_D1N2_D2N2 = [
+  {
+    name: 'Top Level',
+    attributes: {
+      keyA: 'val A',
+      keyB: 'val B',
+      keyC: 'val C',
+    },
+    children: [
+      {
+        name: 'Level 2: A',
+        attributes: {
+          keyA: 'val A',
+          keyB: 'val B',
+          keyC: 'val C',
+        },
+        children: [
+          {
+            name: '3: Son of A',
+            attributes: {
+              keyA: 'val A',
+              keyB: 'val B',
+              keyC: 'val C',
+            },
+          },
+          {
+            name: '3: Daughter of A',
+            attributes: {
+              keyA: 'val A',
+              keyB: 'val B',
+              keyC: 'val C',
+            },
+          },
+        ],
+      },
+      {
+        name: 'Level 2: B',
+      },
+    ],
+  },
+];
+
 const mockData = [
   {
     name: 'Top Level',