[EuiDataGrid] Fix cells losing focus state when scrolled out of view

CHANGELOG.md

@@ -1,6 +1,8 @@
 ## [`main`](https://github.com/elastic/eui/tree/main)
 
-No public interface changes since `43.1.1`.
+**Bug fixes**
+
+- Fixed `EuiDataGrid` issue where a focused cell would lose focus when scrolled out of and back into view ([#5488](https://github.com/elastic/eui/pull/5488))
 
 ## [`43.1.1`](https://github.com/elastic/eui/tree/v43.1.1)
 

src-docs/src/views/datagrid/datagrid_example.js

@@ -358,8 +358,8 @@ export const DataGridExample = {
                 in memory level
               </Link>{' '}
               to have the grid automatically handle updating your columns.
-              Depending upon the level choosen you may need to manage the
-              content order separate from the grid.
+              Depending upon the level chosen, you may need to manage the
+              content order separately from the grid.
             </li>
             <li>
               <Link to="/tabular-content/data-grid-schemas-and-popovers/">

src-docs/src/views/datagrid/datagrid_focus_example.js

@@ -93,7 +93,7 @@ export const DataGridFocusExample = {
             color="warning"
             title="A caution about turning off cell expansion when the width of the column is unknown"
           >
-            In general, you should turn <EuiCode>isExpandible</EuiCode> to false
+            In general, you should turn <EuiCode>isExpandable</EuiCode> to false
             only when you know the exact width and number of items that a cell
             will include. Control columns that contain row actions are a good
             example of when to use them. In certain scenarios, allowing multiple

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 '../__mocks__/row_height_utils';
+import { DataGridFocusContext } from '../data_grid_context';
 
 import { EuiDataGridCell } from './data_grid_cell';
 
@@ -31,31 +32,24 @@ describe('EuiDataGridCell', () => {
     rowHeightUtils: mockRowHeightUtils,
   };
 
-  const mountEuiDataGridCellWithContext = ({ ...props } = {}) => {
-    const focusContext = {
-      setFocusedCell: jest.fn(),
-      onFocusUpdate: jest.fn(),
-    };
-    return mount(<EuiDataGridCell {...requiredProps} {...props} />, {
-      context: focusContext,
-    });
-  };
-
   beforeEach(() => jest.clearAllMocks());
 
   it('renders', () => {
-    const component = mountEuiDataGridCellWithContext();
+    const component = mount(<EuiDataGridCell {...requiredProps} />);
     expect(component).toMatchSnapshot();
   });
 
   it('renders cell buttons', () => {
-    const component = mountEuiDataGridCellWithContext({
-      isExpandable: false,
-      column: {
-        id: 'someColumn',
-        cellActions: [() => <button />],
-      },
-    });
+    const component = mount(
+      <EuiDataGridCell
+        {...requiredProps}
+        isExpandable={false}
+        column={{
+          id: 'someColumn',
+          cellActions: [() => <button />],
+        }}
+      />
+    );
     component.setState({ popoverIsOpen: true });
 
     const cellButtons = component.find('EuiDataGridCellButtons');
@@ -80,7 +74,7 @@ describe('EuiDataGridCell', () => {
         EuiDataGridCell.prototype,
         'shouldComponentUpdate'
       );
-      component = mountEuiDataGridCellWithContext();
+      component = mount(<EuiDataGridCell {...requiredProps} />);
     });
     afterEach(() => {
       shouldComponentUpdate.mockRestore();
@@ -164,13 +158,47 @@ describe('EuiDataGridCell', () => {
   describe('componentDidUpdate', () => {
     it('resets cell props when the cell columnId changes', () => {
       const setState = jest.spyOn(EuiDataGridCell.prototype, 'setState');
-      const component = mountEuiDataGridCellWithContext();
+      const component = mount(<EuiDataGridCell {...requiredProps} />);
 
       component.setProps({ columnId: 'newColumnId' });
       expect(setState).toHaveBeenCalledWith({ cellProps: {} });
     });
   });
 
+  describe('componentDidMount', () => {
+    const focusContext = {
+      focusedCell: undefined,
+      onFocusUpdate: jest.fn(),
+      setFocusedCell: jest.fn(),
+    };
+
+    it('creates an onFocusUpdate subscription', () => {
+      mount(
+        <DataGridFocusContext.Provider value={focusContext}>
+          <EuiDataGridCell {...requiredProps} />
+        </DataGridFocusContext.Provider>
+      );
+
+      expect(focusContext.onFocusUpdate).toHaveBeenCalled();
+    });
+
+    it('mounts the cell with focus state if the current cell should be focused', () => {
+      const component = mount(
+        <DataGridFocusContext.Provider
+          value={{ ...focusContext, focusedCell: [3, 3] }}
+        >
+          <EuiDataGridCell
+            {...requiredProps}
+            colIndex={3}
+            visibleRowIndex={3}
+          />
+        </DataGridFocusContext.Provider>
+      );
+
+      expect((component.instance().state as any).isFocused).toEqual(true);
+    });
+  });
+
   // TODO: Test ResizeObserver logic in Cypress alongside Jest
   describe('row height logic & resize observers', () => {
     describe('recalculateAutoHeight', () => {
@@ -184,9 +212,12 @@ describe('EuiDataGridCell', () => {
       it('sets the row height cache with cell heights on update', () => {
         (mockRowHeightUtils.isAutoHeight as jest.Mock).mockReturnValue(true);
 
-        const component = mountEuiDataGridCellWithContext({
-          rowHeightsOptions: { defaultHeight: 'auto' },
-        });
+        const component = mount(
+          <EuiDataGridCell
+            {...requiredProps}
+            rowHeightsOptions={{ defaultHeight: 'auto' }}
+          />
+        );
 
         triggerUpdate(component);
         expect(mockRowHeightUtils.setRowHeight).toHaveBeenCalled();
@@ -195,9 +226,12 @@ describe('EuiDataGridCell', () => {
       it('does not update the cache if cell height is not auto', () => {
         (mockRowHeightUtils.isAutoHeight as jest.Mock).mockReturnValue(false);
 
-        const component = mountEuiDataGridCellWithContext({
-          rowHeightsOptions: { defaultHeight: 34 },
-        });
+        const component = mount(
+          <EuiDataGridCell
+            {...requiredProps}
+            rowHeightsOptions={{ defaultHeight: 34 }}
+          />
+        );
 
         triggerUpdate(component);
         expect(mockRowHeightUtils.setRowHeight).not.toHaveBeenCalled();
@@ -212,10 +246,13 @@ describe('EuiDataGridCell', () => {
 
       describe('default height', () => {
         it('observes the first cell for size changes and calls this.props.setRowHeight on change', () => {
-          const component = mountEuiDataGridCellWithContext({
-            rowHeightsOptions: { defaultHeight: { lineCount: 3 } },
-            setRowHeight,
-          });
+          const component = mount(
+            <EuiDataGridCell
+              {...requiredProps}
+              rowHeightsOptions={{ defaultHeight: { lineCount: 3 } }}
+              setRowHeight={setRowHeight}
+            />
+          );
 
           callMethod(component);
           expect(
@@ -227,14 +264,17 @@ describe('EuiDataGridCell', () => {
 
       describe('row height overrides', () => {
         it('uses the rowHeightUtils.setRowHeight cache instead of this.props.setRowHeight', () => {
-          const component = mountEuiDataGridCellWithContext({
-            rowHeightsOptions: {
-              defaultHeight: { lineCount: 3 },
-              rowHeights: { 10: { lineCount: 10 } },
-            },
-            rowIndex: 10,
-            setRowHeight,
-          });
+          const component = mount(
+            <EuiDataGridCell
+              {...requiredProps}
+              rowHeightsOptions={{
+                defaultHeight: { lineCount: 3 },
+                rowHeights: { 10: { lineCount: 10 } },
+              }}
+              rowIndex={10}
+              setRowHeight={setRowHeight}
+            />
+          );
 
           callMethod(component);
           expect(
@@ -246,10 +286,13 @@ describe('EuiDataGridCell', () => {
       });
 
       it('recalculates when rowHeightsOptions.defaultHeight.lineCount changes', () => {
-        const component = mountEuiDataGridCellWithContext({
-          rowHeightsOptions: { defaultHeight: { lineCount: 7 } },
-          setRowHeight,
-        });
+        const component = mount(
+          <EuiDataGridCell
+            {...requiredProps}
+            rowHeightsOptions={{ defaultHeight: { lineCount: 7 } }}
+            setRowHeight={setRowHeight}
+          />
+        );
 
         component.setProps({
           rowHeightsOptions: { defaultHeight: { lineCount: 6 } },
@@ -258,10 +301,13 @@ describe('EuiDataGridCell', () => {
       });
 
       it('calculates undefined heights as single rows with a lineCount of 1', () => {
-        const component = mountEuiDataGridCellWithContext({
-          rowHeightsOptions: { defaultHeight: undefined },
-          setRowHeight,
-        });
+        const component = mount(
+          <EuiDataGridCell
+            {...requiredProps}
+            rowHeightsOptions={{ defaultHeight: undefined }}
+            setRowHeight={setRowHeight}
+          />
+        );
 
         callMethod(component);
         expect(
@@ -271,10 +317,13 @@ describe('EuiDataGridCell', () => {
       });
 
       it('does nothing if cell height is not lineCount or undefined', () => {
-        const component = mountEuiDataGridCellWithContext({
-          rowHeightsOptions: { defaultHeight: 34 },
-          setRowHeight,
-        });
+        const component = mount(
+          <EuiDataGridCell
+            {...requiredProps}
+            rowHeightsOptions={{ defaultHeight: 34 }}
+            setRowHeight={setRowHeight}
+          />
+        );
 
         callMethod(component);
         expect(setRowHeight).not.toHaveBeenCalled();
@@ -286,7 +335,7 @@ describe('EuiDataGridCell', () => {
   describe('interactions', () => {
     describe('keyboard events', () => {
       it('when cell is expandable', () => {
-        const component = mountEuiDataGridCellWithContext();
+        const component = mount(<EuiDataGridCell {...requiredProps} />);
         const preventDefault = jest.fn();
 
         component.simulate('keyDown', { preventDefault, key: keys.ENTER });
@@ -296,9 +345,9 @@ describe('EuiDataGridCell', () => {
       });
 
       it('when cell is not expandable', () => {
-        const component = mountEuiDataGridCellWithContext({
-          isExpandable: false,
-        });
+        const component = mount(
+          <EuiDataGridCell {...requiredProps} isExpandable={false} />
+        );
         const preventDefault = jest.fn();
 
         component.simulate('keyDown', { preventDefault, key: keys.ENTER });
@@ -321,28 +370,31 @@ describe('EuiDataGridCell', () => {
     });
 
     it('mouse events', () => {
-      const component = mountEuiDataGridCellWithContext();
+      const component = mount(<EuiDataGridCell {...requiredProps} />);
       component.simulate('mouseEnter');
       expect(component.state('enableInteractions')).toEqual(true);
       component.simulate('mouseLeave');
       expect(component.state('enableInteractions')).toEqual(false);
     });
 
     it('focus/blur events', () => {
-      const component = mountEuiDataGridCellWithContext();
+      const component = mount(<EuiDataGridCell {...requiredProps} />);
       component.simulate('focus');
       component.simulate('blur');
       expect(component.state('disableCellTabIndex')).toEqual(false);
     });
   });
 
   it('renders certain classes/styles if rowHeightOptions is passed', () => {
-    const component = mountEuiDataGridCellWithContext({
-      rowHeightsOptions: {
-        defaultHeight: 20,
-        rowHeights: { 0: 10 },
-      },
-    });
+    const component = mount(
+      <EuiDataGridCell
+        {...requiredProps}
+        rowHeightsOptions={{
+          defaultHeight: 20,
+          rowHeights: { 0: 10 },
+        }}
+      />
+    );
     component.setState({ popoverIsOpen: true });
 
     expect(

src/components/datagrid/body/data_grid_cell.tsx

@@ -225,10 +225,21 @@ export class EuiDataGridCell extends Component<
   };
 
   componentDidMount() {
+    const { colIndex, visibleRowIndex } = this.props;
+
     this.unsubscribeCell = this.context.onFocusUpdate(
-      [this.props.colIndex, this.props.visibleRowIndex],
+      [colIndex, visibleRowIndex],
       this.onFocusUpdate
     );
+
+    // 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
+    ) {
+      this.onFocusUpdate(true);
+    }
   }
 
   onFocusUpdate = (isFocused: boolean) => {

src/components/datagrid/data_grid.tsx

@@ -748,10 +748,11 @@ export const EuiDataGrid: FunctionComponent<EuiDataGridProps> = (props) => {
   );
   const datagridFocusContext = useMemo<DataGridFocusContextShape>(() => {
     return {
+      focusedCell,
       setFocusedCell,
       onFocusUpdate,
     };
-  }, [setFocusedCell, onFocusUpdate]);
+  }, [focusedCell, setFocusedCell, onFocusUpdate]);
 
   const gridId = useGeneratedHtmlId();
   const ariaLabelledById = useGeneratedHtmlId();

src/components/datagrid/data_grid_context.tsx

@@ -16,6 +16,7 @@ import {
 export const DataGridFocusContext = React.createContext<
   DataGridFocusContextShape
 >({
+  focusedCell: undefined,
   setFocusedCell: () => {},
   onFocusUpdate: () => () => {},
 });

src/components/datagrid/data_grid_types.ts

@@ -167,6 +167,7 @@ export type EuiDataGridFooterRowProps = CommonProps &
   };
 
 export interface DataGridFocusContextShape {
+  focusedCell?: EuiDataGridFocusedCell;
   setFocusedCell: (cell: EuiDataGridFocusedCell) => void;
   onFocusUpdate: (
     cell: EuiDataGridFocusedCell,