From 8a035de8d4ddb75d45feaacab181f55b1d7a57d7 Mon Sep 17 00:00:00 2001
From: Constance Chen <constance.chen.3@gmail.com>
Date: Tue, 11 Jan 2022 12:05:45 -0800
Subject: [PATCH] [#5476] Fix keyboard navigation after hiding a column via
 header actions

---
 .../datagrid/body/header/column_actions.test.tsx          | 6 +++++-
 src/components/datagrid/body/header/column_actions.tsx    | 8 +++++++-
 .../datagrid/body/header/data_grid_header_cell.tsx        | 5 ++++-
 3 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/components/datagrid/body/header/column_actions.test.tsx b/src/components/datagrid/body/header/column_actions.test.tsx
index 058353544bb..15b39f51e63 100644
--- a/src/components/datagrid/body/header/column_actions.test.tsx
+++ b/src/components/datagrid/body/header/column_actions.test.tsx
@@ -22,6 +22,7 @@ describe('getColumnActions', () => {
   const focusFirstVisibleInteractiveCell = jest.fn();
   const setIsPopoverOpen = jest.fn();
   const switchColumnPos = jest.fn();
+  const setFocusedCell = jest.fn();
 
   const testArgs = {
     column: { id: 'B' },
@@ -33,6 +34,7 @@ describe('getColumnActions', () => {
     setIsPopoverOpen,
     sorting: undefined,
     switchColumnPos,
+    setFocusedCell,
   };
 
   // DRY test helper
@@ -185,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]);
       });
     });
 
diff --git a/src/components/datagrid/body/header/column_actions.tsx b/src/components/datagrid/body/header/column_actions.tsx
index 2a557bdcd0a..3f594ac0703 100644
--- a/src/components/datagrid/body/header/column_actions.tsx
+++ b/src/components/datagrid/body/header/column_actions.tsx
@@ -33,6 +33,7 @@ interface GetColumnActions {
   setIsPopoverOpen: (value: boolean) => void;
   sorting: EuiDataGridSorting | undefined;
   switchColumnPos: (colFromId: string, colToId: string) => void;
+  setFocusedCell: DataGridFocusContextShape['setFocusedCell'];
 }
 
 export const getColumnActions = ({
@@ -45,6 +46,7 @@ export const getColumnActions = ({
   setIsPopoverOpen,
   sorting,
   switchColumnPos,
+  setFocusedCell,
 }: GetColumnActions): EuiListGroupItemProps[] => {
   if (column.actions === false) {
     return [];
@@ -67,6 +69,7 @@ export const getColumnActions = ({
       column,
       columns,
       switchColumnPos,
+      setFocusedCell,
     }),
     ...(column.actions?.additional || []),
   ];
@@ -133,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 = [];
 
@@ -150,6 +154,7 @@ const getMoveColumnActions = ({
       const targetCol = columns[colIdx - 1];
       if (targetCol) {
         switchColumnPos(column.id, targetCol.id);
+        setFocusedCell([colIdx - 1, -1]);
       }
     };
     const action = {
@@ -169,6 +174,7 @@ const getMoveColumnActions = ({
       const targetCol = columns[colIdx + 1];
       if (targetCol) {
         switchColumnPos(column.id, targetCol.id);
+        setFocusedCell([colIdx + 1, -1]);
       }
     };
     const action = {
diff --git a/src/components/datagrid/body/header/data_grid_header_cell.tsx b/src/components/datagrid/body/header/data_grid_header_cell.tsx
index d7cbb919618..880506e00df 100644
--- a/src/components/datagrid/body/header/data_grid_header_cell.tsx
+++ b/src/components/datagrid/body/header/data_grid_header_cell.tsx
@@ -59,7 +59,9 @@ export const EuiDataGridHeaderCell: FunctionComponent<EuiDataGridHeaderCellProps
   } = {};
   const screenReaderId = useGeneratedHtmlId();
 
-  const { focusFirstVisibleInteractiveCell } = useContext(DataGridFocusContext);
+  const { setFocusedCell, focusFirstVisibleInteractiveCell } = useContext(
+    DataGridFocusContext
+  );
 
   const { sorting } = useContext(DataGridSortingContext);
   let sortString;
@@ -99,6 +101,7 @@ export const EuiDataGridHeaderCell: FunctionComponent<EuiDataGridHeaderCellProps
     setIsPopoverOpen,
     sorting,
     switchColumnPos,
+    setFocusedCell,
   });
 
   const showColumnActions = columnActions && columnActions.length > 0;