Lazy loaded tree children should inherit selected parent state

.changeset/legal-symbols-build.md

@@ -0,0 +1,5 @@
+---
+'@sl-design-system/tree': patch
+---
+
+Fix lazy loaded children to inherit the selected state of the parent node

packages/components/tree/src/flat-tree-data-source.spec.ts

@@ -1,6 +1,16 @@
+import { spy } from 'sinon';
 import { beforeEach, describe, expect, it } from 'vitest';
 import { FlatTreeDataSource } from './flat-tree-data-source.js';
 
+type TestItem = {
+  id: number | string;
+  name: string;
+  level: number;
+  expandable: boolean;
+  expanded?: boolean;
+  selected?: boolean;
+};
+
 describe('FlatTreeDataSource', () => {
   let ds: FlatTreeDataSource;
 
@@ -280,4 +290,190 @@ describe('FlatTreeDataSource', () => {
       expect(ds.items.map(n => n.label)).to.deep.equal(['C', 'Z', 'Y', 'B', 'A']);
     });
   });
+
+  describe('lazy loading children', () => {
+    let loadChildrenSpy: ReturnType<typeof spy>;
+
+    beforeEach(() => {
+      loadChildrenSpy = spy((node: TestItem) => {
+        if (node.id === 1) {
+          return Promise.resolve([
+            { id: 11, name: 'Child 1.1', level: 1, expandable: false },
+            { id: 12, name: 'Child 1.2', level: 1, expandable: false }
+          ]);
+        } else if (node.id === 3) {
+          return Promise.resolve([{ id: 31, name: 'Child 3.1', level: 1, expandable: false }]);
+        }
+
+        return Promise.resolve([]);
+      });
+
+      ds = new FlatTreeDataSource<TestItem>(
+        [
+          { id: 1, name: 'Parent 1', level: 0, expandable: true },
+          { id: 2, name: 'Leaf node at level 0', level: 0, expandable: false },
+          { id: 3, name: 'Parent 3', level: 0, expandable: true }
+        ],
+        {
+          getId: ({ id }) => id,
+          getLabel: ({ name }) => name,
+          getLevel: ({ level }) => level,
+          isExpandable: ({ expandable }) => expandable,
+          isExpanded: ({ expanded }) => expanded ?? false,
+          isSelected: ({ selected }) => selected ?? false,
+          loadChildren: loadChildrenSpy as (node: TestItem) => Promise<TestItem[]>,
+          multiple: true
+        }
+      );
+      ds.update();
+    });
+
+    it('should call loadChildren when expanding a node without children', () => {
+      const parentNode = ds.items[0];
+
+      expect(parentNode.expanded).to.be.false;
+      expect(parentNode.children).to.be.undefined;
+      expect(loadChildrenSpy).not.to.have.been.called;
+
+      ds.expand(parentNode);
+
+      expect(loadChildrenSpy).to.have.been.calledOnce;
+      expect(loadChildrenSpy).to.have.been.calledWith(parentNode.dataNode);
+    });
+
+    it('should set childrenLoading promise when expanding a node', () => {
+      const parentNode = ds.items[0];
+
+      ds.expand(parentNode);
+
+      expect(parentNode.childrenLoading).to.be.instanceof(Promise);
+    });
+
+    it('should load and display children after expansion', async () => {
+      const parentNode = ds.items[0];
+
+      expect(ds.items).to.have.length(3); // All 3 nodes at level 0
+
+      ds.expand(parentNode);
+      await parentNode.childrenLoading;
+
+      expect(parentNode.children).to.have.length(2);
+      expect(parentNode.children![0].label).to.equal('Child 1.1');
+      expect(parentNode.children![1].label).to.equal('Child 1.2');
+      expect(parentNode.childrenLoading).to.be.undefined;
+      expect(ds.items).to.have.length(5); // Now includes the 2 loaded children
+    });
+
+    it('should not call loadChildren again if children are already loaded', async () => {
+      const parentNode = ds.items[0];
+
+      // First expansion
+      ds.expand(parentNode);
+      await parentNode.childrenLoading;
+
+      expect(loadChildrenSpy).to.have.been.calledOnce;
+
+      // Collapse and expand again
+      ds.collapse(parentNode);
+      ds.expand(parentNode);
+
+      expect(loadChildrenSpy).to.have.been.calledOnce; // Should not be called again
+    });
+
+    it('should handle multiple lazy loading operations simultaneously', async () => {
+      const parentNode1 = ds.items[0]; // Parent 1
+      const parentNode3 = ds.items[2]; // Parent 3
+
+      ds.expand(parentNode1);
+      ds.expand(parentNode3);
+
+      await Promise.all([parentNode1.childrenLoading, parentNode3.childrenLoading]);
+
+      expect(loadChildrenSpy).to.have.been.calledTwice;
+      expect(parentNode1.children).to.have.length(2);
+      expect(parentNode3.children).to.have.length(1);
+      expect(ds.items).to.have.length(6); // 3 original + 2 from parent1 + 1 from parent3
+    });
+
+    it('should handle expandAll with lazy loading', async () => {
+      expect(ds.items).to.have.length(3);
+
+      await ds.expandAll();
+
+      expect(loadChildrenSpy).to.have.been.calledTwice;
+      expect(ds.items).to.have.length(6);
+      expect(ds.items.filter(item => item.expandable).every(item => item.expanded)).to.be.true;
+    });
+
+    it('should maintain correct tree structure after lazy loading', async () => {
+      const parentNode = ds.items[0];
+
+      ds.expand(parentNode);
+      await parentNode.childrenLoading;
+
+      const children = parentNode.children!;
+      expect(children[0].parent).to.equal(parentNode);
+      expect(children[1].parent).to.equal(parentNode);
+      expect(children[0].level).to.equal(1);
+      expect(children[1].level).to.equal(1);
+      expect(children[0].lastNodeInLevel).to.be.false;
+      expect(children[1].lastNodeInLevel).to.be.true;
+    });
+
+    it('should handle lazy loading errors gracefully', async () => {
+      const errorSpy = spy(() => Promise.reject(new Error('Failed to load children')));
+
+      const errorDs = new FlatTreeDataSource<TestItem>([{ id: 1, name: 'Parent 1', level: 0, expandable: true }], {
+        getId: ({ id }) => id,
+        getLabel: ({ name }) => name,
+        getLevel: ({ level }) => level,
+        isExpandable: ({ expandable }) => expandable,
+        isExpanded: ({ expanded }) => expanded ?? false,
+        loadChildren: errorSpy as (node: TestItem) => Promise<TestItem[]>
+      });
+      errorDs.update();
+
+      const parentNode = errorDs.items[0];
+
+      errorDs.expand(parentNode);
+
+      try {
+        await parentNode.childrenLoading;
+      } catch (error) {
+        expect(error).to.be.instanceof(Error);
+        expect((error as Error).message).to.equal('Failed to load children');
+      }
+
+      expect(parentNode.children).to.be.undefined;
+      expect(parentNode.childrenLoading).to.be.instanceof(Promise);
+    });
+
+    it('should select all lazy loaded children when parent is already selected', async () => {
+      const parentNode = ds.items[0];
+
+      // Select the parent node before expanding
+      ds.selection.add(parentNode);
+      parentNode.selected = true;
+
+      expect(parentNode.selected).to.be.true;
+      expect(parentNode.children).to.be.undefined;
+
+      // Expand the parent to load children
+      ds.expand(parentNode);
+      await parentNode.childrenLoading;
+
+      // Check that children exist
+      const children = parentNode.children!;
+      expect(children).to.have.length(2);
+
+      // When a selected parent is expanded and lazy-loads children, all children
+      // should automatically be selected to maintain consistent selection state
+      expect(children[0].selected).to.be.true; // Child 1.1 should be selected
+      expect(children[1].selected).to.be.true; // Child 1.2 should also be selected
+
+      // All children should be in the selection
+      expect(ds.selection.has(children[0])).to.be.true;
+      expect(ds.selection.has(children[1])).to.be.true;
+    });
+  });
 });

packages/components/tree/src/flat-tree-data-source.ts

@@ -50,7 +50,18 @@ export class FlatTreeDataSource<T = any> extends TreeDataSource<T> {
       loadChildren = async (node: TreeDataSourceNode<T>) => {
         const children = await options.loadChildren!(node.dataNode);
 
-        return children.map((child, index) => this.#mapToTreeNode(child, node, index === children.length - 1));
+        return children.map((child, index) => {
+          const childNode = this.#mapToTreeNode(child, node, index === children.length - 1);
+
+          // If the parent is selected and we have multiple selection enabled,
+          // ensure all lazy-loaded children are also selected
+          if (this.multiple && node.selected) {
+            childNode.selected = true;
+            this.selection.add(childNode);
+          }
+
+          return childNode;
+        });
       };
     }