fix: tree focus management (adobe#9531)

* fix: tree focus management

* one possible fix

* fix issue in treeview directly

* add table test because behavior should be different

* fix lint

Author: snowystinger

packages/@react-aria/selection/src/useSelectableItem.ts

@@ -200,6 +200,12 @@ export function useSelectableItem(options: SelectableItemOptions): SelectableIte
     };
   }
 
+  useEffect(() => {
+    if (isDisabled && manager.focusedKey === key) {
+      manager.setFocusedKey(null);
+    }
+  }, [manager, isDisabled, key]);
+
   // With checkbox selection, onAction (i.e. navigation) becomes primary, and occurs on a single click of the row.
   // Clicking the checkbox enters selection mode, after which clicking anywhere on any row toggles selection for that row.
   // With highlight selection, onAction is secondary, and occurs on double click. Single click selects the row.

packages/@react-spectrum/s2/test/TableView.test.tsx

@@ -24,17 +24,21 @@ import {
   TableView,
   Text
 } from '../src';
+import {DisabledBehavior} from '@react-types/shared';
 import Filter from '../s2wf-icons/S2_Icon_Filter_20_N.svg';
-import React from 'react';
-import {User} from '@react-aria/test-utils';
+import {pointerMap, User} from '@react-aria/test-utils';
+import React, {useState} from 'react';
+import userEvent from '@testing-library/user-event';
 
 // @ts-ignore
 window.getComputedStyle = (el) => el.style;
 
 describe('TableView', () => {
   let offsetWidth, offsetHeight;
+  let user;
   let testUtilUser = new User({advanceTimer: jest.advanceTimersByTime});
   beforeAll(function () {
+    user = userEvent.setup({delay: null, pointerMap});
     offsetWidth = jest.spyOn(window.HTMLElement.prototype, 'clientWidth', 'get').mockImplementation(() => 400);
     offsetHeight = jest.spyOn(window.HTMLElement.prototype, 'clientHeight', 'get').mockImplementation(() => 200);
     jest.useFakeTimers();
@@ -108,4 +112,40 @@ describe('TableView', () => {
     await tableTester.triggerColumnHeaderAction({column: 1, action: 0, interactionType: 'keyboard'});
     expect(onAction).toHaveBeenCalledTimes(1);
   });
+
+  it('if the previously focused cell\'s row is disabled, the focus should still be restored to the cell when the disabled behavior is changed and the user navigates to the collection', async () => {
+    function Example() {
+      let [disabledBehavior, setDisabledBehavior] = useState<DisabledBehavior>('selection');
+      return (
+        <>
+          <button>Before</button>
+          <TableView aria-label="Dynamic table" disabledBehavior={disabledBehavior} disabledKeys={['2']}>
+            <TableHeader columns={columns}>
+              {(column) => <Column {...column}>{column.name}</Column>}
+            </TableHeader>
+            <TableBody>
+              <Row id="1"><Cell>Foo 1</Cell><Cell>Bar 1</Cell><Cell>Baz 1</Cell><Cell>Yah 1</Cell></Row>
+              <Row id="2"><Cell>Foo 2</Cell><Cell>Bar 2</Cell><Cell>Baz 2</Cell><Cell>Yah 2</Cell></Row>
+              <Row id="3"><Cell>Foo 3</Cell><Cell>Bar 3</Cell><Cell>Baz 3</Cell><Cell>Yah 3</Cell></Row>
+            </TableBody>
+          </TableView>
+          <button onClick={() => setDisabledBehavior('all')}>After</button>
+        </>
+      );
+    }
+
+    let {getAllByRole, getByRole} = render(<Example />);
+    await user.click(document.body);
+
+    let cells = getAllByRole('gridcell');
+    let afterButton = getByRole('button', {name: 'After'});
+    await user.click(cells[3]); // Bar 2
+    expect(document.activeElement).toBe(cells[3]);
+
+    await user.click(afterButton);
+    await user.tab({shift: true});
+    await user.tab({shift: true});
+    await user.tab();
+    expect(document.activeElement).toBe(cells[3]);
+  });
 });

packages/@react-spectrum/s2/test/TreeView.test.tsx

@@ -9,9 +9,10 @@ import {
   TreeViewItemContent
 } from '../src';
 import {AriaTreeTests} from '../../../react-aria-components/test/AriaTree.test-util';
+import {DisabledBehavior} from '@react-types/shared';
 import FileTxt from '../s2wf-icons/S2_Icon_FileText_20_N.svg';
 import Folder from '../s2wf-icons/S2_Icon_Folder_20_N.svg';
-import React from 'react';
+import React, {useState} from 'react';
 import userEvent from '@testing-library/user-event';
 
 AriaTreeTests({
@@ -493,4 +494,74 @@ describe('TreeView', () => {
     let menu = queryByRole('menu');
     expect(menu).not.toBeInTheDocument();
   });
+
+  it('if the previously focused item is disabled, the focus should move to the first item when coming back to the collection', async () => {
+    function Example() {
+      let [disabledBehavior, setDisabledBehavior] = useState<DisabledBehavior>('selection');
+      return (
+        <>
+          <button>Before</button>
+          <TreeView aria-label="Test tree" disabledBehavior={disabledBehavior} disabledKeys={['school']}>
+            <TreeViewItem id="projects" textValue="Projects">
+              <TreeViewItemContent>
+                <Text>Projects</Text>o
+                <Folder />
+                <ActionMenu>
+                  <MenuItem id="edit">
+                    <Text>Edit</Text>
+                  </MenuItem>
+                  <MenuItem id="delete">
+                    <Text>Delete</Text>
+                  </MenuItem>
+                </ActionMenu>
+              </TreeViewItemContent>
+            </TreeViewItem>
+            <TreeViewItem id="school" textValue="School">
+              <TreeViewItemContent>
+                <Text>School</Text>o
+                <Folder />
+                <ActionMenu>
+                  <MenuItem id="edit">
+                    <Text>Edit</Text>
+                  </MenuItem>
+                  <MenuItem id="delete">
+                    <Text>Delete</Text>
+                  </MenuItem>
+                </ActionMenu>
+              </TreeViewItemContent>
+              <TreeViewItem id="homework-1" textValue="Homework-1" isDisabled>
+                <TreeViewItemContent>
+                  <Text>Homework-1</Text>
+                  <Folder />
+                  <ActionMenu>
+                    <MenuItem id="edit">
+                      <Text>Edit</Text>
+                    </MenuItem>
+                    <MenuItem id="delete">
+                      <Text>Delete</Text>
+                    </MenuItem>
+                  </ActionMenu>
+                </TreeViewItemContent>
+              </TreeViewItem>
+            </TreeViewItem>
+          </TreeView>
+          <button onClick={() => setDisabledBehavior('all')}>After</button>
+        </>
+      );
+    }
+
+    let {getAllByRole, getByRole} = render(<Example />);
+    await user.click(document.body);
+
+    let rows = getAllByRole('row');
+    let afterButton = getByRole('button', {name: 'After'});
+    await user.click(rows[1]);
+    expect(document.activeElement).toBe(rows[1]);
+
+    await user.click(afterButton);
+    await user.tab({shift: true});
+    await user.tab({shift: true});
+    await user.tab();
+    expect(document.activeElement).toBe(rows[0]);
+  });
 });