fix(components): [tree] up/down keys stop working after expanding nodes (#21526)

* fix(components): [tree] Up/Down keys stop working after expanding nodes

Fix issue #21514

* fix(components): [tree] Up/Down keys stop working after expanding nodes

* fix(components): [tree] Up/Down keys stop working after expanding nodes

packages/components/tree/__tests__/tree.test.ts

@@ -1577,6 +1577,111 @@ describe('Tree.vue', () => {
     }
   })
 
+  test('expand and render-after-expand set to false', async () => {
+    const { wrapper } = getTreeVm(``, {
+      template: `
+        <div>
+          <el-tree ref="tree1" :data="data" node-key="id" :render-after-expand="false"></el-tree>
+        </div>
+      `,
+      data() {
+        return {
+          data: [
+            {
+              id: 1,
+              label: '一级 1',
+              children: [
+                {
+                  id: 11,
+                  label: '二级 1-1',
+                  children: [
+                    {
+                      id: 111,
+                      label: '三级 1-1',
+                    },
+                  ],
+                },
+              ],
+            },
+            {
+              id: 2,
+              label: '一级 2',
+              children: [
+                {
+                  id: 21,
+                  label: '二级 2-1',
+                  children: [
+                    {
+                      id: 211,
+                      label: '三级 2-1-1',
+                    },
+                  ],
+                },
+                {
+                  id: 22,
+                  label: '二级 2-2',
+                },
+              ],
+            },
+            {
+              id: 3,
+              label: '一级 3',
+              children: [
+                {
+                  id: 31,
+                  label: '二级 3-1',
+                },
+                {
+                  id: 32,
+                  label: '二级 3-2',
+                },
+              ],
+            },
+          ],
+        }
+      },
+    })
+    await nextTick()
+    let flag = false
+    function handleFocus() {
+      return () => (flag = true)
+    }
+
+    const fromElement = wrapper.find('div[data-key="2"]')
+    await fromElement.trigger('click')
+
+    expect(fromElement.classes('is-expanded')).toBe(true) // 判断是否已展开
+
+    const treeWrapper = wrapper.findComponent(Tree)
+
+    await sleep()
+    const tree = wrapper.findComponent({ name: 'ElTree' })
+    ;(tree.vm as InstanceType<typeof Tree>).setCurrentKey(2)
+
+    const allNodes = treeWrapper.findAll('.el-tree-node')
+    const visibleNodes = allNodes.filter((val) => {
+      const node = wrapper.vm.$refs.tree1.getNode(val.element.dataset.key)!
+      return node.parent?.expanded || node.parent?.level === 0
+    })
+    const len = visibleNodes.length
+    for (let i = 0; i < len; i++) {
+      if (visibleNodes[i + 1]) {
+        defineGetter(visibleNodes[i + 1].element, 'focus', handleFocus)
+        // 模拟按下下箭头键
+        visibleNodes[i].element.dispatchEvent(
+          new KeyboardEvent('keydown', {
+            code: 'ArrowDown',
+            bubbles: true,
+            cancelable: false,
+          })
+        )
+
+        expect(flag).toBe(true)
+        flag = false
+      }
+    }
+  })
+
   test('navigate with disabled', async () => {
     const wrapper = mount({
       template: `

packages/components/tree/src/model/useKeydown.ts

@@ -26,6 +26,15 @@ export function useKeydown({ el$ }: UseKeydownOption, store: Ref<TreeStore>) {
     })
   })
 
+  function canNodeFocus(treeItems: HTMLElement[], nextIndex: number): boolean {
+    const currentNode = store.value.getNode(treeItems[nextIndex].dataset.key!)
+    return (
+      currentNode.canFocus &&
+      currentNode.visible &&
+      (currentNode.parent?.expanded || currentNode.parent?.level === 0)
+    )
+  }
+
   const handleKeydown = (ev: KeyboardEvent): void => {
     const currentItem = ev.target as HTMLDivElement
     if (!currentItem.className.includes(ns.b('node'))) return
@@ -46,10 +55,7 @@ export function useKeydown({ el$ }: UseKeydownOption, store: Ref<TreeStore>) {
             : treeItems.length - 1
         const startIndex = nextIndex
         while (true) {
-          const currentNode = store.value.getNode(
-            treeItems[nextIndex].dataset.key!
-          )
-          if (currentNode.canFocus && currentNode.visible) {
+          if (canNodeFocus(treeItems, nextIndex)) {
             break
           }
 
@@ -71,10 +77,7 @@ export function useKeydown({ el$ }: UseKeydownOption, store: Ref<TreeStore>) {
             : 0
         const startIndex = nextIndex
         while (true) {
-          const currentNode = store.value.getNode(
-            treeItems[nextIndex].dataset.key!
-          )
-          if (currentNode.canFocus && currentNode.visible) {
+          if (canNodeFocus(treeItems, nextIndex)) {
             break
           }