From dcaa44365efbd9a0d854d5890091b572f74e746f Mon Sep 17 00:00:00 2001
From: Laura Silva <91160746+silvalaura@users.noreply.github.com>
Date: Fri, 5 Apr 2024 17:21:25 -0400
Subject: [PATCH] feat(TreeView): Fixing initialSelectedItems (#1225)

---
 .changeset/treeview-4.md                      |   5 +
 .../components/TreeView/TreeView.stories.tsx  | 389 +++++++++++++++---
 .../src/components/TreeView/TreeView.test.js  | 131 +++++-
 .../components/TreeView/TreeViewContext.ts    |  14 +-
 .../src/components/TreeView/useTreeItem.ts    | 125 +++---
 .../src/components/TreeView/useTreeView.ts    |   7 +
 .../src/components/TreeView/utils.ts          |  58 ++-
 7 files changed, 584 insertions(+), 145 deletions(-)
 create mode 100644 .changeset/treeview-4.md

diff --git a/.changeset/treeview-4.md b/.changeset/treeview-4.md
new file mode 100644
index 000000000..2d5e6c2d0
--- /dev/null
+++ b/.changeset/treeview-4.md
@@ -0,0 +1,5 @@
+---
+"react-magma-dom": minor
+---
+
+feat(TreeView): TreeView & TreeItem updates
\ No newline at end of file
diff --git a/packages/react-magma-dom/src/components/TreeView/TreeView.stories.tsx b/packages/react-magma-dom/src/components/TreeView/TreeView.stories.tsx
index 1359aa2ba..1097ae25d 100644
--- a/packages/react-magma-dom/src/components/TreeView/TreeView.stories.tsx
+++ b/packages/react-magma-dom/src/components/TreeView/TreeView.stories.tsx
@@ -90,25 +90,25 @@ export const Default = args => {
     setSelectedItems(selected);
     setIndeterminateItems(indet);
     setTotal(items.length);
-    // console.log('onSelection event:', items);
+    console.log('onSelection event:', items);
   }
 
   return (
     <>
       <Card isInverse={args.isInverse}>
         <TreeView {...args} onSelectedItemChange={onSelection}>
-          {/* <TreeItem label={<>Part 1: Introduction</>} itemId="pt1" testId="pt1">
+          <TreeItem label={<>Part 1: Introduction</>} itemId="pt1" testId="pt1">
             <TreeItem
               icon={<ArticleIcon />}
               label={<>Chapter 1: I love tiramisu jelly beans soufflé</>}
               itemId="pt1ch1"
               testId="pt1ch1"
+              isDisabled
             >
               <TreeItem
                 icon={<ArticleIcon />}
                 label={<>Section 1.1: Cake donut lemon drops gingerbread</>}
                 itemId="pt1ch1.1"
-                // isDisabled
               />
             </TreeItem>
             <TreeItem
@@ -155,7 +155,7 @@ export const Default = args => {
                 itemId="pt1ch3.3"
               />
             </TreeItem>
-          </TreeItem> */}
+          </TreeItem>
           <TreeItem
             icon={<FolderIcon />}
             label={
@@ -310,7 +310,9 @@ export const Default = args => {
         <>
           <p>{total} total</p>
           <p>Selected: {selectedItems}</p>
-          {args.selectable === TreeViewSelectable.multi && (<p>Indeterminate: {indeterminateItems}</p>)}
+          {args.selectable === TreeViewSelectable.multi && (
+            <p>Indeterminate: {indeterminateItems}</p>
+          )}
         </>
       )}
     </>
@@ -320,13 +322,16 @@ export const Default = args => {
 Default.args = {
   selectable: TreeViewSelectable.multi,
   ariaLabel: 'Textbook tree',
-  initialExpandedItems: [],
+  initialExpandedItems: ['pt1', 'pt1ch1'],
   initialSelectedItems: [
-    // { itemId: 'pt1ch1', checkedStatus: IndeterminateCheckboxStatus.checked },
-    // { itemId: 'pt1', checkedStatus: IndeterminateCheckboxStatus.indeterminate },
+    { itemId: 'pt1ch1', checkedStatus: IndeterminateCheckboxStatus.checked },
+    { itemId: 'pt1', checkedStatus: IndeterminateCheckboxStatus.indeterminate },
     { itemId: 'pt2ch4', checkedStatus: IndeterminateCheckboxStatus.checked },
     // { itemId: 'pt2ch5', checkedStatus: IndeterminateCheckboxStatus.checked },
-    { itemId: 'pt2ch5.1.1', checkedStatus: IndeterminateCheckboxStatus.checked },
+    {
+      itemId: 'pt2ch5.1.1',
+      checkedStatus: IndeterminateCheckboxStatus.checked,
+    },
     { itemId: 'pt2ch5.2', checkedStatus: IndeterminateCheckboxStatus.checked },
     { itemId: 'pt2ch5.3', checkedStatus: IndeterminateCheckboxStatus.checked },
   ],
@@ -405,7 +410,9 @@ export const NoIcons = args => {
       {args.selectable !== TreeViewSelectable.off && (
         <>
           <p>Selected: {selectedItems}</p>
-          {args.selectable === TreeViewSelectable.multi && (<p>Indeterminate: {indeterminateItems}</p>)}
+          {args.selectable === TreeViewSelectable.multi && (
+            <p>Indeterminate: {indeterminateItems}</p>
+          )}
         </>
       )}
     </>
@@ -421,16 +428,32 @@ NoIcons.args = {
 
 export const Textbook = args => {
   const [selectedItems, setSelectedItems] = React.useState(null);
+  const [indeterminateItems, setIndeterminateItems] = React.useState(null);
+  const [total, setTotal] = React.useState(null);
 
   function onSelection(items) {
-    const allTags = items.map((i: any, key) => {
-      return (
-        <Tag key={key} size={TagSize.small}>
-          {i}
+    const selected = items
+      .filter(i => i.checkedStatus === IndeterminateCheckboxStatus.checked)
+      .map((i, key) => (
+        <Tag key={key} size={TagSize.small} color={TagColor.primary}>
+          {i.itemId}
+        </Tag>
+      ));
+
+    const indet = items
+      .filter(
+        i => i.checkedStatus === IndeterminateCheckboxStatus.indeterminate
+      )
+      .map((i, key) => (
+        <Tag key={key} size={TagSize.small} color={TagColor.default}>
+          {i.itemId}
         </Tag>
-      );
-    });
-    setSelectedItems(allTags);
+      ));
+
+    setSelectedItems(selected);
+    setIndeterminateItems(indet);
+    setTotal(items.length);
+    // console.log('onSelection event:', items);
   }
 
   return (
@@ -503,7 +526,13 @@ export const Textbook = args => {
       </Card>
       <br />
       {args.selectable !== TreeViewSelectable.off && (
-        <>Selected: {selectedItems}</>
+        <>
+          <p>{total} total</p>
+          <p>Selected: {selectedItems}</p>
+          {args.selectable === TreeViewSelectable.multi && (
+            <p>Indeterminate: {indeterminateItems}</p>
+          )}
+        </>
       )}
     </>
   );
@@ -516,6 +545,7 @@ Textbook.args = {
 export const Simple = args => {
   const [selectedItems, setSelectedItems] = React.useState(null);
   const [indeterminateItems, setIndeterminateItems] = React.useState(null);
+  const [total, setTotal] = React.useState(null);
 
   function onSelection(items) {
     const selected = items
@@ -538,7 +568,8 @@ export const Simple = args => {
 
     setSelectedItems(selected);
     setIndeterminateItems(indet);
-    console.log('onSelection event:', items);
+    setTotal(items.length);
+    // console.log('onSelection event:', items);
   }
 
   return (
@@ -564,8 +595,11 @@ export const Simple = args => {
       <br />
       {args.selectable !== TreeViewSelectable.off && (
         <>
+          <p>{total} total</p>
           <p>Selected: {selectedItems}</p>
-          {args.selectable === TreeViewSelectable.multi && (<p>Indeterminate: {indeterminateItems}</p>)}
+          {args.selectable === TreeViewSelectable.multi && (
+            <p>Indeterminate: {indeterminateItems}</p>
+          )}
         </>
       )}
     </>
@@ -637,7 +671,9 @@ export const DefaultIcon = args => {
       {args.selectable !== TreeViewSelectable.off && (
         <>
           <p>Selected: {selectedItems}</p>
-          {args.selectable === TreeViewSelectable.multi && (<p>Indeterminate: {indeterminateItems}</p>)}
+          {args.selectable === TreeViewSelectable.multi && (
+            <p>Indeterminate: {indeterminateItems}</p>
+          )}
         </>
       )}
     </>
@@ -648,16 +684,30 @@ DefaultIcon.parameters = { controls: { exclude: ['isInverse'] } };
 
 export const FirstItemLeaf = args => {
   const [selectedItems, setSelectedItems] = React.useState(null);
+  const [indeterminateItems, setIndeterminateItems] = React.useState(null);
 
   function onSelection(items) {
-    const allTags = items.map((i: any, key) => {
-      return (
-        <Tag key={key} size={TagSize.small}>
-          {i}
+    const selected = items
+      .filter(i => i.checkedStatus === IndeterminateCheckboxStatus.checked)
+      .map((i, key) => (
+        <Tag key={key} size={TagSize.small} color={TagColor.primary}>
+          {i.itemId}
         </Tag>
-      );
-    });
-    setSelectedItems(allTags);
+      ));
+
+    const indet = items
+      .filter(
+        i => i.checkedStatus === IndeterminateCheckboxStatus.indeterminate
+      )
+      .map((i, key) => (
+        <Tag key={key} size={TagSize.small} color={TagColor.default}>
+          {i.itemId}
+        </Tag>
+      ));
+
+    setSelectedItems(selected);
+    setIndeterminateItems(indet);
+    // console.log('onSelection event:', items);
   }
 
   return (
@@ -683,7 +733,12 @@ export const FirstItemLeaf = args => {
       </TreeView>
       <br />
       {args.selectable !== TreeViewSelectable.off && (
-        <>Selected: {selectedItems}</>
+        <>
+          <p>Selected: {selectedItems}</p>
+          {args.selectable === TreeViewSelectable.multi && (
+            <p>Indeterminate: {indeterminateItems}</p>
+          )}
+        </>
       )}
     </>
   );
@@ -764,34 +819,254 @@ Flat.args = {
 
 Flat.parameters = { controls: { exclude: ['isInverse'] } };
 
-export const UnitTest = () => {
+export const UnitTest = args => {
+  const [selectedItems, setSelectedItems] = React.useState(null);
+  const [indeterminateItems, setIndeterminateItems] = React.useState(null);
+
+  function onSelection(items) {
+    const selected = items
+      .filter(i => i.checkedStatus === IndeterminateCheckboxStatus.checked)
+      .map((i, key) => (
+        <Tag key={key} size={TagSize.small} color={TagColor.primary}>
+          {i.itemId}
+        </Tag>
+      ));
+
+    const indet = items
+      .filter(
+        i => i.checkedStatus === IndeterminateCheckboxStatus.indeterminate
+      )
+      .map((i, key) => (
+        <Tag key={key} size={TagSize.small} color={TagColor.default}>
+          {i.itemId}
+        </Tag>
+      ));
+
+    setSelectedItems(selected);
+    setIndeterminateItems(indet);
+    console.log('onSelection event:', items);
+  }
+
   return (
-    <TreeView
-      testId={'tree-view'}
-      initialSelectedItems={[
-        { itemId: 'item0', checkedStatus: 'checked' },
-        { itemId: 'item1', checkedStatus: 'checked' },
-        // { itemId: 'item2', checkedStatus: 'checked' },
-        // { itemId: 'item-child2', checkedStatus: 'checked' },
-      ]}
-      initialExpandedItems={['item1', 'item2']}
-      selectable={TreeViewSelectable.multi}
-    >
-      <TreeItem label="Node 0" itemId="item0" testId="item0" />
-      <TreeItem label="Node 1" itemId="item1" testId="item1">
-        <TreeItem label="Child 1" itemId="item-child1" testId="item-child1" />
-      </TreeItem>
-      <TreeItem label="Node 2" itemId="item2" testId="item2">
-        <TreeItem label="Child 2.1" itemId="item-child2" testId="item-child2" />
-        <TreeItem
-          label="Child 2.2"
-          itemId="item-child2.2"
-          testId="item-child2.2"
-        />
-      </TreeItem>
-      <TreeItem label="Node 3" itemId="item3" testId="item3">
-        <TreeItem label="Child 3" itemId="item-child3" testId="item-child3" />
-      </TreeItem>
-    </TreeView>
+    <>
+      {/* <TreeView
+        onSelectedItemChange={onSelection}
+        selectable={TreeViewSelectable.multi}
+        initialExpandedItems={['item2', 'item-child2.1',]}
+        initialSelectedItems={[
+          {
+            itemId: 'item-ggchild3',
+            checkedStatus: IndeterminateCheckboxStatus.checked,
+          },
+          {
+            itemId: 'item-child3',
+            checkedStatus: IndeterminateCheckboxStatus.checked,
+          },
+        ]}
+      >
+        <TreeItem label="Node 0" itemId="item0" testId="item0" />
+        <TreeItem label="Node 1" itemId="item1" testId="item1">
+          <TreeItem label="Child 1" itemId="item-child1" testId="item-child1" />
+        </TreeItem>
+        <TreeItem label="Node 2" itemId="item2" testId="item2">
+          <TreeItem
+            label="Child 2"
+            itemId="item-child2.1"
+            testId="item-child2.1"
+          >
+            <TreeItem
+              label="Grandchild 2"
+              itemId="item-gchild2"
+              testId="item-gchild2"
+            >
+              <TreeItem
+                label="Great-grandchild 2"
+                itemId="item-ggchild2"
+                testId="item-ggchild2"
+              />
+              <TreeItem
+                label="Great-grandchild 3"
+                itemId="item-ggchild3"
+                testId="item-ggchild3"
+              />
+            </TreeItem>
+          </TreeItem>
+        </TreeItem>
+        <TreeItem label="Node 3" itemId="item3" testId="item3">
+          <TreeItem label="Child 3" itemId="item-child3" testId="item-child3" />
+          <TreeItem
+            label="Child 3.1"
+            itemId="item-child3.1"
+            testId="item-child3.1"
+          />
+        </TreeItem>
+      </TreeView> */}
+
+      {/* one level */}
+      <TreeView
+        testId={'tree-view'}
+        onSelectedItemChange={onSelection}
+        initialSelectedItems={[
+          {
+            itemId: 'item2',
+            checkedStatus: IndeterminateCheckboxStatus.indeterminate,
+          },
+          {
+            itemId: 'item-child2.1',
+            checkedStatus: IndeterminateCheckboxStatus.checked,
+          },
+          {
+            itemId: 'item-child3',
+            checkedStatus: IndeterminateCheckboxStatus.checked,
+          },
+        ]}
+        initialExpandedItems={['item2', 'item1']}
+        selectable={TreeViewSelectable.multi}
+      >
+        <TreeItem label="Node 0" itemId="item0" testId="item0" />
+        <TreeItem label="Node 1" itemId="item1" testId="item1">
+          <TreeItem label="Child 1" itemId="item-child1" testId="item-child1" />
+        </TreeItem>
+        <TreeItem label="Node 2" itemId="item2" testId="item2">
+          <TreeItem
+            label="Child 2.1"
+            itemId="item-child2.1"
+            testId="item-child2.1"
+          />
+          <TreeItem
+            label="Child 2.2"
+            itemId="item-child2.2"
+            testId="item-child2.2"
+          />
+        </TreeItem>
+        <TreeItem label="Node 3" itemId="item3" testId="item3">
+          <TreeItem label="Child 3" itemId="item-child3" testId="item-child3" />
+        </TreeItem>
+      </TreeView>
+
+      {/* single item */}
+      {/* <TreeItem label={'labelText'} itemId="parent" testId={'testId'} isDisabled>
+          <TreeItem
+            label={`${'labelText'}-child`}
+            testId={`${'testId'}-child`}
+            itemId="child"
+          />
+        </TreeItem> */}
+
+      <br />
+
+      <p>Selected: {selectedItems}</p>
+      <p>Indeterminate: {indeterminateItems}</p>
+    </>
+  );
+};
+
+export const Animals = () => {
+  const [selectedItems, setSelectedItems] = React.useState(null);
+  const [indeterminateItems, setIndeterminateItems] = React.useState(null);
+  const [total, setTotal] = React.useState(null);
+
+  function onSelection(items) {
+    const selected = items
+      .filter(i => i.checkedStatus === IndeterminateCheckboxStatus.checked)
+      .map((i, key) => (
+        <Tag key={key} size={TagSize.small} color={TagColor.primary}>
+          {i.itemId}
+        </Tag>
+      ));
+
+    const indet = items
+      .filter(
+        i => i.checkedStatus === IndeterminateCheckboxStatus.indeterminate
+      )
+      .map((i, key) => (
+        <Tag key={key} size={TagSize.small} color={TagColor.default}>
+          {i.itemId}
+        </Tag>
+      ));
+
+    setSelectedItems(selected);
+    setIndeterminateItems(indet);
+    setTotal(items.length);
+    console.log('onSelection event:', items);
+  }
+
+  return (
+    <>
+      <TreeView
+        ariaLabel="initial-selected-treeview"
+        selectable={TreeViewSelectable.multi}
+        initialSelectedItems={[
+          {
+            itemId: 'Parrots',
+            checkedStatus: IndeterminateCheckboxStatus.checked,
+          },
+          {
+            itemId: 'African Grey',
+            checkedStatus: IndeterminateCheckboxStatus.checked,
+          },
+          {
+            itemId: 'Cockatiel',
+            checkedStatus: IndeterminateCheckboxStatus.checked,
+          },
+          {
+            itemId: 'Budgerigar',
+            checkedStatus: IndeterminateCheckboxStatus.checked,
+          },
+          {
+            itemId: 'Birds of Prey',
+            checkedStatus: IndeterminateCheckboxStatus.indeterminate,
+          },
+          {
+            itemId: 'Eagles',
+            checkedStatus: IndeterminateCheckboxStatus.checked,
+          },
+          {
+            itemId: 'Falcons',
+            checkedStatus: IndeterminateCheckboxStatus.checked,
+          },
+          //   {
+          //     "itemId": "Hawks",
+          //     "checkedStatus": IndeterminateCheckboxStatus.checked
+          // },
+          {
+            itemId: 'Birds',
+            checkedStatus: IndeterminateCheckboxStatus.indeterminate,
+          },
+        ]}
+        onSelectedItemChange={onSelection}
+      >
+        <TreeItem label="Mammals" itemId="Mammals">
+          <TreeItem label="Dogs" itemId="Dogs">
+            <TreeItem label="German Shepherd" itemId="German Shepherd" />
+            <TreeItem label="Labrador Retriever" itemId="Labrador Retriever" />
+            <TreeItem label="American Bully" itemId="American Bully" />
+          </TreeItem>
+          <TreeItem label="Cats" itemId="Cats">
+            <TreeItem label="Siamese" itemId="Siamese" />
+            <TreeItem label="Persian" itemId="Persian" />
+            <TreeItem label="Bengal" itemId="Bengal" />
+          </TreeItem>
+        </TreeItem>
+        <TreeItem label="Birds" itemId="Birds">
+          <TreeItem label="Parrots" itemId="Parrots">
+            <TreeItem label="African Grey" itemId="African Grey" />
+            <TreeItem label="Cockatiel" itemId="Cockatiel" />
+            <TreeItem label="Budgerigar" itemId="Budgerigar" />
+          </TreeItem>
+          <TreeItem label="Birds of Prey" itemId="Birds of Prey">
+            <TreeItem label="Eagles" itemId="Eagles" />
+            <TreeItem label="Hawks" itemId="Hawks" />
+            <TreeItem label="Falcons" itemId="Falcons" />
+          </TreeItem>
+        </TreeItem>
+        <TreeItem label="Amphibians" itemId="Amphibians" />
+      </TreeView>
+      <>
+        <p>{total} total</p>
+        <p>Selected: {selectedItems}</p>
+        <p>Indeterminate: {indeterminateItems}</p>
+      </>
+    </>
   );
 };
diff --git a/packages/react-magma-dom/src/components/TreeView/TreeView.test.js b/packages/react-magma-dom/src/components/TreeView/TreeView.test.js
index b8762ed62..85b2e5dfe 100644
--- a/packages/react-magma-dom/src/components/TreeView/TreeView.test.js
+++ b/packages/react-magma-dom/src/components/TreeView/TreeView.test.js
@@ -6,6 +6,7 @@ import { magma } from '../../theme/magma';
 import userEvent from '@testing-library/user-event';
 import { FavoriteIcon } from 'react-magma-icons';
 import { transparentize } from 'polished';
+import { IndeterminateCheckboxStatus } from '../IndeterminateCheckbox';
 
 const TEXT = 'Test Text Tree Item';
 const testId = 'tree-view';
@@ -17,7 +18,7 @@ const getTreeItemsOneLevel = props => (
       <TreeItem label="Child 1" itemId="item-child1" testId="item-child1" />
     </TreeItem>
     <TreeItem label="Node 2" itemId="item2" testId="item2">
-      <TreeItem label="Child 2.1" itemId="item-child2" testId="item-child2" />
+      <TreeItem label="Child 2.1" itemId="item-child2.1" testId="item-child2.1" />
       <TreeItem label="Child 2.2" itemId="item-child2.2" testId="item-child2.2" />
     </TreeItem>
     <TreeItem label="Node 3" itemId="item3" testId="item3">
@@ -42,7 +43,7 @@ const getTreeItemsMultiLevel = props => (
       <TreeItem label="Child 1" itemId="item-child1" testId="item-child1" />
     </TreeItem>
     <TreeItem label="Node 2" itemId="item2" testId="item2">
-      <TreeItem label="Child 2" itemId="item-child2" testId="item-child2">
+      <TreeItem label="Child 2" itemId="item-child2.1" testId="item-child2.1">
         <TreeItem
           label="Grandchild 2"
           itemId="item-gchild2"
@@ -73,7 +74,7 @@ const getTreeItemsWithDisabled = props => (
       <TreeItem label="Child 1" itemId="item-child1" testId="item-child1" />
     </TreeItem>
     <TreeItem label="Node 2" itemId="item2" testId="item2" isDisabled>
-      <TreeItem label="Child 2" itemId="item-child2" testId="item-child2" />
+      <TreeItem label="Child 2" itemId="item-child2.1" testId="item-child2.1" />
     </TreeItem>
   </TreeView>
 );
@@ -137,14 +138,14 @@ describe('TreeView', () => {
     it('when child item is part of the array, that item is expanded', () => {
       const { getByTestId } = render(
         getTreeItemsMultiLevel({
-          initialExpandedItems: ['item2', 'item-child2'],
+          initialExpandedItems: ['item2', 'item-child2.1'],
         })
       );
 
       expect(getByTestId('item1')).toHaveAttribute('aria-expanded', 'false');
       expect(getByTestId('item2')).toHaveAttribute('aria-expanded', 'true');
       expect(getByTestId('item3')).toHaveAttribute('aria-expanded', 'false');
-      expect(getByTestId('item-child2')).toHaveAttribute(
+      expect(getByTestId('item-child2.1')).toHaveAttribute(
         'aria-expanded',
         'true'
       );
@@ -217,7 +218,7 @@ describe('TreeView', () => {
         const { getByTestId } = render(
           getTreeItemsOneLevel({
             initialSelectedItems: [
-              { itemId: 'item2', checkedStatus: 'checked' },
+              { itemId: 'item2', checkedStatus: IndeterminateCheckboxStatus.checked },
             ],
             selectable: TreeViewSelectable.off,
           })
@@ -257,7 +258,7 @@ describe('TreeView', () => {
         const { getByTestId } = render(
           getTreeItemsOneLevel({
             initialSelectedItems: [
-              { itemId: 'item2', checkedStatus: 'checked' },
+              { itemId: 'item2', checkedStatus: IndeterminateCheckboxStatus.checked },
             ],
             selectable: TreeViewSelectable.single,
           })
@@ -277,8 +278,8 @@ describe('TreeView', () => {
         const { getByTestId } = render(
           getTreeItemsOneLevel({
             initialSelectedItems: [
-              { itemId: 'item2', checkedStatus: 'checked' },
-              { itemId: 'item0', checkedStatus: 'checked' },
+              {itemId: 'item2', checkedStatus: IndeterminateCheckboxStatus.checked},
+              { itemId: 'item0', checkedStatus: IndeterminateCheckboxStatus.checked },
             ],
             selectable: TreeViewSelectable.single,
           })
@@ -318,7 +319,7 @@ describe('TreeView', () => {
         const { getByTestId } = render(
           getTreeItemsOneLevel({
             initialSelectedItems: [
-              { itemId: 'item2', checkedStatus: 'checked' },
+              {itemId: 'item2', checkedStatus: IndeterminateCheckboxStatus.checked},
             ],
             selectable: TreeViewSelectable.multi,
           })
@@ -338,9 +339,9 @@ describe('TreeView', () => {
         const { getByTestId } = render(
           getTreeItemsOneLevel({
             initialSelectedItems: [
-              { itemId: 'item0', checkedStatus: 'checked' },
-              { itemId: 'item1', checkedStatus: 'checked' },
-              { itemId: 'item2', checkedStatus: 'checked' },
+              { itemId: 'item0', checkedStatus: IndeterminateCheckboxStatus.checked },
+              { itemId: 'item1', checkedStatus: IndeterminateCheckboxStatus.checked },
+              { itemId: 'item2', checkedStatus: IndeterminateCheckboxStatus.checked },
             ],
             initialExpandedItems: ['item1', 'item2'],
             selectable: TreeViewSelectable.multi,
@@ -351,7 +352,7 @@ describe('TreeView', () => {
         expect(getByTestId('item1')).toHaveAttribute('aria-checked', 'true');
         expect(getByTestId('item-child1')).toHaveAttribute('aria-checked', 'true');
         expect(getByTestId('item2')).toHaveAttribute('aria-checked', 'true');
-        expect(getByTestId('item-child2')).toHaveAttribute('aria-checked', 'true');
+        expect(getByTestId('item-child2.1')).toHaveAttribute('aria-checked', 'true');
         expect(getByTestId('item3')).toHaveAttribute('aria-checked', 'false');
 
         expect(getByTestId('item0')).not.toHaveAttribute('aria-selected');
@@ -359,6 +360,28 @@ describe('TreeView', () => {
         expect(getByTestId('item2')).not.toHaveAttribute('aria-selected');
         expect(getByTestId('item3')).not.toHaveAttribute('aria-selected');
       });
+
+      it('and initialSelectedItems is set to multiple items at different depths, all those TreeItems are selected', () => {
+        const { getByTestId } = render(
+          getTreeItemsOneLevel({
+            initialSelectedItems: [
+              { itemId: 'item2', checkedStatus: IndeterminateCheckboxStatus.indeterminate },
+              { itemId: 'item-child2.1', checkedStatus: IndeterminateCheckboxStatus.checked },
+              { itemId: 'item-child3', checkedStatus: IndeterminateCheckboxStatus.checked },
+            ],
+            initialExpandedItems: ['item1', 'item2'],
+            selectable: TreeViewSelectable.multi,
+          })
+        );
+
+        expect(getByTestId('item2')).toHaveAttribute('aria-checked', 'false');
+        expect(getByTestId('item-child2.1')).toHaveAttribute('aria-checked', 'true');
+        expect(getByTestId('item-child2.2')).toHaveAttribute('aria-checked', 'false');
+        expect(getByTestId('item3')).toHaveAttribute('aria-checked', 'true');
+
+        userEvent.click(getByTestId('item3-expand'));
+        expect(getByTestId('item-child3')).toHaveAttribute('aria-checked', 'true');
+      });
     });
   });
 
@@ -426,7 +449,7 @@ describe('TreeView', () => {
 
         userEvent.click(getByTestId('item0-label'));
         expect(onSelectedItemChange).toHaveBeenCalledWith([
-          { itemId: 'item0', checkedStatus: 'checked' },
+          { itemId: 'item0', checkedStatus: IndeterminateCheckboxStatus.checked },
         ]);
       });
 
@@ -441,7 +464,7 @@ describe('TreeView', () => {
 
         userEvent.click(getByTestId('item2-label'));
         expect(onSelectedItemChange).toHaveBeenCalledWith([
-          { itemId: 'item2', checkedStatus: 'checked' },
+          { itemId: 'item2', checkedStatus: IndeterminateCheckboxStatus.checked },
         ]);
       });
 
@@ -490,7 +513,7 @@ describe('TreeView', () => {
 
         userEvent.click(getByTestId('item0-checkbox'));
         expect(onSelectedItemChange).toHaveBeenCalledWith([
-          { itemId: 'item0', checkedStatus: 'checked' },
+          { itemId: 'item0', checkedStatus: IndeterminateCheckboxStatus.checked },
         ]);
       });
 
@@ -505,11 +528,22 @@ describe('TreeView', () => {
 
         userEvent.click(getByTestId('item2-checkbox'));
         expect(onSelectedItemChange).toHaveBeenCalledWith([
-          { itemId: 'item-child2', checkedStatus: 'checked' },
-          { itemId: 'item-gchild2', checkedStatus: 'checked' },
-          { itemId: 'item-ggchild2', checkedStatus: 'checked' },
-          { itemId: 'item-ggchild3', checkedStatus: 'checked' },
-          { itemId: 'item2', checkedStatus: 'checked' },
+          { itemId: 'item-child2.1', checkedStatus: IndeterminateCheckboxStatus.checked },
+          { itemId: 'item-gchild2', checkedStatus: IndeterminateCheckboxStatus.checked },
+          { itemId: 'item-ggchild2', checkedStatus: IndeterminateCheckboxStatus.checked },
+          { itemId: 'item-ggchild3', checkedStatus: IndeterminateCheckboxStatus.checked },
+          { itemId: 'item2', checkedStatus: IndeterminateCheckboxStatus.checked },
+        ]);
+        userEvent.click(getByTestId('item2-expand'));
+        userEvent.click(getByTestId('item-child2.1-expand'));
+        userEvent.click(getByTestId('item-gchild2-expand'));
+
+        userEvent.click(getByTestId('item-ggchild2-checkbox'));
+        expect(onSelectedItemChange).toHaveBeenCalledWith([
+          { itemId: 'item-ggchild3', checkedStatus: IndeterminateCheckboxStatus.checked },
+          { itemId: 'item-gchild2', checkedStatus: IndeterminateCheckboxStatus.indeterminate },
+          { itemId: 'item-child2.1', checkedStatus: IndeterminateCheckboxStatus.indeterminate },
+          { itemId: 'item2', checkedStatus: IndeterminateCheckboxStatus.indeterminate },
         ]);
       });
 
@@ -529,6 +563,59 @@ describe('TreeView', () => {
     });
   });
 
+  describe('initialExpandedItems and initialSelectedItems', () => {
+    it('when initialExpandedItems and initialSelectedItems are empty, no TreeItem is expanded or selected', () => {
+      const { getByTestId } = render(
+        getTreeItemsOneLevel({  selectable: TreeViewSelectable.multi, initialExpandedItems: [], initialSelectedItems: []})
+      );
+
+      expect(getByTestId('item1')).toHaveAttribute('aria-expanded', 'false');
+      expect(getByTestId('item2')).toHaveAttribute('aria-expanded', 'false');
+      expect(getByTestId('item3')).toHaveAttribute('aria-expanded', 'false');
+
+      expect(getByTestId('item0')).toHaveAttribute('aria-checked', 'false');
+      expect(getByTestId('item1')).toHaveAttribute('aria-checked', 'false');
+      expect(getByTestId('item2')).toHaveAttribute('aria-checked', 'false');
+      expect(getByTestId('item3')).toHaveAttribute('aria-checked', 'false');
+    });
+
+    it('when initialExpandedItems is set and initialSelectedItems is set, the items are expanded and selected', () => {
+      const { getByTestId } = render(
+        getTreeItemsOneLevel({
+          selectable: TreeViewSelectable.multi,
+          initialExpandedItems: ['item2', 'item1'],
+          initialSelectedItems: [
+            {
+              itemId: 'item2',
+              checkedStatus: IndeterminateCheckboxStatus.indeterminate,
+            },
+            {
+              itemId: 'item-child2.1',
+              checkedStatus: IndeterminateCheckboxStatus.checked,
+            },
+            {
+              itemId: 'item-child3',
+              checkedStatus: IndeterminateCheckboxStatus.checked,
+            },
+          ],
+        })
+      );
+
+      expect(getByTestId('item1')).toHaveAttribute('aria-expanded', 'true');
+      expect(getByTestId('item2')).toHaveAttribute('aria-expanded', 'true');
+      expect(getByTestId('item3')).toHaveAttribute('aria-expanded', 'false');
+
+      expect(getByTestId('item0')).toHaveAttribute('aria-checked', 'false');
+      expect(getByTestId('item1')).toHaveAttribute('aria-checked', 'false');
+      expect(getByTestId('item2')).toHaveAttribute('aria-checked', 'false');
+      expect(getByTestId('item-child2.1')).toHaveAttribute('aria-checked', 'true');
+      expect(getByTestId('item3')).toHaveAttribute('aria-checked', 'true');
+
+      userEvent.click(getByTestId('item3-expand'));
+      expect(getByTestId('item-child3')).toHaveAttribute('aria-checked', 'true');
+    });
+  });
+
   describe('a11y', () => {
     it('sets the ariaLabel', () => {
       const testId = 'ariaLabelId';
diff --git a/packages/react-magma-dom/src/components/TreeView/TreeViewContext.ts b/packages/react-magma-dom/src/components/TreeView/TreeViewContext.ts
index 115480457..956c0e588 100644
--- a/packages/react-magma-dom/src/components/TreeView/TreeViewContext.ts
+++ b/packages/react-magma-dom/src/components/TreeView/TreeViewContext.ts
@@ -4,13 +4,15 @@ import { IndeterminateCheckboxStatus } from '../IndeterminateCheckbox';
 
 export interface TreeItemSelectedInterface {
   itemId?: string;
-  checkedStatus: IndeterminateCheckboxStatus
-};
+  checkedStatus: IndeterminateCheckboxStatus;
+}
 
 export interface TreeViewContextInterface {
   children?: React.ReactNode | React.ReactNode[];
   hasIcons: boolean;
-  onSelectedItemChange?: (selectedItems: Array<TreeItemSelectedInterface>) => void;
+  onSelectedItemChange?: (
+    selectedItems: Array<TreeItemSelectedInterface>
+  ) => void;
   onExpandedChange?: (event: React.SyntheticEvent) => void;
   selectable: TreeViewSelectable;
   setHasIcons: React.Dispatch<React.SetStateAction<boolean>>;
@@ -24,7 +26,9 @@ export interface TreeViewContextInterface {
     itemRef: React.MutableRefObject<Element>
   ) => void;
   initialSelectedItemsNeedUpdate: boolean;
-    setInitialSelectedItemsNeedUpdate: React.Dispatch<React.SetStateAction<any>>;
+  setInitialSelectedItemsNeedUpdate: React.Dispatch<React.SetStateAction<any>>;
+  initialExpandedItemsNeedUpdate: boolean;
+  setInitialExpandedItemsNeedUpdate: React.Dispatch<React.SetStateAction<any>>;
 }
 
 export const TreeViewContext = React.createContext<TreeViewContextInterface>({
@@ -38,4 +42,6 @@ export const TreeViewContext = React.createContext<TreeViewContextInterface>({
   registerTreeItem: (elements, element) => {},
   initialSelectedItemsNeedUpdate: false,
   setInitialSelectedItemsNeedUpdate: () => {},
+  initialExpandedItemsNeedUpdate: false,
+  setInitialExpandedItemsNeedUpdate: () => {},
 });
diff --git a/packages/react-magma-dom/src/components/TreeView/useTreeItem.ts b/packages/react-magma-dom/src/components/TreeView/useTreeItem.ts
index 84f68e547..54787a469 100644
--- a/packages/react-magma-dom/src/components/TreeView/useTreeItem.ts
+++ b/packages/react-magma-dom/src/components/TreeView/useTreeItem.ts
@@ -14,13 +14,14 @@ import {
   getChildrenCheckedStatus,
   // getEnabledTreeItemChildrenLength,
   getUniqueSelectedItemsArray,
-  selectedItemsIncludesId,
+  arrayIncludesId,
   getUpdatedSelectedItems,
   findCommonItems,
   areArraysEqual,
   findChildByItemId,
   getChildrenItemIdsInTree,
   getAllParentIds,
+  getChildrenItemIdsFlat,
 } from './utils';
 
 export interface UseTreeItemProps extends React.HTMLAttributes<HTMLLIElement> {
@@ -121,9 +122,10 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
     treeItemRefArray,
     initialSelectedItemsNeedUpdate,
     setInitialSelectedItemsNeedUpdate,
+    initialExpandedItemsNeedUpdate,
   } = React.useContext(TreeViewContext);
 
-  const [expanded, setExpanded] = React.useState(isDisabled);
+  const [expanded, setExpanded] = React.useState(false);
   const [checkedStatus, setCheckedStatus] =
     React.useState<IndeterminateCheckboxStatus>(
       IndeterminateCheckboxStatus.unchecked
@@ -165,8 +167,8 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
           updateParentCheckStatus(index, newStatus);
           setSelectedItems(prev => {
             return getUniqueSelectedItemsArray(
-              prev,
               [{ itemId: itemId, checkedStatus: newStatus }],
+              prev,
               []
             );
           });
@@ -196,22 +198,18 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
     });
   }
 
-  React.useEffect(() => {
-    if (isDisabled || initialExpandedItems?.length === 0) {
-      setExpanded(false);
-    } else if (initialExpandedItems?.includes(itemId)) {
-      setExpanded(true);
-    } else {
-      setExpanded(false);
-    }
-  }, [initialExpandedItems]);
-
   React.useEffect(() => {
     if (initialSelectedItemsNeedUpdate) {
       updateInitialSelected();
     }
   }, [initialSelectedItemsNeedUpdate]);
 
+  React.useEffect(() => {
+    if (initialExpandedItemsNeedUpdate) {
+      updateInitialExpanded();
+    }
+  }, [initialExpandedItemsNeedUpdate]);
+
   const updateCheckedStatusFromChild = (
     index: number,
     status: IndeterminateCheckboxStatus
@@ -230,10 +228,13 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
     itemId,
     itemIdChildren
   ) => {
-    const item = initialSelectedChildrenItems.find(child => child.itemId === itemId);
+    const item = initialSelectedChildrenItems.find(
+      child => child.itemId === itemId
+    );
+
     const parentStatus =
       item?.checkedStatus ||
-      areArraysEqual(childrenItemIds, initialSelectedChildrenItems)
+      areArraysEqual(initialSelectedChildrenItems, childrenItemIds)
         ? IndeterminateCheckboxStatus.checked
         : IndeterminateCheckboxStatus.indeterminate;
 
@@ -243,31 +244,51 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
 
     if (!item) {
       setSelectedItems(prev => {
-        return getUniqueSelectedItemsArray(prev, initialSelectedChildrenItems, [
-          { itemId: itemId, checkedStatus: parentStatus },
-        ]);
+        return getUniqueSelectedItemsArray(
+          [{ itemId: itemId, checkedStatus: parentStatus }],
+          initialSelectedChildrenItems,
+          prev
+        );
       });
-      return 
+      return;
     }
 
-    const itemThing = itemIdChildren.find(child => child.itemId === itemId);
+    const thisItem = itemIdChildren.find(child => child.itemId === itemId);
 
     if (
-      itemThing?.children.length > 0 &&
+      thisItem?.children.length > 0 &&
       item?.checkedStatus === IndeterminateCheckboxStatus.checked
     ) {
-      // TODO cleanup
-      const itemNode = findChildByItemId(treeItemChildren, itemThing?.itemId);
-      const newChildren = getChildrenItemIds(itemNode?.props.children);
+      const itemNode = findChildByItemId(treeItemChildren, thisItem?.itemId);
+      const newChildren = getChildrenItemIds(
+        itemNode?.props.children,
+        checkedStatus
+      );
 
       setSelectedItems(prev => {
-        return getUniqueSelectedItemsArray(prev, newChildren, [
-          { itemId: itemId, checkedStatus: parentStatus },
-        ]);
+        return getUniqueSelectedItemsArray(
+          [{ itemId: itemId, checkedStatus: parentStatus }],
+          newChildren,
+          prev
+        );
       });
     }
   };
 
+  const updateInitialExpanded = () => {
+    if (initialExpandedItems?.length !== 0 && !isDisabled) {
+      const childrenItemIds = getChildrenItemIdsFlat(treeItemChildren);
+      const allExpanded = [...initialExpandedItems, ...childrenItemIds];
+      if (allExpanded?.some(item => item === itemId)) {
+        setExpanded(true);
+      } else {
+        setExpanded(false);
+      }
+    } else {
+      setExpanded(false);
+    }
+  };
+
   const updateInitialSelected = () => {
     if (selectable === TreeViewSelectable.single && initialSelectedItems) {
       if (
@@ -285,17 +306,19 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
     ) {
       const item = initialSelectedItems.find(obj => obj.itemId === itemId);
       const status = item?.checkedStatus;
-
-      const childrenItemIds = getChildrenItemIds(treeItemChildren);
+      const childrenItemIds = getChildrenItemIds(
+        treeItemChildren,
+        status || IndeterminateCheckboxStatus.checked
+      );
       // Items from initialSelectedItems that are children
       const initialSelectedChildrenItems = findCommonItems(
-        initialSelectedItems,
-        childrenItemIds
+        childrenItemIds,
+        initialSelectedItems
       );
 
       if (
         !isDisabled &&
-        (selectedItemsIncludesId(initialSelectedItems, itemId) ||
+        (arrayIncludesId(initialSelectedItems, itemId) ||
           childrenItemIds?.includes(itemId))
       ) {
         setStatusUpdatedBy(StatusUpdatedByOptions.checkboxChange);
@@ -314,9 +337,9 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
         }
         setSelectedItems(prev => {
           const allItems = getUniqueSelectedItemsArray(
-            prev,
+            childrenItemIds,
             initialSelectedItems,
-            childrenItemIds
+            prev
           );
           return allItems;
         });
@@ -325,9 +348,10 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
         const itemIdChildren = getChildrenItemIdsInTree(treeItemChildren);
 
         for (const i of initialSelectedChildrenItems) {
-          const itemIdNode = findChildByItemId(treeItemChildren, i.itemId);
+          const itemIdNode = findChildByItemId(treeItemChildren, i.itemId);          
           const childrenOfItemId = getChildrenItemIds(
-            itemIdNode?.props?.children
+            itemIdNode?.props?.children,
+            status
           );
 
           const parentIds = getAllParentIds(itemIdChildren, i.itemId);
@@ -378,6 +402,7 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
           );
         } else {
           const childrenIds = getChildrenItemIds(treeItemChildren);
+
           const newChildrenCheckedStatus = getChildrenCheckedStatus(
             childrenIds,
             parentCheckedStatus
@@ -412,11 +437,13 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
           statusFromChildren === IndeterminateCheckboxStatus.checked ||
           statusFromChildren === IndeterminateCheckboxStatus.indeterminate
         ) {
-          if (itemId && !selectedItemsIncludesId(selectedItems, itemId)) {
+          if (itemId && !arrayIncludesId(selectedItems, itemId)) {
             setSelectedItems([
               ...selectedItems,
               { itemId, checkedStatus: statusFromChildren },
             ]);
+          } else {
+            setSelectedItems(updateItemStatus);
           }
         } else if (
           statusFromChildren === IndeterminateCheckboxStatus.unchecked
@@ -428,7 +455,7 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
         statusUpdatedBy !== StatusUpdatedByOptions.parent &&
         statusFromChildren === IndeterminateCheckboxStatus.indeterminate
       ) {
-        if (!selectedItemsIncludesId(selectedItems, itemId)) {
+        if (!arrayIncludesId(selectedItems, itemId)) {
           setSelectedItems([
             ...selectedItems,
             { itemId, checkedStatus: statusFromChildren },
@@ -459,7 +486,8 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
             Array(numberOfTreeItemChildren).fill(status)
           );
         } else {
-          const childrenIds = getChildrenItemIds(treeItemChildren);
+          const childrenIds = getChildrenItemIds(treeItemChildren, 'something');
+
           const newChildrenCheckedStatus = getChildrenCheckedStatus(
             childrenIds,
             status
@@ -475,7 +503,7 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
     event: React.ChangeEvent<HTMLInputElement>,
     itemId: any
   ): void => {
-    if (!selectedItemsIncludesId(selectedItems, itemId)) {
+    if (!arrayIncludesId(selectedItems, itemId)) {
       setSelectedItems([
         { itemId, checkedStatus: IndeterminateCheckboxStatus.checked },
       ]);
@@ -486,11 +514,14 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
     event: React.ChangeEvent<HTMLInputElement>
   ): void => {
     if (hasOwnTreeItems) {
-      const childrenIds = getChildrenItemIds(treeItemChildren);
+      const childrenIds = getChildrenItemIds(
+        treeItemChildren,
+        IndeterminateCheckboxStatus.checked
+      );
       if (event.target.checked) {
         updateParentCheckStatus(index, IndeterminateCheckboxStatus.checked);
 
-        if (!selectedItemsIncludesId(selectedItems, itemId)) {
+        if (!arrayIncludesId(selectedItems, itemId)) {
           setSelectedItems([
             ...selectedItems,
             ...childrenIds,
@@ -514,7 +545,7 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
       }
     } else {
       if (event.target.checked) {
-        if (!selectedItemsIncludesId(selectedItems, itemId)) {
+        if (!arrayIncludesId(selectedItems, itemId)) {
           setSelectedItems([
             ...selectedItems,
             { itemId, checkedStatus: IndeterminateCheckboxStatus.checked },
@@ -607,7 +638,7 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
   };
 
   const toggleMultiSelectItems = () => {
-    const status = selectedItemsIncludesId(selectedItems, itemId)
+    const status = arrayIncludesId(selectedItems, itemId)
       ? IndeterminateCheckboxStatus.unchecked
       : IndeterminateCheckboxStatus.checked;
     setStatusUpdatedBy(StatusUpdatedByOptions.checkboxChange);
@@ -615,8 +646,8 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
     updateParentCheckStatus(index, status);
 
     if (hasOwnTreeItems) {
-      const childrenIds = getChildrenItemIds(treeItemChildren);
-      if (!selectedItemsIncludesId(selectedItems, itemId)) {
+      const childrenIds = getChildrenItemIds(treeItemChildren, status);
+      if (!arrayIncludesId(selectedItems, itemId)) {
         setSelectedItems([
           ...selectedItems,
           ...childrenIds,
@@ -631,7 +662,7 @@ export function useTreeItem(props: UseTreeItemProps, forwardedRef) {
         setSelectedItems(newSelectedItems);
       }
     } else {
-      if (!selectedItemsIncludesId(selectedItems, itemId)) {
+      if (!arrayIncludesId(selectedItems, itemId)) {
         setSelectedItems([...selectedItems, { itemId, checkedStatus: status }]);
       } else {
         setSelectedItems(selectedItems.filter(obj => obj.itemId !== itemId));
diff --git a/packages/react-magma-dom/src/components/TreeView/useTreeView.ts b/packages/react-magma-dom/src/components/TreeView/useTreeView.ts
index 4c9aec636..bef343a5c 100644
--- a/packages/react-magma-dom/src/components/TreeView/useTreeView.ts
+++ b/packages/react-magma-dom/src/components/TreeView/useTreeView.ts
@@ -66,6 +66,8 @@ export function useTreeView(props: UseTreeViewProps) {
   const [selectedItems, setSelectedItems] = React.useState([]);
   const [initialSelectedItemsNeedUpdate, setInitialSelectedItemsNeedUpdate] =
     React.useState(false);
+  const [initialExpandedItemsNeedUpdate, setInitialExpandedItemsNeedUpdate] =
+    React.useState(false);
 
   const [treeItemRefArray, registerTreeItem] = useDescendants();
 
@@ -82,6 +84,9 @@ export function useTreeView(props: UseTreeViewProps) {
     if (selectable !== TreeViewSelectable.off && initialSelectedItems) {
       setInitialSelectedItemsNeedUpdate(true);
     }
+    if (initialExpandedItems) {
+      setInitialExpandedItemsNeedUpdate(true)
+    }
   }, []);
 
   const contextValue = {
@@ -98,6 +103,8 @@ export function useTreeView(props: UseTreeViewProps) {
     registerTreeItem,
     initialSelectedItemsNeedUpdate,
     setInitialSelectedItemsNeedUpdate,
+    initialExpandedItemsNeedUpdate,
+    setInitialExpandedItemsNeedUpdate,
   };
 
   return { contextValue };
diff --git a/packages/react-magma-dom/src/components/TreeView/utils.ts b/packages/react-magma-dom/src/components/TreeView/utils.ts
index 14d95ccec..54d449b98 100644
--- a/packages/react-magma-dom/src/components/TreeView/utils.ts
+++ b/packages/react-magma-dom/src/components/TreeView/utils.ts
@@ -92,26 +92,34 @@ export function getTreeItemWrapperCursor(
   return 'default';
 }
 
-// Returns boolean if selectedItems has itemId
-export function selectedItemsIncludesId(selectedItems, itemId) {
-  return selectedItems?.some(item => item.itemId === itemId);
+// Returns boolean if itemsArray has itemId
+export function arrayIncludesId(itemsArray, itemId) {
+  return itemsArray?.some(item => item.itemId === itemId);
 }
 
-// Return an array of all the enabled children ids recursively
-export function getChildrenItemIds(children) {
+// Return an array of objects of all the enabled children ids recursively
+export function getChildrenItemIds(children, status = '') {
   let itemIds = [];
 
   React.Children.forEach(children, child => {
     if (!child.props?.isDisabled) {
+      const childStatus =
+        status === IndeterminateCheckboxStatus.checked
+          ? IndeterminateCheckboxStatus.checked
+          : IndeterminateCheckboxStatus.unchecked;
+
       if (child.props?.itemId) {
         itemIds.push({
           itemId: child.props.itemId,
-          checkedStatus: IndeterminateCheckboxStatus.checked,
+          checkedStatus: childStatus,
         });
       }
 
       if (child.props?.children) {
-        const nestedItemIds = getChildrenItemIds(child.props.children);
+        const nestedItemIds = getChildrenItemIds(
+          child.props.children,
+          childStatus
+        );
         itemIds = itemIds.concat(nestedItemIds);
       }
     }
@@ -120,6 +128,27 @@ export function getChildrenItemIds(children) {
   return itemIds;
 }
 
+// Return an array of strings of all enabled children ids recursively
+export function getChildrenItemIdsFlat(children) {
+  let itemIds = [];
+
+  React.Children.forEach(children, child => {
+    if (!child.props?.isDisabled) {
+      if (child.props?.itemId) {
+        itemIds.push(child.props.itemId);
+      }
+
+      if (child.props?.children) {
+        const nestedItemIds = getChildrenItemIdsFlat(child.props.children);
+        itemIds = itemIds.concat(nestedItemIds);
+      }
+    }
+  });
+
+  return itemIds;
+}
+
+// Return an array of objects where all children are items are nested in the parents
 export function getChildrenItemIdsInTree(children) {
   let itemIds = [];
 
@@ -169,7 +198,8 @@ export function findChildByItemId(children, itemId) {
     }
 
     if (child?.props?.children) {
-      const nestedChild = findChildByItemId(child?.props?.children, itemId);
+      const nestedChild = findChildByItemId([child?.props?.children], itemId);
+
       if (nestedChild) {
         return nestedChild;
       }
@@ -200,7 +230,9 @@ export function getMissingChildrenIds(selectedItems, childrenIds) {
 
 // Return an array of statuses for all enabled children
 export function getChildrenCheckedStatus(childrenIds, status) {
-  return childrenIds.map(child => (child.isDisabled ? IndeterminateCheckboxStatus.unchecked : status));
+  return childrenIds.map(child =>
+    child.isDisabled ? IndeterminateCheckboxStatus.unchecked : status
+  );
 }
 
 // Return the length of enabled children
@@ -226,12 +258,8 @@ export function getUpdatedSelectedItems(selectedItems, itemId, checkedStatus) {
 }
 
 // Return an array of unique items from the previous state, initially selected items and the childrem item ids
-export function getUniqueSelectedItemsArray(
-  prev,
-  initialSelectedItems,
-  childrenItemIds
-) {
-  const combinedArray = [...prev, ...initialSelectedItems, ...childrenItemIds];
+export function getUniqueSelectedItemsArray(itemArr0, itemArr1, itemArr2) {
+  const combinedArray = [...itemArr0, ...itemArr2, ...itemArr1];
   const uniqueItemsMap = new Map();
   for (const item of combinedArray) {
     uniqueItemsMap.set(item.itemId, item);