[tree view] Fix collapsed children not selected

packages/x-tree-view/src/internals/TreeViewProvider/TreeViewChildrenItemProvider.tsx

@@ -20,6 +20,11 @@ export function TreeViewChildrenItemProvider(props: TreeViewChildrenItemProvider
 
   const { store, rootRef } = useTreeViewContext<SimpleTreeViewStore<any>>();
   const childrenIdAttrToIdRef = React.useRef<Map<string, string>>(new Map());
+  // Re-run the children-discovery effect below whenever a child registers/unregisters, so it
+  // picks up children once they are mounted in the DOM. This is required because some transition
+  // components (for example `@mui/material`'s `Collapse`) mount the children in a deferred render
+  // that does not re-render this provider on its own.
+  const [, rerender] = React.useReducer((s: number) => s + 1, 0);
 
   React.useEffect(() => {
     if (!rootRef.current) {
@@ -57,9 +62,14 @@ export function TreeViewChildrenItemProvider(props: TreeViewChildrenItemProvider
 
   const value = React.useMemo<TreeViewChildrenItemContextValue>(
     () => ({
-      registerChild: (childIdAttribute, childItemId) =>
-        childrenIdAttrToIdRef.current.set(childIdAttribute, childItemId),
-      unregisterChild: (childIdAttribute) => childrenIdAttrToIdRef.current.delete(childIdAttribute),
+      registerChild: (childIdAttribute, childItemId) => {
+        childrenIdAttrToIdRef.current.set(childIdAttribute, childItemId);
+        rerender();
+      },
+      unregisterChild: (childIdAttribute) => {
+        childrenIdAttrToIdRef.current.delete(childIdAttribute);
+        rerender();
+      },
       parentId: itemId,
     }),
     [itemId],

packages/x-tree-view/src/internals/plugins/selection/TreeViewSelectionPlugin.test.tsx

@@ -1,5 +1,5 @@
 import { spy } from 'sinon';
-import { fireEvent, act } from '@mui/internal-test-utils';
+import { fireEvent, act, waitFor } from '@mui/internal-test-utils';
 import { describeTreeView } from 'test/utils/tree-view/describeTreeView';
 import { clearWarningsCache } from '@mui/x-internals/warning';
 import { TreeViewAnyStore } from '../../models';
@@ -880,7 +880,7 @@ describeTreeView<TreeViewAnyStore>(
           expect(view.getSelectedTreeItems()).to.deep.equal(['1', '1.2']);
         });
 
-        it('should select all the children when selecting a collapsed parent and then expanding', () => {
+        it('should select all the children when selecting a collapsed parent and then expanding', async () => {
           const view = render({
             multiSelect: true,
             checkboxSelection: true,
@@ -890,7 +890,9 @@ describeTreeView<TreeViewAnyStore>(
 
           fireEvent.click(view.getItemCheckboxInput('1'));
           fireEvent.click(view.getItemContent('1'));
-          expect(view.getSelectedTreeItems()).to.deep.equal(['1', '1.1', '1.2']);
+          await waitFor(() => {
+            expect(view.getSelectedTreeItems()).to.deep.equal(['1', '1.1', '1.2']);
+          });
         });
       });