fix: (WIP) Fix test utils for selectionMode replace

packages/@react-aria/test-utils/src/gridlist.ts

@@ -64,6 +64,11 @@ export class GridListTester {
     if (targetIndex === -1) {
       throw new Error('Option provided is not in the gridlist');
     }
+
+    if (document.activeElement !== this._gridlist || !this._gridlist.contains(document.activeElement)) {
+      act(() => this._gridlist.focus());
+    }
+
     if (document.activeElement === this._gridlist) {
       await this.user.keyboard('[ArrowDown]');
     } else if (this._gridlist.contains(document.activeElement) && document.activeElement!.getAttribute('role') !== 'row') {
@@ -161,10 +166,6 @@ export class GridListTester {
         return;
       }
 
-      if (document.activeElement !== this._gridlist || !this._gridlist.contains(document.activeElement)) {
-        act(() => this._gridlist.focus());
-      }
-
       await this.keyboardNavigateToRow({row});
       await this.user.keyboard('[Enter]');
     } else {

packages/@react-aria/test-utils/src/listbox.ts

@@ -93,6 +93,10 @@ export class ListBoxTester {
       throw new Error('Option provided is not in the listbox');
     }
 
+    if (document.activeElement !== this._listbox || !this._listbox.contains(document.activeElement)) {
+      act(() => this._listbox.focus());
+    }
+
     if (document.activeElement === this._listbox) {
       await this.user.keyboard('[ArrowDown]');
     }
@@ -135,10 +139,6 @@ export class ListBoxTester {
         return;
       }
 
-      if (document.activeElement !== this._listbox || !this._listbox.contains(document.activeElement)) {
-        act(() => this._listbox.focus());
-      }
-
       await this.keyboardNavigateToOption({option});
       await this.user.keyboard(`[${keyboardActivation}]`);
     } else {
@@ -179,10 +179,6 @@ export class ListBoxTester {
         return;
       }
 
-      if (document.activeElement !== this._listbox || !this._listbox.contains(document.activeElement)) {
-        act(() => this._listbox.focus());
-      }
-
       await this.keyboardNavigateToOption({option});
       await this.user.keyboard('[Enter]');
     } else {

packages/@react-aria/test-utils/src/tabs.ts

@@ -84,6 +84,10 @@ export class TabsTester {
       throw new Error('Tab provided is not in the tablist');
     }
 
+    if (document.activeElement !== this._tablist || !this._tablist.contains(document.activeElement)) {
+      act(() => this._tablist.focus());
+    }
+
     if (!this._tablist.contains(document.activeElement)) {
       let selectedTab = this.selectedTab;
       if (selectedTab != null) {
@@ -137,10 +141,6 @@ export class TabsTester {
     }
 
     if (interactionType === 'keyboard') {
-      if (document.activeElement !== this._tablist || !this._tablist.contains(document.activeElement)) {
-        act(() => this._tablist.focus());
-      }
-
       let tabsOrientation = this._tablist.getAttribute('aria-orientation') || 'horizontal';
       await this.keyboardNavigateToTab({tab, orientation: tabsOrientation as Orientation});
       if (manualActivation) {

packages/@react-aria/test-utils/src/tree.ts

@@ -71,6 +71,11 @@ export class TreeTester {
     if (targetIndex === -1) {
       throw new Error('Option provided is not in the tree');
     }
+
+    if (document.activeElement !== this._tree || !this._tree.contains(document.activeElement)) {
+      act(() => this._tree.focus());
+    }
+
     if (document.activeElement === this.tree) {
       await this.user.keyboard('[ArrowDown]');
     } else if (this._tree.contains(document.activeElement) && document.activeElement!.getAttribute('role') !== 'row') {
@@ -89,6 +94,8 @@ export class TreeTester {
     }
   };
 
+  // TODO: we'll need to support selectionBehavior="replace", where clicks/keyboard will need to go through different flows
+  // for both single selection and multiple selection due to the nature of the selection replacement on focus
   /**
    * Toggles the selection for the specified tree row. Defaults to using the interaction type set on the tree tester.
    */
@@ -135,7 +142,7 @@ export class TreeTester {
         // Note that long press interactions with rows is strictly touch only for grid rows
         await triggerLongPress({element: cell, advanceTimer: this._advanceTimer, pointerOpts: {pointerType: 'touch'}});
       } else {
-        await pressElement(this.user, cell, interactionType);
+        await pressElement(this.user, row, interactionType);
       }
     }
   };
@@ -206,10 +213,6 @@ export class TreeTester {
         return;
       }
 
-      if (document.activeElement !== this._tree || !this._tree.contains(document.activeElement)) {
-        act(() => this._tree.focus());
-      }
-
       await this.keyboardNavigateToRow({row});
       await this.user.keyboard('[Enter]');
     } else {

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

@@ -463,6 +463,102 @@ describe('Tree', () => {
     expect(treeTester.selectedRows[0]).toBe(row1);
   });
 
+  it('highlight selection TreeView can select a row via keyboard', async function () {
+    let items = [
+      {
+        id: 'projects',
+        name: 'Projects',
+        childItems: [
+          {id: 'project-1', name: 'Project 1'},
+          {
+            id: 'project-2',
+            name: 'Project 2',
+            childItems: [
+              {id: 'document-a', name: 'Document A'},
+              {id: 'document-b', name: 'Document B'}
+            ]
+          }
+        ]
+      },
+      {
+        id: 'reports',
+        name: 'Reports',
+        childItems: [
+          {id: 'report-1', name: 'Reports 1'}
+        ]
+      }
+    ];
+
+    let {getByTestId} = render(
+      <TreeView
+        aria-label="Example tree with dynamic content"
+        height="size-3000"
+        maxWidth="size-6000"
+        defaultExpandedKeys={['projects', 'project-2']}
+        items={items}
+        data-testid="action-rail-tree"
+        selectionMode="single"
+        selectionStyle="highlight">
+        {(item: any) => (
+          <DynamicTreeItem id={item.id} childItems={item.childItems} textValue={item.name} name={item.name} />
+        )}
+      </TreeView>
+    );
+    let treeTester = testUtilUser.createTester('Tree', {
+      user,
+      root: getByTestId('action-rail-tree'),
+      interactionType: 'keyboard'
+    });
+
+    let rows = treeTester.rows;
+    await treeTester.toggleRowSelection({row: 0});
+    expect(treeTester.selectedRows).toHaveLength(1);
+    expect(rows[0]).toHaveAttribute('aria-selected', 'true');
+
+    await treeTester.toggleRowSelection({row: 1});
+    expect(treeTester.selectedRows).toHaveLength(1);
+    expect(rows[1]).toHaveAttribute('aria-selected', 'true');
+
+    await treeTester.toggleRowSelection({row: 0});
+    expect(treeTester.selectedRows).toHaveLength(1);
+    expect(rows[0]).toHaveAttribute('aria-selected', 'true');
+  });
+
+  // TODO: replace the test above this one with a similar set up as the below
+  it('should perform selection for highlight mode with single selection', async () => {
+    let {getByRole} = render(<StaticTree treeProps={{selectionMode: 'single', selectionStyle: 'highlight'}} />);
+    let treeTester = testUtilUser.createTester('Tree', {user, root: getByRole('treegrid')});
+    let rows = treeTester.rows;
+
+    for (let row of treeTester.rows) {
+      let checkbox = within(row).queryByRole('checkbox');
+      expect(checkbox).toBeNull();
+      expect(row).toHaveAttribute('aria-selected', 'false');
+      expect(row).not.toHaveAttribute('data-selected');
+      expect(row).toHaveAttribute('data-selection-mode', 'multiple');
+    }
+
+    let row2 = rows[2];
+    await treeTester.toggleRowSelection({row: 'Projects-1'});
+    expect(row2).toHaveAttribute('aria-selected', 'true');
+    expect(row2).toHaveAttribute('data-selected', 'true');
+    expect(onSelectionChange).toHaveBeenCalledTimes(1);
+    expect(new Set(onSelectionChange.mock.calls[0][0])).toEqual(new Set(['Projects-1']));
+    expect(treeTester.selectedRows).toHaveLength(1);
+    expect(treeTester.selectedRows[0]).toBe(row2);
+
+    let row1 = rows[1];
+    await treeTester.toggleRowSelection({row: row1});
+    expect(row1).toHaveAttribute('aria-selected', 'true');
+    expect(row1).toHaveAttribute('data-selected', 'true');
+    expect(row2).toHaveAttribute('aria-selected', 'false');
+    expect(row2).not.toHaveAttribute('data-selected');
+    expect(onSelectionChange).toHaveBeenCalledTimes(2);
+    expect(new Set(onSelectionChange.mock.calls[1][0])).toEqual(new Set(['Projects']));
+    expect(treeTester.selectedRows).toHaveLength(1);
+    expect(treeTester.selectedRows[0]).toBe(row1);
+  });
+
   it('should render a chevron for an expandable row marked with hasChildItems', () => {
     let {getAllByRole} = render(
       <TreeView aria-label="test tree">