restoring focus back to the tree if the user was focused on the loader

Author: LFDanLu

packages/@react-stately/tree/src/useTreeState.ts

@@ -65,8 +65,13 @@ export function useTreeState<T extends object>(props: TreeProps<T>): TreeState<T
 
   // Reset focused key if that item is deleted from the collection.
   useEffect(() => {
-    if (selectionState.focusedKey != null && !tree.getItem(selectionState.focusedKey)) {
-      selectionState.setFocusedKey(null);
+    if (selectionState.focusedKey != null) {
+      let focusedItem = tree.getItem(selectionState.focusedKey);
+      // TODO: do we want to have the same logic as useListState/useGridState where it tries to find the nearest row?
+      // We could possibly special case this loader case and have it try to find the item just before it/the parent
+      if (!focusedItem || focusedItem.type === 'loader' && !focusedItem.props.isLoading) {
+        selectionState.setFocusedKey(null);
+      }
     }
   // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [tree, selectionState.focusedKey]);

packages/react-aria-components/test/Tree.test.tsx

@@ -1592,6 +1592,27 @@ describe('Tree', () => {
         expect(documentLoader1ParentStyles.top).toBe('900px');
         expect(documentLoader1ParentStyles.height).toBe('30px');
       });
+
+      it('should restore focus to the tree if the loader is keyboard focused when loading finishes', async () => {
+        let tree = render(
+          <VirtualizedLoadingSentinelTree rootIsLoading />
+        );
+        let treeTester = testUtilUser.createTester('Tree', {root: tree.getByRole('treegrid')});
+        let rows = treeTester.rows;
+        expect(rows).toHaveLength(8);
+        let rootLoaderRow = rows[7];
+        expect(rootLoaderRow).toHaveTextContent('Loading...');
+
+        await user.tab();
+        await user.keyboard('{End}');
+        expect(document.activeElement).toBe(rootLoaderRow);
+
+        tree.rerender(
+          <VirtualizedLoadingSentinelTree />
+        );
+
+        expect(document.activeElement).toBe(treeTester.tree);
+      });
     });
   });