[DataGridPremium] Update column state correctly when grouping mode is updated with one grouping column (#17069)

Author: arminmeh

packages/x-data-grid-premium/src/hooks/features/rowGrouping/useGridRowGroupingPreProcessors.ts

@@ -117,13 +117,10 @@ export const useGridRowGroupingPreProcessors = (
       const groupingColDefs = getGroupingColDefs(columnsState);
       let newColumnFields: string[] = [];
       const newColumnsLookup: GridColumnRawLookup = {};
-      const prevGroupingfields: string[] = [];
 
       // We only keep the non-grouping columns
       columnsState.orderedFields.forEach((field) => {
-        if (isGroupingColumn(field)) {
-          prevGroupingfields.push(field);
-        } else {
+        if (!isGroupingColumn(field)) {
           newColumnFields.push(field);
           newColumnsLookup[field] = columnsState.lookup[field];
         }
@@ -140,16 +137,19 @@ export const useGridRowGroupingPreProcessors = (
         newColumnsLookup[groupingColDef.field] = groupingColDef;
       });
 
-      if (prevGroupingfields.length !== groupingColDefs.length) {
-        const startIndex = newColumnFields[0] === GRID_CHECKBOX_SELECTION_FIELD ? 1 : 0;
-        newColumnFields = [
-          ...newColumnFields.slice(0, startIndex),
-          ...groupingColDefs.map((colDef) => colDef.field),
-          ...newColumnFields.slice(startIndex),
-        ];
-        columnsState.orderedFields = newColumnFields;
-      }
+      const checkBoxFieldIndex = newColumnFields.findIndex(
+        (field) => field === GRID_CHECKBOX_SELECTION_FIELD,
+      );
+      const checkBoxColumn =
+        checkBoxFieldIndex !== -1 ? newColumnFields.splice(checkBoxFieldIndex, 1) : [];
+
+      newColumnFields = [
+        ...checkBoxColumn,
+        ...groupingColDefs.map((colDef) => colDef.field),
+        ...newColumnFields,
+      ];
 
+      columnsState.orderedFields = newColumnFields;
       columnsState.lookup = newColumnsLookup;
 
       return columnsState;

packages/x-data-grid-premium/src/tests/rowGrouping.DataGridPremium.test.tsx

@@ -468,6 +468,34 @@ describe('<DataGridPremium /> - Row grouping', () => {
       ]);
     });
 
+    // https://github.com/mui/mui-x/issues/17046
+    it('should support rowGroupingColumnMode switch with one grouping column', () => {
+      const { setProps } = render(
+        <Test
+          rowGroupingModel={['category1']}
+          rowGroupingColumnMode="multiple"
+          defaultGroupingExpansionDepth={-1}
+        />,
+      );
+
+      expect(getColumnHeadersTextContent()).to.deep.equal([
+        'category1',
+        'id',
+        'category1',
+        'category2',
+      ]);
+      expect(getColumnValues(0)).to.deep.equal(['Cat A (3)', '', '', '', 'Cat B (2)', '', '']);
+
+      setProps({ rowGroupingColumnMode: 'single' });
+      expect(getColumnHeadersTextContent()).to.deep.equal([
+        'category1',
+        'id',
+        'category1',
+        'category2',
+      ]);
+      expect(getColumnValues(0)).to.deep.equal(['Cat A (3)', '', '', '', 'Cat B (2)', '', '']);
+    });
+
     it('should respect the model grouping order when rowGroupingColumnMode = "single"', () => {
       render(
         <Test