From 58cc2588e6a12d96004053a035e33d931a21fce5 Mon Sep 17 00:00:00 2001
From: Matheus Wichman <matheushw@outlook.com>
Date: Thu, 29 Jul 2021 23:02:11 -0300
Subject: [PATCH] [DataGrid] Drop autoFocus

---
 .../grid/components/cell/GridCell.tsx         | 18 +++--
 .../components/cell/GridEditBooleanCell.tsx   | 10 ++-
 .../grid/components/cell/GridEditDateCell.tsx | 10 ++-
 .../components/cell/GridEditInputCell.tsx     | 10 ++-
 .../cell/GridEditSingleSelectCell.tsx         | 13 ++-
 .../_modules_/grid/components/cell/index.ts   |  1 +
 .../grid/hooks/features/focus/useGridFocus.ts | 78 +++++++++++-------
 .../hooks/features/rows/useGridEditRows.ts    |  1 +
 .../hooks/features/rows/useGridParamsApi.ts   | 14 ++--
 .../src/tests/keyboard.DataGrid.test.tsx      |  1 +
 .../x-grid/src/tests/editRows.XGrid.test.tsx  | 80 ++++++++++++++++++-
 .../grid/x-grid/src/tests/rows.XGrid.test.tsx | 57 ++++++++++++-
 .../src/stories/grid-rows.stories.tsx         | 15 +++-
 13 files changed, 256 insertions(+), 52 deletions(-)

diff --git a/packages/grid/_modules_/grid/components/cell/GridCell.tsx b/packages/grid/_modules_/grid/components/cell/GridCell.tsx
index 2fba02dd00e7..520350f69b38 100644
--- a/packages/grid/_modules_/grid/components/cell/GridCell.tsx
+++ b/packages/grid/_modules_/grid/components/cell/GridCell.tsx
@@ -145,13 +145,17 @@ export const GridCell = React.memo(function GridCell(props: GridCellProps) {
   };
 
   React.useLayoutEffect(() => {
-    const doc = ownerDocument(apiRef!.current.rootElementRef!.current as HTMLElement);
+    if (!hasFocus) {
+      return;
+    }
+
+    if (cellMode === 'edit' && isEditable) {
+      return;
+    }
+
+    const doc = ownerDocument(apiRef!.current.rootElementRef!.current as HTMLElement)!;
 
-    if (
-      hasFocus &&
-      cellRef.current &&
-      (!doc.activeElement || !cellRef.current!.contains(doc.activeElement))
-    ) {
+    if (cellRef.current && (doc.activeElement || !cellRef.current!.contains(doc.activeElement))) {
       const focusableElement = cellRef.current!.querySelector('[tabindex="0"]') as HTMLElement;
       if (focusableElement) {
         focusableElement!.focus();
@@ -175,7 +179,7 @@ export const GridCell = React.memo(function GridCell(props: GridCellProps) {
       data-mode={cellMode}
       aria-colindex={colIndex + 1}
       style={style}
-      tabIndex={tabIndex}
+      tabIndex={cellMode === 'view' || !isEditable ? tabIndex : -1}
       {...eventsHandlers}
     >
       {children || valueToRender?.toString()}
diff --git a/packages/grid/_modules_/grid/components/cell/GridEditBooleanCell.tsx b/packages/grid/_modules_/grid/components/cell/GridEditBooleanCell.tsx
index 8ca3887b9fd1..48b0141ab94a 100644
--- a/packages/grid/_modules_/grid/components/cell/GridEditBooleanCell.tsx
+++ b/packages/grid/_modules_/grid/components/cell/GridEditBooleanCell.tsx
@@ -19,12 +19,14 @@ export function GridEditBooleanCell(
     colDef,
     cellMode,
     isEditable,
+    tabIndex,
     className,
     getValue,
     hasFocus,
     ...other
   } = props;
 
+  const inputRef = React.useRef<HTMLInputElement>(null);
   const id = useId();
   const [valueState, setValueState] = React.useState(value);
 
@@ -41,11 +43,17 @@ export function GridEditBooleanCell(
     setValueState(value);
   }, [value]);
 
+  React.useLayoutEffect(() => {
+    if (hasFocus && inputRef.current) {
+      inputRef.current.focus();
+    }
+  }, [hasFocus]);
+
   return (
     <label htmlFor={id} className={clsx('MuiDataGrid-editBooleanCell', className)} {...other}>
       <Checkbox
-        autoFocus
         id={id}
+        inputRef={inputRef}
         checked={Boolean(valueState)}
         onChange={handleChange}
         size="small"
diff --git a/packages/grid/_modules_/grid/components/cell/GridEditDateCell.tsx b/packages/grid/_modules_/grid/components/cell/GridEditDateCell.tsx
index 0bdab8d6efd5..1fcebf2af00a 100644
--- a/packages/grid/_modules_/grid/components/cell/GridEditDateCell.tsx
+++ b/packages/grid/_modules_/grid/components/cell/GridEditDateCell.tsx
@@ -13,11 +13,13 @@ export function GridEditDateCell(props: GridCellParams & InputBaseProps) {
     colDef,
     cellMode,
     isEditable,
+    tabIndex,
     hasFocus,
     getValue,
     ...other
   } = props;
 
+  const inputRef = React.useRef<any>();
   const [valueState, setValueState] = React.useState(value);
 
   const handleChange = React.useCallback(
@@ -54,6 +56,12 @@ export function GridEditDateCell(props: GridCellParams & InputBaseProps) {
     setValueState(value);
   }, [value]);
 
+  React.useLayoutEffect(() => {
+    if (hasFocus && inputRef.current) {
+      inputRef.current.focus();
+    }
+  }, [hasFocus]);
+
   let valueToDisplay = valueState || '';
   if (valueState instanceof Date) {
     const offset = valueState.getTimezoneOffset();
@@ -63,7 +71,7 @@ export function GridEditDateCell(props: GridCellParams & InputBaseProps) {
 
   return (
     <InputBase
-      autoFocus
+      inputRef={inputRef}
       fullWidth
       className="MuiDataGrid-editInputCell"
       type={isDateTime ? 'datetime-local' : 'date'}
diff --git a/packages/grid/_modules_/grid/components/cell/GridEditInputCell.tsx b/packages/grid/_modules_/grid/components/cell/GridEditInputCell.tsx
index f84c9773aed7..896982e8094a 100644
--- a/packages/grid/_modules_/grid/components/cell/GridEditInputCell.tsx
+++ b/packages/grid/_modules_/grid/components/cell/GridEditInputCell.tsx
@@ -13,11 +13,13 @@ export function GridEditInputCell(props: GridCellParams & InputBaseProps) {
     colDef,
     cellMode,
     isEditable,
+    tabIndex,
     hasFocus,
     getValue,
     ...other
   } = props;
 
+  const inputRef = React.useRef<any>();
   const [valueState, setValueState] = React.useState(value);
 
   const handleChange = React.useCallback(
@@ -33,9 +35,15 @@ export function GridEditInputCell(props: GridCellParams & InputBaseProps) {
     setValueState(value);
   }, [value]);
 
+  React.useLayoutEffect(() => {
+    if (hasFocus && inputRef.current) {
+      inputRef.current.focus();
+    }
+  }, [hasFocus]);
+
   return (
     <InputBase
-      autoFocus
+      inputRef={inputRef}
       className="MuiDataGrid-editInputCell"
       fullWidth
       type={colDef.type === 'number' ? colDef.type : 'text'}
diff --git a/packages/grid/_modules_/grid/components/cell/GridEditSingleSelectCell.tsx b/packages/grid/_modules_/grid/components/cell/GridEditSingleSelectCell.tsx
index 6891d827b443..177bb409d9f4 100644
--- a/packages/grid/_modules_/grid/components/cell/GridEditSingleSelectCell.tsx
+++ b/packages/grid/_modules_/grid/components/cell/GridEditSingleSelectCell.tsx
@@ -26,12 +26,15 @@ export function GridEditSingleSelectCell(props: GridCellParams & SelectProps) {
     colDef,
     cellMode,
     isEditable,
+    tabIndex,
     className,
     getValue,
     hasFocus,
     ...other
   } = props;
 
+  const ref = React.useRef<any>();
+
   const handleChange = (event) => {
     api.setEditCellValue({ id, field, value: event.target.value }, event);
     if (!event.key) {
@@ -46,14 +49,22 @@ export function GridEditSingleSelectCell(props: GridCellParams & SelectProps) {
     }
   };
 
+  React.useLayoutEffect(() => {
+    if (hasFocus && ref.current) {
+      // TODO v5: replace with inputRef.current.focus()
+      // See https://github.com/mui-org/material-ui/issues/21441
+      ref.current.querySelector('[role="button"]').focus();
+    }
+  }, [hasFocus]);
+
   return (
     <Select
+      ref={ref}
       value={value}
       onChange={handleChange}
       MenuProps={{
         onClose: handleClose,
       }}
-      autoFocus
       fullWidth
       open
       {...other}
diff --git a/packages/grid/_modules_/grid/components/cell/index.ts b/packages/grid/_modules_/grid/components/cell/index.ts
index 8d5f70eb20a5..dadbada0b95a 100644
--- a/packages/grid/_modules_/grid/components/cell/index.ts
+++ b/packages/grid/_modules_/grid/components/cell/index.ts
@@ -1,4 +1,5 @@
 export * from './GridCell';
 export * from './GridEditInputCell';
+export * from './GridEditSingleSelectCell';
 export * from './GridEmptyCell';
 export * from './GridRowCells';
diff --git a/packages/grid/_modules_/grid/hooks/features/focus/useGridFocus.ts b/packages/grid/_modules_/grid/hooks/features/focus/useGridFocus.ts
index 0e64a7983c44..82c18618aca3 100644
--- a/packages/grid/_modules_/grid/hooks/features/focus/useGridFocus.ts
+++ b/packages/grid/_modules_/grid/hooks/features/focus/useGridFocus.ts
@@ -1,10 +1,10 @@
 import * as React from 'react';
 import { ownerDocument } from '@material-ui/core/utils';
 import {
-  GRID_CELL_CLICK,
   GRID_CELL_DOUBLE_CLICK,
   GRID_CELL_MOUSE_UP,
   GRID_CELL_FOCUS_OUT,
+  GRID_CELL_KEY_DOWN,
   GRID_COLUMN_HEADER_BLUR,
   GRID_COLUMN_HEADER_FOCUS,
   GRID_CELL_MODE_CHANGE,
@@ -22,7 +22,7 @@ import { GridComponentProps } from '../../../GridComponentProps';
 export const useGridFocus = (apiRef: GridApiRef, props: Pick<GridComponentProps, 'rows'>): void => {
   const logger = useLogger('useGridFocus');
   const [, setGridState, forceUpdate] = useGridState(apiRef);
-  const insideFocusedCell = React.useRef(false);
+  const lastClickedCell = React.useRef<GridCellParams | null>(null);
 
   const setCellFocus = React.useCallback(
     (id: GridRowId, field: string) => {
@@ -65,13 +65,24 @@ export const useGridFocus = (apiRef: GridApiRef, props: Pick<GridComponentProps,
     [apiRef, forceUpdate, logger, setGridState],
   );
 
-  const updateFocus = React.useCallback(
+  const handleCellDoubleClick = React.useCallback(
     ({ id, field }: GridCellParams) => {
       apiRef.current.setCellFocus(id, field);
     },
     [apiRef],
   );
 
+  const handleCellKeyDown = React.useCallback(
+    (params: GridCellParams, event: React.KeyboardEvent) => {
+      // Ignore Enter and Tab because GRID_CELL_NAVIGATION_KEY_DOWN handles the focus
+      if (event.key === 'Enter' || event.key === 'Tab') {
+        return;
+      }
+      apiRef.current.setCellFocus(params.id, params.field);
+    },
+    [apiRef],
+  );
+
   const handleColumnHeaderFocus = React.useCallback(
     ({ field }: GridCellParams, event: React.FocusEvent) => {
       if (event.target !== event.currentTarget) {
@@ -90,47 +101,52 @@ export const useGridFocus = (apiRef: GridApiRef, props: Pick<GridComponentProps,
     }));
   }, [logger, setGridState]);
 
-  const handleCellMouseUp = React.useCallback(
-    (params: GridCellParams) => {
-      const { cell } = apiRef.current.getState().focus;
-      if (!cell) {
-        return;
-      }
-
-      if (params.id === cell.id && params.field === cell.field) {
-        insideFocusedCell.current = true;
-      }
-    },
-    [apiRef],
-  );
+  const handleCellMouseUp = React.useCallback((params: GridCellParams) => {
+    lastClickedCell.current = params;
+  }, []);
 
   const handleDocumentClick = React.useCallback(
-    (event: DocumentEventMap['click']) => {
-      const isInsideFocusedCell = insideFocusedCell.current;
-      insideFocusedCell.current = false;
+    (event: MouseEvent) => {
+      const cellParams = lastClickedCell.current;
+      lastClickedCell.current = null;
 
-      const { cell } = apiRef.current.getState().focus;
-      if (!cell || isInsideFocusedCell) {
+      const { cell: focusedCell } = apiRef.current.getState().focus;
+
+      if (!focusedCell) {
+        if (cellParams) {
+          apiRef.current.setCellFocus(cellParams.id, cellParams.field);
+        }
         return;
       }
 
-      const cellElement = apiRef.current.getCellElement(cell.id, cell.field);
-      if (cellElement?.contains(event.target as HTMLElement)) {
+      if (cellParams?.id === focusedCell.id && cellParams?.field === focusedCell.field) {
         return;
       }
 
-      setGridState((previousState) => ({
-        ...previousState,
-        focus: { cell: null, columnHeader: null },
-      }));
+      const cellElement = apiRef.current.getCellElement(focusedCell.id, focusedCell.field);
+      if (cellElement?.contains(event.target as HTMLElement)) {
+        return;
+      }
 
+      // There's a focused cell but another cell was clicked
+      // Publishes an event to notify that the focus was lost
       apiRef.current.publishEvent(
         GRID_CELL_FOCUS_OUT,
-        apiRef.current.getCellParams(cell.id, cell.field),
+        apiRef.current.getCellParams(focusedCell.id, focusedCell.field),
         event,
       );
+
+      if (cellParams) {
+        apiRef.current.setCellFocus(cellParams.id, cellParams.field);
+      } else {
+        setGridState((previousState) => ({
+          ...previousState,
+          focus: { cell: null, columnHeader: null },
+        }));
+        forceUpdate();
+      }
     },
-    [apiRef, setGridState],
+    [apiRef, forceUpdate, setGridState],
   );
 
   const handleCellModeChange = React.useCallback(
@@ -180,9 +196,9 @@ export const useGridFocus = (apiRef: GridApiRef, props: Pick<GridComponentProps,
   }, [apiRef, handleDocumentClick]);
 
   useGridApiEventHandler(apiRef, GRID_COLUMN_HEADER_BLUR, handleBlur);
-  useGridApiEventHandler(apiRef, GRID_CELL_CLICK, updateFocus);
-  useGridApiEventHandler(apiRef, GRID_CELL_DOUBLE_CLICK, updateFocus);
+  useGridApiEventHandler(apiRef, GRID_CELL_DOUBLE_CLICK, handleCellDoubleClick);
   useGridApiEventHandler(apiRef, GRID_CELL_MOUSE_UP, handleCellMouseUp);
+  useGridApiEventHandler(apiRef, GRID_CELL_KEY_DOWN, handleCellKeyDown);
   useGridApiEventHandler(apiRef, GRID_CELL_MODE_CHANGE, handleCellModeChange);
   useGridApiEventHandler(apiRef, GRID_COLUMN_HEADER_FOCUS, handleColumnHeaderFocus);
 };
diff --git a/packages/grid/_modules_/grid/hooks/features/rows/useGridEditRows.ts b/packages/grid/_modules_/grid/hooks/features/rows/useGridEditRows.ts
index 42271ec315ea..d0c8e7d672b2 100644
--- a/packages/grid/_modules_/grid/hooks/features/rows/useGridEditRows.ts
+++ b/packages/grid/_modules_/grid/hooks/features/rows/useGridEditRows.ts
@@ -113,6 +113,7 @@ export function useGridEditRows(
     [apiRef],
   );
 
+  // TODO it's returning undefined when colDef.editable is undefined
   const isCellEditable = React.useCallback(
     (params: GridCellParams) =>
       params.colDef.editable &&
diff --git a/packages/grid/_modules_/grid/hooks/features/rows/useGridParamsApi.ts b/packages/grid/_modules_/grid/hooks/features/rows/useGridParamsApi.ts
index a329bd570219..6d3aa4ac7f7c 100644
--- a/packages/grid/_modules_/grid/hooks/features/rows/useGridParamsApi.ts
+++ b/packages/grid/_modules_/grid/hooks/features/rows/useGridParamsApi.ts
@@ -11,7 +11,6 @@ import {
   getGridRowElement,
 } from '../../../utils/domUtils';
 import { useGridApiMethod } from '../../root/useGridApiMethod';
-import { useGridSelector } from '../core/useGridSelector';
 import { gridFocusCellSelector, gridTabIndexCellSelector } from '../focus/gridFocusStateSelector';
 
 let warnedOnce = false;
@@ -26,9 +25,6 @@ function warnMissingColumn(field) {
 }
 
 export function useGridParamsApi(apiRef: GridApiRef) {
-  const cellFocus = useGridSelector(apiRef, gridFocusCellSelector);
-  const cellTabIndex = useGridSelector(apiRef, gridTabIndexCellSelector);
-
   const getColumnHeaderParams = React.useCallback(
     (field: string): GridColumnHeaderParams => ({
       field,
@@ -66,6 +62,9 @@ export function useGridParamsApi(apiRef: GridApiRef) {
         throw new Error(`No row with id #${id} found`);
       }
 
+      const cellFocus = gridFocusCellSelector(apiRef.current.getState());
+      const cellTabIndex = gridTabIndexCellSelector(apiRef.current.getState());
+
       const params: GridValueGetterParams = {
         id,
         field,
@@ -81,7 +80,7 @@ export function useGridParamsApi(apiRef: GridApiRef) {
 
       return params;
     },
-    [apiRef, cellFocus, cellTabIndex],
+    [apiRef],
   );
 
   const getCellParams = React.useCallback(
@@ -94,6 +93,9 @@ export function useGridParamsApi(apiRef: GridApiRef) {
         throw new Error(`No row with id #${id} found`);
       }
 
+      const cellFocus = gridFocusCellSelector(apiRef.current.getState());
+      const cellTabIndex = gridTabIndexCellSelector(apiRef.current.getState());
+
       const params: GridCellParams = {
         id,
         field,
@@ -114,7 +116,7 @@ export function useGridParamsApi(apiRef: GridApiRef) {
 
       return params;
     },
-    [apiRef, cellFocus, cellTabIndex],
+    [apiRef],
   );
 
   const getCellValue = React.useCallback(
diff --git a/packages/grid/data-grid/src/tests/keyboard.DataGrid.test.tsx b/packages/grid/data-grid/src/tests/keyboard.DataGrid.test.tsx
index 9454e6031961..66c4b9f189b6 100644
--- a/packages/grid/data-grid/src/tests/keyboard.DataGrid.test.tsx
+++ b/packages/grid/data-grid/src/tests/keyboard.DataGrid.test.tsx
@@ -93,6 +93,7 @@ describe('<DataGrid /> - Keyboard', () => {
     const input = screen.getByTestId('custom-input');
     input.focus();
     expect(getActiveCell()).to.equal('0-1');
+    input.focus(); // The focus moves back to the cell
     fireEvent.keyDown(input, { key: 'ArrowLeft' });
     expect(getActiveCell()).to.equal('0-1');
   });
diff --git a/packages/grid/x-grid/src/tests/editRows.XGrid.test.tsx b/packages/grid/x-grid/src/tests/editRows.XGrid.test.tsx
index 42f4a7e135ec..fe2dc100d955 100644
--- a/packages/grid/x-grid/src/tests/editRows.XGrid.test.tsx
+++ b/packages/grid/x-grid/src/tests/editRows.XGrid.test.tsx
@@ -4,6 +4,7 @@ import {
   GridComponentProps,
   useGridApiRef,
   XGrid,
+  GridEditSingleSelectCell,
 } from '@material-ui/x-grid';
 import Portal from '@material-ui/unstyled/Portal';
 import { expect } from 'chai';
@@ -207,7 +208,7 @@ describe('<XGrid /> - Edit Rows', () => {
     expect(cell).to.have.text('Adidas');
   });
 
-  it('should allow to save an edit changes using Enter', () => {
+  it('should allow to save changes using Enter', () => {
     render(<TestCase />);
     const cell = getCell(1, 0);
     cell.focus();
@@ -273,6 +274,7 @@ describe('<XGrid /> - Edit Rows', () => {
     expect(cell.querySelector('input')!.value).to.equal('1970');
 
     const otherCell = getCell(2, 1);
+    fireEvent.mouseUp(otherCell);
     fireEvent.click(otherCell);
     fireEvent.focus(otherCell);
     expect(cell).not.to.have.class('MuiDataGrid-cell--editing');
@@ -318,8 +320,9 @@ describe('<XGrid /> - Edit Rows', () => {
   it('should work correctly when the cell editing was initiated programmatically', () => {
     render(<TestCase />);
     apiRef.current.setCellMode(1, 'year', 'edit');
-    expect(getActiveCell()).to.equal('1-1');
     const cell = getCell(1, 1);
+    cell.focus();
+    expect(getActiveCell()).to.equal('1-1');
     const input = cell.querySelector('input')!;
     expect(input.value).to.equal('1961');
 
@@ -327,6 +330,7 @@ describe('<XGrid /> - Edit Rows', () => {
     expect(cell.querySelector('input')!.value).to.equal('1970');
 
     const otherCell = getCell(2, 1);
+    fireEvent.mouseUp(otherCell);
     fireEvent.click(otherCell);
     fireEvent.focus(otherCell);
     expect(cell).not.to.have.class('MuiDataGrid-cell--editing');
@@ -457,6 +461,13 @@ describe('<XGrid /> - Edit Rows', () => {
     });
   });
 
+  it('should set the focus correctly', () => {
+    render(<TestCase />);
+    const cell = getCell(0, 0);
+    fireEvent.doubleClick(cell);
+    expect(document.activeElement).to.equal(screen.getByRole('textbox'));
+  });
+
   describe('column type: singleSelect', () => {
     it('should change cell value correctly when the valueOptions is array of strings', () => {
       render(
@@ -532,6 +543,33 @@ describe('<XGrid /> - Edit Rows', () => {
       expect(cell).not.to.have.class('MuiDataGrid-cell--editing');
       expect(cell).to.have.text('Italy');
     });
+
+    it('should set the focus correctly', () => {
+      render(
+        <TestCase
+          columns={[
+            {
+              field: 'brand',
+              type: 'singleSelect',
+              valueOptions: ['Nike', 'Adidas'],
+              editable: true,
+              renderEditCell: (params: any) => (
+                <GridEditSingleSelectCell {...params} open={false} /> // Force to appear closed
+              ),
+            },
+          ]}
+          rows={[
+            {
+              id: 0,
+              brand: 'Nike',
+            },
+          ]}
+        />,
+      );
+      const cell = getCell(0, 0);
+      fireEvent.doubleClick(cell);
+      expect(document.activeElement).to.equal(screen.getByRole('button', { name: 'Nike' }));
+    });
   });
 
   it('should keep the right type', () => {
@@ -599,6 +637,18 @@ describe('<XGrid /> - Edit Rows', () => {
       fireEvent.change(input, { target: { value: '' } });
       expect(onEditCellPropsChange.args[0][0].props.value).to.equal(null);
     });
+
+    it('should set the focus correctly', () => {
+      render(
+        <TestCase
+          rows={[{ id: 0, date: new Date(2021, 6, 5) }]}
+          columns={[{ field: 'date', type: 'date', editable: true }]}
+        />,
+      );
+      const cell = getCell(0, 0);
+      fireEvent.doubleClick(cell);
+      expect(document.activeElement).to.equal(screen.getByRole('cell').querySelector('input'));
+    });
   });
 
   describe('column type: dateTime', () => {
@@ -637,6 +687,18 @@ describe('<XGrid /> - Edit Rows', () => {
       fireEvent.change(input, { target: { value: '' } });
       expect(onEditCellPropsChange.args[0][0].props.value).to.equal(null);
     });
+
+    it('should set the focus correctly', () => {
+      render(
+        <TestCase
+          rows={[{ id: 0, date: new Date(2021, 6, 5, 14, 30) }]}
+          columns={[{ field: 'date', type: 'dateTime', editable: true }]}
+        />,
+      );
+      const cell = getCell(0, 0);
+      fireEvent.doubleClick(cell);
+      expect(document.activeElement).to.equal(screen.getByRole('cell').querySelector('input'));
+    });
   });
 
   it('should call onCellEditCommit with the correct params', () => {
@@ -673,6 +735,18 @@ describe('<XGrid /> - Edit Rows', () => {
         props: { value: true },
       });
     });
+
+    it('should set the focus correctly', () => {
+      render(
+        <TestCase
+          rows={[{ id: 0, isAdmin: false }]}
+          columns={[{ field: 'isAdmin', type: 'boolean', editable: true }]}
+        />,
+      );
+      const cell = getCell(0, 0);
+      fireEvent.doubleClick(cell);
+      expect(document.activeElement).to.equal(screen.getByRole('checkbox'));
+    });
   });
 
   describe('validation', () => {
@@ -738,6 +812,7 @@ describe('<XGrid /> - Edit Rows', () => {
       render(<TestCase editRowsModel={{ 1: { year: { value: 1961 } } }} />);
       const cell = getCell(1, 1);
       const input = cell.querySelector('input')!;
+      input.focus();
       expect(input.value).to.equal('1961');
       fireEvent.change(input, { target: { value: '1970' } });
       fireEvent.keyDown(input, { key: 'Enter' });
@@ -776,6 +851,7 @@ describe('<XGrid /> - Edit Rows', () => {
       );
       const cell = getCell(1, 1);
       const input = cell.querySelector('input')!;
+      input.focus();
       fireEvent.change(input, { target: { value: 1970 } });
       expect(onEditRowsModelChange.lastCall.firstArg).to.deep.equal({
         1: { year: { value: '1970' } },
diff --git a/packages/grid/x-grid/src/tests/rows.XGrid.test.tsx b/packages/grid/x-grid/src/tests/rows.XGrid.test.tsx
index 0b93a31b1296..83b1acb7c44e 100644
--- a/packages/grid/x-grid/src/tests/rows.XGrid.test.tsx
+++ b/packages/grid/x-grid/src/tests/rows.XGrid.test.tsx
@@ -4,7 +4,7 @@ import {
   // @ts-expect-error need to migrate helpers to TypeScript
   fireEvent,
 } from 'test/utils';
-import { useFakeTimers } from 'sinon';
+import { useFakeTimers, spy } from 'sinon';
 import { expect } from 'chai';
 import { getCell, getColumnValues } from 'test/utils/helperFn';
 import {
@@ -14,6 +14,7 @@ import {
   useGridApiRef,
   XGrid,
   XGridProps,
+  GRID_CELL_FOCUS_OUT,
 } from '@material-ui/x-grid';
 import { useData } from 'packages/storybook/src/hooks/useData';
 
@@ -597,6 +598,7 @@ describe('<XGrid /> - Rows', () => {
     it('should focus the clicked cell in the state', () => {
       render(<TestCase rows={baselineProps.rows} />);
 
+      fireEvent.mouseUp(getCell(0, 0));
       fireEvent.click(getCell(0, 0));
       expect(apiRef.current.getState().focus.cell).to.deep.equal({
         id: baselineProps.rows[0].id,
@@ -615,6 +617,7 @@ describe('<XGrid /> - Rows', () => {
     it('should not reset focus when removing a row not containing the focus cell', () => {
       const { setProps } = render(<TestCase rows={baselineProps.rows} />);
 
+      fireEvent.mouseUp(getCell(1, 0));
       fireEvent.click(getCell(1, 0));
       setProps({ rows: baselineProps.rows.slice(1) });
       expect(apiRef.current.getState().focus.cell).to.deep.equal({
@@ -622,5 +625,57 @@ describe('<XGrid /> - Rows', () => {
         field: baselineProps.columns[0].field,
       });
     });
+
+    it('should set the focus when pressing a key inside a cell', () => {
+      render(<TestCase rows={baselineProps.rows} />);
+      const cell = getCell(1, 0);
+      cell.focus();
+      fireEvent.keyDown(cell, { key: 'a' });
+      expect(apiRef.current.getState().focus.cell).to.deep.equal({
+        id: baselineProps.rows[1].id,
+        field: baselineProps.columns[0].field,
+      });
+    });
+
+    it('should update the focus when clicking from one cell to another', () => {
+      render(<TestCase rows={baselineProps.rows} />);
+      fireEvent.mouseUp(getCell(1, 0));
+      fireEvent.click(getCell(1, 0));
+      expect(apiRef.current.getState().focus.cell).to.deep.equal({
+        id: baselineProps.rows[1].id,
+        field: baselineProps.columns[0].field,
+      });
+      fireEvent.mouseUp(getCell(2, 1));
+      fireEvent.click(getCell(2, 1));
+      expect(apiRef.current.getState().focus.cell).to.deep.equal({
+        id: baselineProps.rows[2].id,
+        field: baselineProps.columns[1].field,
+      });
+    });
+
+    it('should reset focus when clicking outside the focused cell', () => {
+      render(<TestCase rows={baselineProps.rows} />);
+      fireEvent.mouseUp(getCell(1, 0));
+      fireEvent.click(getCell(1, 0));
+      expect(apiRef.current.getState().focus.cell).to.deep.equal({
+        id: baselineProps.rows[1].id,
+        field: baselineProps.columns[0].field,
+      });
+      fireEvent.click(document.body);
+      expect(apiRef.current.getState().focus.cell).to.deep.equal(null);
+    });
+
+    it('should publish GRID_CELL_FOCUS_OUT when clicking outside the focused cell', () => {
+      const onCellFocusOut = spy();
+      render(<TestCase rows={baselineProps.rows} />);
+      apiRef.current.subscribeEvent(GRID_CELL_FOCUS_OUT, onCellFocusOut);
+      fireEvent.mouseUp(getCell(1, 0));
+      fireEvent.click(getCell(1, 0));
+      expect(onCellFocusOut.callCount).to.equal(0);
+      fireEvent.click(document.body);
+      expect(onCellFocusOut.callCount).to.equal(1);
+      expect(onCellFocusOut.args[0][0].id).to.equal(baselineProps.rows[1].id);
+      expect(onCellFocusOut.args[0][0].field).to.equal(baselineProps.columns[0].field);
+    });
   });
 });
diff --git a/packages/storybook/src/stories/grid-rows.stories.tsx b/packages/storybook/src/stories/grid-rows.stories.tsx
index 9a68be194355..46b4c13eaf56 100644
--- a/packages/storybook/src/stories/grid-rows.stories.tsx
+++ b/packages/storybook/src/stories/grid-rows.stories.tsx
@@ -61,7 +61,10 @@ const baselineProps = {
       isPublished: true,
     },
   ],
-  columns: [{ field: 'brand' }, { field: 'isPublished', type: 'boolean' }],
+  columns: [
+    { field: 'brand', editable: true },
+    { field: 'isPublished', type: 'boolean' },
+  ],
 };
 
 function getStoryRowId(row) {
@@ -379,6 +382,7 @@ const baselineEditProps = {
       DOB: new Date(1996, 10, 2),
       meetup: new Date(2020, 2, 25, 10, 50, 0),
       isAdmin: true,
+      country: 'Spain',
     },
     {
       id: 1,
@@ -391,6 +395,7 @@ const baselineEditProps = {
       DOB: new Date(1992, 1, 20),
       meetup: new Date(2020, 4, 15, 10, 50, 0),
       isAdmin: true,
+      country: 'Netherlands',
     },
     {
       id: 2,
@@ -403,6 +408,7 @@ const baselineEditProps = {
       DOB: new Date(1986, 0, 12),
       meetup: new Date(2020, 3, 5, 10, 50, 0),
       isAdmin: false,
+      country: 'Brazil',
     },
   ],
   columns: [
@@ -414,6 +420,13 @@ const baselineEditProps = {
       valueGetter: ({ row }) => `${row.firstname || ''} ${row.lastname || ''}`,
     },
     { field: 'isAdmin', width: 120, type: 'boolean', editable: true },
+    {
+      field: 'country',
+      width: 120,
+      type: 'singleSelect',
+      editable: true,
+      valueOptions: ['Bulgaria', 'Netherlands', 'France', 'Italy', 'Brazil', 'Spain'],
+    },
     { field: 'username', editable: true },
     { field: 'email', editable: true, width: 150 },
     { field: 'age', width: 50, type: 'number', editable: true },