[EuiDataGrid] Fix multiple grid focus restoration issues

CHANGELOG.md

@@ -1,5 +1,9 @@
 ## [`main`](https://github.com/elastic/eui/tree/main)
 
+**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-docs/src/views/datagrid/focus.js

@@ -17,7 +17,7 @@ import {
 
 const data = [];
 
-for (let i = 0; i < 10; i++) {
+for (let i = 0; i < 25; i++) {
   data.push([
     <span>{fake('{{name.firstName}}')}</span>,
     <span>{fake('{{name.firstName}}')}</span>,
@@ -214,7 +214,7 @@ export default () => {
   );
 
   const [pagination, setPagination] = useState({
-    pageSize: 4,
+    pageSize: 20,
     pageIndex: 0,
     pageSizeOptions: [4],
   });
@@ -250,6 +250,7 @@ export default () => {
           onChangeItemsPerPage,
           onChangePage,
         }}
+        height={300}
       />
     </>
   );

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,14 @@ 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]);
       });
     });
 

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;