From e8377b80533508c9383010a25617334e452e5ab2 Mon Sep 17 00:00:00 2001
From: Zach Arend <zachzach@google.com>
Date: Thu, 10 Aug 2023 20:29:07 +0000
Subject: [PATCH] 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.
---
 src/cdk/tree/toggle.ts             |  6 +++++
 src/cdk/tree/tree-redesign.spec.ts | 39 +++++++++++++++++++++++++++---
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/src/cdk/tree/toggle.ts b/src/cdk/tree/toggle.ts
index b7004e60f1cb..5853bf64f937 100644
--- a/src/cdk/tree/toggle.ts
+++ b/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();
   }
 
diff --git a/src/cdk/tree/tree-redesign.spec.ts b/src/cdk/tree/tree-redesign.spec.ts
index 13f0098b745c..a43a1dc4e999 100644
--- a/src/cdk/tree/tree-redesign.spec.ts
+++ b/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,15 +1334,21 @@ 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[] {
@@ -1328,6 +1356,11 @@ class FakeDataSource extends DataSource<TestData> {
       ...new Set(nodes.flatMap(parent => [parent, ...this.getRecursiveData(parent.children)])),
     ];
   }
+
+  clear() {
+    this.data = [];
+    this.dataIndex = 0;
+  }
 }
 
 function getNodes(treeElement: Element): HTMLElement[] {