From c9151bd4fd655c20d347c09ecfccd3bcc5392c0f Mon Sep 17 00:00:00 2001
From: Olivier Tassinari <olivier.tassinari@gmail.com>
Date: Thu, 17 Jun 2021 02:18:52 +0200
Subject: [PATCH] [DataGrid] Fix double-click issue

---
 .../columnSelection/GridHeaderCheckbox.tsx    | 18 ++--------
 .../x-grid/src/tests/selection.XGrid.test.tsx | 35 ++++++++++++++++++-
 2 files changed, 37 insertions(+), 16 deletions(-)

diff --git a/packages/grid/_modules_/grid/components/columnSelection/GridHeaderCheckbox.tsx b/packages/grid/_modules_/grid/components/columnSelection/GridHeaderCheckbox.tsx
index c1a1041db4538..3de90aa85c00a 100644
--- a/packages/grid/_modules_/grid/components/columnSelection/GridHeaderCheckbox.tsx
+++ b/packages/grid/_modules_/grid/components/columnSelection/GridHeaderCheckbox.tsx
@@ -23,24 +23,12 @@ export const GridHeaderCheckbox = React.forwardRef<HTMLInputElement, GridColumnH
     const totalSelectedRows = useGridSelector(apiRef, selectedGridRowsCountSelector);
     const totalRows = useGridSelector(apiRef, gridRowCountSelector);
 
-    const [isIndeterminate, setIsIndeterminate] = React.useState(
-      totalSelectedRows > 0 && totalSelectedRows !== totalRows,
-    );
-    const [isChecked, setChecked] = React.useState(
-      totalSelectedRows === totalRows || isIndeterminate,
-    );
-
-    React.useEffect(() => {
-      const isNewIndeterminate = totalSelectedRows > 0 && totalSelectedRows !== totalRows;
-      const isNewChecked = (totalRows > 0 && totalSelectedRows === totalRows) || isIndeterminate;
-      setChecked(isNewChecked);
-      setIsIndeterminate(isNewIndeterminate);
-    }, [isIndeterminate, totalRows, totalSelectedRows]);
+    const isIndeterminate = totalSelectedRows > 0 && totalSelectedRows !== totalRows;
+    const isChecked = totalSelectedRows === totalRows || isIndeterminate;
 
     const handleChange = (event: React.ChangeEvent<HTMLInputElement>) => {
       const checked = event.target.checked;
-      setChecked(checked);
-      apiRef!.current.selectRows(visibleRowIds, checked);
+      apiRef!.current.selectRows(visibleRowIds, checked, !event.target.indeterminate);
     };
 
     const tabIndex = tabIndexState !== null && tabIndexState.field === props.field ? 0 : -1;
diff --git a/packages/grid/x-grid/src/tests/selection.XGrid.test.tsx b/packages/grid/x-grid/src/tests/selection.XGrid.test.tsx
index 20ad02ddd45c1..b82504e742791 100644
--- a/packages/grid/x-grid/src/tests/selection.XGrid.test.tsx
+++ b/packages/grid/x-grid/src/tests/selection.XGrid.test.tsx
@@ -1,7 +1,14 @@
 import * as React from 'react';
-import { createClientRenderStrictMode } from 'test/utils';
 import { expect } from 'chai';
 import { spy } from 'sinon';
+import { getColumnValues } from 'test/utils/helperFn';
+import {
+  // @ts-expect-error need to migrate helpers to TypeScript
+  screen,
+  createClientRenderStrictMode,
+  // @ts-expect-error need to migrate helpers to TypeScript
+  fireEvent,
+} from 'test/utils';
 import { GridApiRef, GridComponentProps, useGridApiRef, XGrid } from '@material-ui/x-grid';
 
 const isJSDOM = /jsdom/.test(window.navigator.userAgent);
@@ -126,4 +133,30 @@ describe('<XGrid /> - Selection', () => {
     });
     expect(apiRef.current.getSelectedRows()).to.have.keys([0]);
   });
+
+  it('should select only filtered rows after filter is applied', () => {
+    render(<Test checkboxSelection />);
+    const selectAll = screen.getByRole('checkbox', {
+      name: /select all rows checkbox/i,
+    });
+    apiRef.current.setFilterModel({
+      items: [
+        {
+          columnField: 'brand',
+          operatorValue: 'contains',
+          value: 'Puma',
+        },
+      ],
+    });
+    expect(getColumnValues(1)).to.deep.equal(['Puma']);
+    fireEvent.click(selectAll);
+    // TODO fix, should be only 2
+    expect(Array.from(apiRef.current.getSelectedRows().keys())).to.deep.equal([0, 1, 2]);
+    fireEvent.click(selectAll);
+    expect(Array.from(apiRef.current.getSelectedRows().keys())).to.deep.equal([]);
+    fireEvent.click(selectAll);
+    expect(Array.from(apiRef.current.getSelectedRows().keys())).to.deep.equal([2]);
+    fireEvent.click(selectAll);
+    expect(Array.from(apiRef.current.getSelectedRows().keys())).to.deep.equal([]);
+  });
 });