[EuiDataGrid] Fix inconsistent cell header focus when clicking another cell header when a cell header popover is already open

CHANGELOG.md

@@ -3,6 +3,7 @@
 **Bug fixes**
 
 - Fixed EuiDataGrid height issue when in full-screen mode and with scrolling content ([#5557](https://github.com/elastic/eui/pull/5557))
+- Fixed a focus bug in `EuiDataGrid` when clicking another cell header with an already-open cell header popover ([#5556](https://github.com/elastic/eui/pull/5556))
 
 ## [`46.1.0`](https://github.com/elastic/eui/tree/v46.1.0)
 

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

@@ -11,6 +11,8 @@ import { mount, shallow } from 'enzyme';
 
 import { EuiDataGridSorting } from '../../data_grid_types';
 import { DataGridSortingContext } from '../../utils/sorting';
+import { DataGridFocusContext } from '../../utils/focus';
+import { mockFocusContext } from '../../utils/__mocks__/focus_context';
 
 import { EuiDataGridHeaderCell } from './data_grid_header_cell';
 
@@ -145,10 +147,15 @@ describe('EuiDataGridHeaderCell', () => {
     });
 
     it('handles popover open', () => {
-      const component = mount(<EuiDataGridHeaderCell {...requiredProps} />);
+      const component = mount(
+        <DataGridFocusContext.Provider value={mockFocusContext}>
+          <EuiDataGridHeaderCell {...requiredProps} />
+        </DataGridFocusContext.Provider>
+      );
       component.find('.euiDataGridHeaderCell__button').simulate('click');
 
       expect(component.find('EuiPopover').prop('isOpen')).toEqual(true);
+      expect(mockFocusContext.setFocusedCell).toHaveBeenCalledWith([0, -1]);
     });
 
     it('handles popover close', () => {

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

@@ -152,9 +152,10 @@ export const EuiDataGridHeaderCell: FunctionComponent<EuiDataGridHeaderCellProps
           button={
             <button
               className="euiDataGridHeaderCell__button"
-              onClick={() =>
-                setIsPopoverOpen((isPopoverOpen) => !isPopoverOpen)
-              }
+              onClick={() => {
+                setFocusedCell([index, -1]);
+                setIsPopoverOpen((isPopoverOpen) => !isPopoverOpen);
+              }}
             >
               {sortingArrow}
               <div className="euiDataGridHeaderCell__content">