fix(cdk/tree): CdkTreeNodeToggle focuses node when toggling it

Fix focus behavior of CdkTreeNodeToggle. When toggling the expanded or
collapsed state of a node, focus that node. Fix issue where end user
cannot tab into tree component when collaping the parent of the active
node.

Before this commit is applied, there is a bug where end user cannot tab
into the tree.

Reproduction steps
1. Active a tree node
2. (focus state renders)
3. Using mouse, collapse parent of node from step 1.
4. (tree node collapses)
5. Press Tab, then shift + tab
6. (item before tree is focused. Can tab into the tree)

With this commit applied, above issue is no longer expected to
reproduce.

This commit message is only for reviewers of this PR and can be deleted
when landing this change in main.

Author: zarend

src/cdk/tree/toggle.ts

@@ -39,11 +39,17 @@ export class CdkTreeNodeToggle<T, K = T> {
 
   constructor(protected _tree: CdkTree<T, K>, protected _treeNode: CdkTreeNode<T, K>) {}
 
+  // Toggle the expanded or collapsed state of this node.
+  //
+  // Focus this node with expanding or collapsing it. This ensures that the active node will always
+  // be visible when expanding and collapsing.
   _toggle(event: Event): void {
     this.recursive
       ? this._tree.toggleDescendants(this._treeNode.data)
       : this._tree.toggle(this._treeNode.data);
 
+    this._tree._keyManager.focusItem(this._treeNode);
+
     event.stopPropagation();
   }
 

src/cdk/tree/tree-redesign.spec.ts

@@ -270,6 +270,28 @@ describe('CdkTree redesign', () => {
           .toBe(0);
       });
 
+      it('should focus a node when collapsing it', () => {
+        // Create a tree with two nodes. A parent node and its child.
+        dataSource.clear();
+        const parent = dataSource.addData();
+        dataSource.addChild(parent);
+
+        component.tree.expandAll();
+        fixture.detectChanges();
+
+        // focus the child node
+        getNodes(treeElement)[1].click();
+        fixture.detectChanges();
+
+        // collapse the parent node
+        getNodes(treeElement)[0].click();
+        fixture.detectChanges();
+
+        expect(getNodes(treeElement).map(x => x.getAttribute('tabindex')))
+          .withContext(`Expecting parent node to be focused since it was collapsed.`)
+          .toEqual(['0', '-1']);
+      });
+
       it('should expand/collapse the node recursively', () => {
         expect(dataSource.data.length).toBe(3);
 
@@ -1312,22 +1334,33 @@ class FakeDataSource extends DataSource<TestData> {
     return child;
   }
 
-  addData(level: number = 1) {
+  addData(level: number = 1): TestData {
     const nextIndex = ++this.dataIndex;
 
     let copiedData = this.data.slice();
-    copiedData.push(
-      new TestData(`topping_${nextIndex}`, `cheese_${nextIndex}`, `base_${nextIndex}`, level),
+    const newData = new TestData(
+      `topping_${nextIndex}`,
+      `cheese_${nextIndex}`,
+      `base_${nextIndex}`,
+      level,
     );
+    copiedData.push(newData);
 
     this.data = copiedData;
+
+    return newData;
   }
 
   getRecursiveData(nodes: TestData[] = this._dataChange.getValue()): TestData[] {
     return [
       ...new Set(nodes.flatMap(parent => [parent, ...this.getRecursiveData(parent.children)])),
     ];
   }
+
+  clear() {
+    this.data = [];
+    this.dataIndex = 0;
+  }
 }
 
 function getNodes(treeElement: Element): HTMLElement[] {