From 127b9230e0bb550894a87b747b5cea14db323662 Mon Sep 17 00:00:00 2001
From: emyarod <emyarod@users.noreply.github.com>
Date: Wed, 9 Sep 2020 13:59:29 -0500
Subject: [PATCH] fix(TreeNode): sync expanded prop and state (#6791)

* fix(TreeNode): sync expanded prop and state

* docs(TreeView): add controlled expansion example

Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>
---
 .../react/src/components/TreeView/TreeNode.js |  4 +--
 .../src/components/TreeView/TreeView-story.js | 31 ++++++++++++++++---
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/packages/react/src/components/TreeView/TreeNode.js b/packages/react/src/components/TreeView/TreeNode.js
index bd5d82aec208..4505f60477de 100644
--- a/packages/react/src/components/TreeView/TreeNode.js
+++ b/packages/react/src/components/TreeView/TreeNode.js
@@ -166,8 +166,8 @@ export default function TreeNode({
     }
 
     // sync props and state
-    setExpanded(expanded);
-  }, [children, depth, expanded, Icon]);
+    setExpanded(isExpanded);
+  }, [children, depth, Icon, isExpanded]);
 
   const treeNodeProps = {
     ...rest,
diff --git a/packages/react/src/components/TreeView/TreeView-story.js b/packages/react/src/components/TreeView/TreeView-story.js
index 306ff9d2f21a..0cdceff624c0 100644
--- a/packages/react/src/components/TreeView/TreeView-story.js
+++ b/packages/react/src/components/TreeView/TreeView-story.js
@@ -5,7 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import React from 'react';
+import React, { useState } from 'react';
 import { Document16, Folder16 } from '@carbon/icons-react';
 import { action } from '@storybook/addon-actions';
 import {
@@ -233,16 +233,17 @@ const nodes = [
   },
 ];
 
-function renderTree({ nodes, withIcons = false }) {
+function renderTree({ nodes, expanded, withIcons = false }) {
   if (!nodes) {
     return;
   }
-  return nodes.map(({ children, renderIcon, ...nodeProps }) => (
+  return nodes.map(({ children, renderIcon, isExpanded, ...nodeProps }) => (
     <TreeNode
       key={nodeProps.id}
       renderIcon={withIcons ? renderIcon : null}
+      isExpanded={expanded ?? isExpanded}
       {...nodeProps}>
-      {renderTree({ nodes: children, withIcons })}
+      {renderTree({ nodes: children, expanded, withIcons })}
     </TreeNode>
   ));
 }
@@ -283,3 +284,25 @@ export const WithIcons = () => (
 );
 
 WithIcons.storyName = 'with icons';
+
+export const WithControlledExpansion = () => {
+  const [expanded, setExpanded] = useState(undefined);
+  return (
+    <>
+      <InlineNotification
+        kind="info"
+        title="Experimental component"
+        subtitle="An accessibility review of this component is in progress"
+      />
+      <div style={{ marginBottom: '1rem' }}>
+        <button type="button" onClick={() => setExpanded(true)}>
+          expand all
+        </button>
+        <button type="button" onClick={() => setExpanded(false)}>
+          collapse all
+        </button>
+      </div>
+      <TreeView {...props()}>{renderTree({ nodes, expanded })}</TreeView>
+    </>
+  );
+};