[EuiDataGrid] Fix multiple grid focus restoration issues (#5530)

* [#5517] Add `isFocusedCellInView` state to track visibility

- and use this to determine whether or not the entire grid should have a tabindex/be focusable, rather than a simple 'hasFocus' check which doesn't account for virtualization

- Add new mockFocusContext for easier testing for various tests that need to set a focus context

- EuiDataGridCell
  - DRY out an `this.isFocusedCell` method
  - Call `setIsFocusedCellInView` on mount and on unmount (`setFocusedCell` handles setting true/false on new cell focus)
  - Write new unit tests for componentWillUnmount
  - Clean up unnecessary uncovered function fallback for this.unsubscribeCell - we're already checking for a falsy value before calling it

* [#5517] Add `onItemsRendered` ref to store currently visible/virtualized refs

- see https://react-window.vercel.app/#/api/FixedSizeGrid for `onItemsRendered` API
- this struck me as the quickest/easiest way to determine which virtualized cells are in view
- I opted to save this as a ref instead of as state as I didn't see the need to cause a rerender

- `focusFirstVisibleInteractiveCell` was split out to its own fn as it will shortly be used elsewhere to fix another focus bug in the grid

* [#4923] Fix focus returning to document body when hiding a column via header

* [#5476] Fix keyboard navigation after hiding a column via header actions

* Add changelog entry

* [PR feedback] Add explanatory comment in unit test

* [PR feedback] Account for empty grids

- gridItemsRendered will be falsey if rowCount is 0 and will cause an error on tab, so we should opt to simply leave focus on the grid body wrapper

* [bug] Handle arrow keydown behavior on sticky header when row 0 is not in view

CHANGELOG.md

@@ -3,6 +3,10 @@
 - Added virtulized rendering option to `EuiSelectableList` with `isVirtualized` ([#5521](https://github.com/elastic/eui/pull/5521))
 - Added expanded option properties to `EuiSelectableOption` with `data` ([#5521](https://github.com/elastic/eui/pull/5521))
 
+**Bug fixes**
+
+- Fixed multiple bugs with `EuiDataGrid` keyboard focus restoration ([#5530](https://github.com/elastic/eui/pull/5530))
+
 **Breaking changes**
 
 - Changed `EuiSearchBar` to preserve phrases with leading and trailing spaces, instead of dropping surrounding whitespace ([#5514](https://github.com/elastic/eui/pull/5514))

src/components/datagrid/body/data_grid_body.test.tsx

@@ -52,6 +52,7 @@ describe('EuiDataGridBody', () => {
  resetAfterRowIndex: jest.fn(),
  } as any,
  },
+ gridItemsRendered: {} as any,
  wrapperRef: { current: document.createElement('div') },
  };
 

src/components/datagrid/body/data_grid_body.tsx

@@ -243,6 +243,7 @@ export const EuiDataGridBody: FunctionComponent<EuiDataGridBodyProps> = (
  gridStyles,
  gridWidth,
  gridRef,
+ gridItemsRendered,
  wrapperRef,
  } = props;
 
@@ -442,6 +443,9 @@ export const EuiDataGridBody: FunctionComponent<EuiDataGridBodyProps> = (
  <Grid
  {...(virtualizationOptions ? virtualizationOptions : {})}
  ref={gridRef}
+ onItemsRendered={(itemsRendered) => {
+ gridItemsRendered.current = itemsRendered;
+ }}
  innerElementType={InnerElement}
  outerRef={outerGridRef}
  innerRef={innerGridRef}

src/components/datagrid/body/data_grid_cell.test.tsx

@@ -10,6 +10,7 @@ import React from 'react';
 import { mount, ReactWrapper } from 'enzyme';
 import { keys } from '../../../services';
 import { mockRowHeightUtils } from '../utils/__mocks__/row_heights';
+import { mockFocusContext } from '../utils/__mocks__/focus_context';
 import { DataGridFocusContext } from '../utils/focus';
 
 import { EuiDataGridCell } from './data_grid_cell';
@@ -166,27 +167,21 @@ describe('EuiDataGridCell', () => {
  });
 
  describe('componentDidMount', () => {
- const focusContext = {
- focusedCell: undefined,
- onFocusUpdate: jest.fn(),
- setFocusedCell: jest.fn(),
- };
-
  it('creates an onFocusUpdate subscription', () => {
  mount(
- <DataGridFocusContext.Provider value={focusContext}>
+ <DataGridFocusContext.Provider value={mockFocusContext}>
  <EuiDataGridCell {...requiredProps} />
  </DataGridFocusContext.Provider>
  );
 
- expect(focusContext.onFocusUpdate).toHaveBeenCalled();
+ expect(mockFocusContext.onFocusUpdate).toHaveBeenCalled();
  });
 
  it('mounts the cell with focus state if the current cell should be focused', () => {
  const focusSpy = jest.spyOn(HTMLElement.prototype, 'focus');
  const component = mount(
  <DataGridFocusContext.Provider
- value={{ ...focusContext, focusedCell: [3, 3] }}
+ value={{ ...mockFocusContext, focusedCell: [3, 3] }}
  >
  <EuiDataGridCell
  {...requiredProps}
@@ -198,6 +193,44 @@ describe('EuiDataGridCell', () => {
 
  expect((component.instance().state as any).isFocused).toEqual(true);
  expect(focusSpy).toHaveBeenCalledWith({ preventScroll: true });
+ expect(mockFocusContext.setIsFocusedCellInView).toHaveBeenCalledWith(
+ true
+ );
+ });
+ });
+
+ describe('componentWillUnmount', () => {
+ it('unsubscribes from the onFocusUpdate subscription', () => {
+ const unsubscribeCellMock = jest.fn();
+ mockFocusContext.onFocusUpdate.mockReturnValueOnce(unsubscribeCellMock);
+
+ const component = mount(
+ <DataGridFocusContext.Provider value={mockFocusContext}>
+ <EuiDataGridCell {...requiredProps} />
+ </DataGridFocusContext.Provider>
+ );
+ component.unmount();
+
+ expect(unsubscribeCellMock).toHaveBeenCalled();
+ });
+
+ it('sets isFocusedCellInView to false if the current cell is focused and unmounting due to being scrolled out of view', () => {
+ const component = mount(
+ <DataGridFocusContext.Provider
+ value={{ ...mockFocusContext, focusedCell: [3, 3] }}
+ >
+ <EuiDataGridCell
+ {...requiredProps}
+ colIndex={3}
+ visibleRowIndex={3}
+ />
+ </DataGridFocusContext.Provider>
+ );
+ component.unmount();
+
+ expect(mockFocusContext.setIsFocusedCellInView).toHaveBeenCalledWith(
+ false
+ );
  });
  });
 

src/components/datagrid/body/data_grid_cell.tsx

@@ -124,7 +124,7 @@ export class EuiDataGridCell extends Component<
  enableInteractions: false,
  disableCellTabIndex: false,
  };
- unsubscribeCell?: Function = () => {};
+ unsubscribeCell?: Function;
  focusTimeout: number | undefined;
  style = null;
 
@@ -234,16 +234,21 @@ export class EuiDataGridCell extends Component<
 
  // Account for virtualization - when a cell unmounts when scrolled out of view
  // and then remounts when scrolled back into view, it should retain focus state
- if (
- this.context.focusedCell?.[0] === colIndex &&
- this.context.focusedCell?.[1] === visibleRowIndex
- ) {
+ if (this.isFocusedCell()) {
  // The second flag sets preventScroll: true as a focus option, which prevents
  // hijacking the user's scroll behavior when the cell re-mounts on scroll
  this.onFocusUpdate(true, true);
+ this.context.setIsFocusedCellInView(true);
  }
  }
 
+ isFocusedCell = () => {
+ return (
+ this.context.focusedCell?.[0] === this.props.colIndex &&
+ this.context.focusedCell?.[1] === this.props.visibleRowIndex
+ );
+ };
+
  onFocusUpdate = (isFocused: boolean, preventScroll = false) => {
  this.setState({ isFocused }, () => {
  if (isFocused) {
@@ -257,6 +262,10 @@ export class EuiDataGridCell extends Component<
  if (this.unsubscribeCell) {
  this.unsubscribeCell();
  }
+
+ if (this.isFocusedCell()) {
+ this.context.setIsFocusedCellInView(false);
+ }
  }
 
  componentDidUpdate(prevProps: EuiDataGridCellProps) {

src/components/datagrid/body/header/column_actions.test.tsx

@@ -19,18 +19,22 @@ import {
 
 describe('getColumnActions', () => {
  const setVisibleColumns = jest.fn();
+ const focusFirstVisibleInteractiveCell = jest.fn();
  const setIsPopoverOpen = jest.fn();
  const switchColumnPos = jest.fn();
+ const setFocusedCell = jest.fn();
 
  const testArgs = {
  column: { id: 'B' },
  columns: [{ id: 'A' }, { id: 'B' }, { id: 'C' }],
  schema: {},
  schemaDetectors,
  setVisibleColumns,
+ focusFirstVisibleInteractiveCell,
  setIsPopoverOpen,
  sorting: undefined,
  switchColumnPos,
+ setFocusedCell,
  };
 
  // DRY test helper
@@ -104,9 +108,10 @@ describe('getColumnActions', () => {
  `);
  });
 
- it('sets column visibility on click', () => {
+ it('hides the current column on click and refocuses into the grid', () => {
  callActionOnClick(hideColumn);
  expect(setVisibleColumns).toHaveBeenCalledWith(['A', 'C']);
+ expect(focusFirstVisibleInteractiveCell).toHaveBeenCalled();
  });
  });
 
@@ -182,12 +187,18 @@ describe('getColumnActions', () => {
  `);
  });
 
- it('calls switchColumnPos on click', () => {
+ it('calls switchColumnPos and updates the focused cell column index on click', () => {
  callActionOnClick(moveLeft);
  expect(switchColumnPos).toHaveBeenCalledWith('B', 'A');
+ expect(setFocusedCell).toHaveBeenLastCalledWith([0, -1]);
 
  callActionOnClick(moveRight);
  expect(switchColumnPos).toHaveBeenCalledWith('B', 'C');
+ expect(setFocusedCell).toHaveBeenLastCalledWith([2, -1]);
+
+ // Quick note about the -1 row index above: `-1` is the index we set for our
+ // sticky header row, and these column actions can only ever be called by an
+ // interactive header cell, so we can safely and statically set the -1 index
  });
  });
 

src/components/datagrid/body/header/column_actions.tsx

@@ -13,6 +13,7 @@ import {
  EuiDataGridSchema,
  EuiDataGridSchemaDetector,
  EuiDataGridSorting,
+ DataGridFocusContextShape,
 } from '../../data_grid_types';
 import { EuiI18n } from '../../../i18n';
 import { EuiListGroupItemProps } from '../../../list_group';
@@ -28,9 +29,11 @@ interface GetColumnActions {
  schema: EuiDataGridSchema;
  schemaDetectors: EuiDataGridSchemaDetector[];
  setVisibleColumns: (columnId: string[]) => void;
+ focusFirstVisibleInteractiveCell: DataGridFocusContextShape['focusFirstVisibleInteractiveCell'];
  setIsPopoverOpen: (value: boolean) => void;
  sorting: EuiDataGridSorting | undefined;
  switchColumnPos: (colFromId: string, colToId: string) => void;
+ setFocusedCell: DataGridFocusContextShape['setFocusedCell'];
 }
 
 export const getColumnActions = ({
@@ -39,9 +42,11 @@ export const getColumnActions = ({
  schema,
  schemaDetectors,
  setVisibleColumns,
+ focusFirstVisibleInteractiveCell,
  setIsPopoverOpen,
  sorting,
  switchColumnPos,
+ setFocusedCell,
 }: GetColumnActions): EuiListGroupItemProps[] => {
  if (column.actions === false) {
  return [];
@@ -52,6 +57,7 @@ export const getColumnActions = ({
  column,
  columns,
  setVisibleColumns,
+ focusFirstVisibleInteractiveCell,
  }),
  ...getSortColumnActions({
  column,
@@ -63,6 +69,7 @@ export const getColumnActions = ({
  column,
  columns,
  switchColumnPos,
+ setFocusedCell,
  }),
  ...(column.actions?.additional || []),
  ];
@@ -85,20 +92,27 @@ export const getColumnActions = ({
  */
 type HideColumnAction = Pick<
  GetColumnActions,
- 'column' | 'columns' | 'setVisibleColumns'
+ | 'column'
+ | 'columns'
+ | 'setVisibleColumns'
+ | 'focusFirstVisibleInteractiveCell'
 >;
 
 export const getHideColumnAction = ({
  column,
  columns,
  setVisibleColumns,
+ focusFirstVisibleInteractiveCell,
 }: HideColumnAction): EuiListGroupItemProps[] => {
  const items = [];
 
- const onClickHideColumn = () =>
+ const onClickHideColumn = () => {
  setVisibleColumns(
  columns.filter((col) => col.id !== column.id).map((col) => col.id)
  );
+ // Since we hid the current column, we need to manually set focus back onto the grid
+ focusFirstVisibleInteractiveCell();
+ };
 
  const action = {
  label: (
@@ -122,13 +136,14 @@ export const getHideColumnAction = ({
  */
 type MoveColumnActions = Pick<
  GetColumnActions,
- 'column' | 'columns' | 'switchColumnPos'
+ 'column' | 'columns' | 'switchColumnPos' | 'setFocusedCell'
 >;
 
 const getMoveColumnActions = ({
  column,
  columns,
  switchColumnPos,
+ setFocusedCell,
 }: MoveColumnActions): EuiListGroupItemProps[] => {
  const items = [];
 
@@ -139,6 +154,7 @@ const getMoveColumnActions = ({
  const targetCol = columns[colIdx - 1];
  if (targetCol) {
  switchColumnPos(column.id, targetCol.id);
+ setFocusedCell([colIdx - 1, -1]);
  }
  };
  const action = {
@@ -158,6 +174,7 @@ const getMoveColumnActions = ({
  const targetCol = columns[colIdx + 1];
  if (targetCol) {
  switchColumnPos(column.id, targetCol.id);
+ setFocusedCell([colIdx + 1, -1]);
  }
  };
  const action = {

src/components/datagrid/body/header/data_grid_header_cell.tsx

@@ -21,6 +21,7 @@ import { EuiIcon } from '../../../icon';
 import { EuiListGroup } from '../../../list_group';
 import { EuiPopover } from '../../../popover';
 import { DataGridSortingContext } from '../../utils/sorting';
+import { DataGridFocusContext } from '../../utils/focus';
 import { EuiDataGridHeaderCellProps } from '../../data_grid_types';
 
 import { getColumnActions } from './column_actions';
@@ -58,6 +59,10 @@ export const EuiDataGridHeaderCell: FunctionComponent<EuiDataGridHeaderCellProps
  } = {};
  const screenReaderId = useGeneratedHtmlId();
 
+ const { setFocusedCell, focusFirstVisibleInteractiveCell } = useContext(
+ DataGridFocusContext
+ );
+
  const { sorting } = useContext(DataGridSortingContext);
  let sortString;
  if (sorting) {
@@ -92,9 +97,11 @@ export const EuiDataGridHeaderCell: FunctionComponent<EuiDataGridHeaderCellProps
  schema,
  schemaDetectors,
  setVisibleColumns,
+ focusFirstVisibleInteractiveCell,
  setIsPopoverOpen,
  sorting,
  switchColumnPos,
+ setFocusedCell,
  });
 
  const showColumnActions = columnActions && columnActions.length > 0;

src/components/datagrid/body/header/data_grid_header_cell_wrapper.test.tsx

@@ -12,6 +12,7 @@ import { mount } from 'enzyme';
 
 import { keys } from '../../../../services';
 import { DataGridFocusContext } from '../../utils/focus';
+import { mockFocusContext } from '../../utils/__mocks__/focus_context';
 
 import { EuiDataGridHeaderCellWrapper } from './data_grid_header_cell_wrapper';
 
@@ -23,16 +24,12 @@ describe('EuiDataGridHeaderCellWrapper', () => {
  children: <button />,
  };
 
- const focusContext = {
- setFocusedCell: jest.fn(),
- onFocusUpdate: jest.fn(),
- };
  const mountWithContext = (props = {}, isFocused = true) => {
- (focusContext.onFocusUpdate as jest.Mock).mockImplementation(
+ (mockFocusContext.onFocusUpdate as jest.Mock).mockImplementation(
  (_, callback) => callback(isFocused) // allows us to mock isFocused state
  );
  return mount(
- <DataGridFocusContext.Provider value={focusContext}>
+ <DataGridFocusContext.Provider value={mockFocusContext}>
  <EuiDataGridHeaderCellWrapper {...requiredProps} {...props} />
  </DataGridFocusContext.Provider>
  );
@@ -186,7 +183,7 @@ describe('EuiDataGridHeaderCellWrapper', () => {
  act(() => {
  headerCell.dispatchEvent(new FocusEvent('focusin'));
  });
- expect(focusContext.setFocusedCell).toHaveBeenCalled();
+ expect(mockFocusContext.setFocusedCell).toHaveBeenCalled();
  });
 
  it('re-enables and focuses cell interactives when already focused', () => {