Fix Grid header adding and removing with select all checkbox (#17592)

Change-Id: I4306635ca334429ef006216c5b35165cf89a938f

3 files changed, 104 insertions, 100 deletions

diff --git a/client/src/com/vaadin/client/connectors/GridConnector.java b/client/src/com/vaadin/client/connectors/GridConnector.java
index f2246b1e5f..b720dc442f 100644
--- a/client/src/com/vaadin/client/connectors/GridConnector.java
+++ b/client/src/com/vaadin/client/connectors/GridConnector.java
@@ -1028,6 +1028,10 @@ public class GridConnector extends AbstractHasComponentsConnector implements
         for (RowState rowState : state.rows) {
             HeaderRow row = getWidget().appendHeaderRow();
 
+            if (rowState.defaultRow) {
+                getWidget().setDefaultHeaderRow(row);
+            }
+
             for (CellState cellState : rowState.cells) {
                 CustomGridColumn column = columnIdToColumn
                         .get(cellState.columnId);
@@ -1049,10 +1053,6 @@ public class GridConnector extends AbstractHasComponentsConnector implements
                 updateHeaderCellFromState(row.join(columns), cellState);
             }
 
-            if (rowState.defaultRow) {
-                getWidget().setDefaultHeaderRow(row);
-            }
-
             row.setStyleName(rowState.styleName);
         }
     }
diff --git a/client/src/com/vaadin/client/widgets/Grid.java b/client/src/com/vaadin/client/widgets/Grid.java
index 56b081f6b3..e47418a740 100644
--- a/client/src/com/vaadin/client/widgets/Grid.java
+++ b/client/src/com/vaadin/client/widgets/Grid.java
@@ -861,6 +861,7 @@ public class Grid<T> extends ResizeComposite implements
             if (row != null) {
                 row.setDefault(true);
             }
+
             defaultRow = row;
             requestSectionRefresh();
         }
@@ -911,6 +912,16 @@ public class Grid<T> extends ResizeComposite implements
                     BrowserEvents.TOUCHMOVE, BrowserEvents.TOUCHEND,
                     BrowserEvents.TOUCHCANCEL, BrowserEvents.CLICK);
         }
+
+        @Override
+        protected void addColumn(Column<?, ?> column) {
+            super.addColumn(column);
+
+            // Add default content for new columns.
+            if (defaultRow != null) {
+                column.setDefaultHeaderContent(defaultRow.getCell(column));
+            }
+        }
     }
 
     /**
@@ -923,6 +934,11 @@ public class Grid<T> extends ResizeComposite implements
 
         protected void setDefault(boolean isDefault) {
             this.isDefault = isDefault;
+            if (isDefault) {
+                for (Column<?, ?> column : getSection().grid.getColumns()) {
+                    column.setDefaultHeaderContent(getCell(column));
+                }
+            }
         }
 
         public boolean isDefault() {
@@ -2341,8 +2357,6 @@ public class Grid<T> extends ResizeComposite implements
         }
 
         void initDone() {
-            addSelectAllToDefaultHeader();
-
             setWidth(-1);
 
             setEditable(false);
@@ -2350,62 +2364,52 @@ public class Grid<T> extends ResizeComposite implements
             initDone = true;
         }
 
-        protected void addSelectAllToDefaultHeader() {
-            if (getSelectionModel() instanceof SelectionModel.Multi
-                    && header.getDefaultRow() != null) {
-                // If selection cell already contains a widget do not
-                // create a new CheckBox
-                HeaderCell selectionCell = header.getDefaultRow().getCell(this);
-                if (selectionCell.getType().equals(GridStaticCellType.WIDGET)
-                        && selectionCell.getWidget() instanceof CheckBox) {
-                    return;
-                }
-                /*
-                 * TODO: Currently the select all check box is shown when multi
-                 * selection is in use. This might result in malfunctions if no
-                 * SelectAllHandlers are present.
-                 * 
-                 * Later on this could be fixed so that it check such handlers
-                 * exist.
-                 */
-                final SelectionModel.Multi<T> model = (Multi<T>) getSelectionModel();
-                final CheckBox checkBox = GWT.create(CheckBox.class);
-                checkBox.addValueChangeHandler(new ValueChangeHandler<Boolean>() {
+        @Override
+        protected void setDefaultHeaderContent(HeaderCell selectionCell) {
+            /*
+             * TODO: Currently the select all check box is shown when multi
+             * selection is in use. This might result in malfunctions if no
+             * SelectAllHandlers are present.
+             * 
+             * Later on this could be fixed so that it check such handlers
+             * exist.
+             */
+            final SelectionModel.Multi<T> model = (Multi<T>) getSelectionModel();
+            final CheckBox checkBox = GWT.create(CheckBox.class);
+            checkBox.addValueChangeHandler(new ValueChangeHandler<Boolean>() {
 
-                    @Override
-                    public void onValueChange(ValueChangeEvent<Boolean> event) {
-                        if (event.getValue()) {
-                            fireEvent(new SelectAllEvent<T>(model));
-                            selected = true;
-                        } else {
-                            model.deselectAll();
-                            selected = false;
-                        }
+                @Override
+                public void onValueChange(ValueChangeEvent<Boolean> event) {
+                    if (event.getValue()) {
+                        fireEvent(new SelectAllEvent<T>(model));
+                        selected = true;
+                    } else {
+                        model.deselectAll();
+                        selected = false;
                     }
-                });
-                checkBox.setValue(selected);
-                selectionCell.setWidget(checkBox);
-                // Select all with space when "select all" cell is active
-                addHeaderKeyUpHandler(new HeaderKeyUpHandler() {
-                    @Override
-                    public void onKeyUp(GridKeyUpEvent event) {
-                        if (event.getNativeKeyCode() != KeyCodes.KEY_SPACE) {
-                            return;
-                        }
-                        HeaderRow targetHeaderRow = getHeader().getRow(
-                                event.getFocusedCell().getRowIndex());
-                        if (!targetHeaderRow.isDefault()) {
-                            return;
-                        }
-                        if (event.getFocusedCell().getColumn() == SelectionColumn.this) {
-                            // Send events to ensure row selection state is
-                            // updated
-                            checkBox.setValue(!checkBox.getValue(), true);
-                        }
+                }
+            });
+            checkBox.setValue(selected);
+            selectionCell.setWidget(checkBox);
+            // Select all with space when "select all" cell is active
+            addHeaderKeyUpHandler(new HeaderKeyUpHandler() {
+                @Override
+                public void onKeyUp(GridKeyUpEvent event) {
+                    if (event.getNativeKeyCode() != KeyCodes.KEY_SPACE) {
+                        return;
                     }
-                });
-
-            }
+                    HeaderRow targetHeaderRow = getHeader().getRow(
+                            event.getFocusedCell().getRowIndex());
+                    if (!targetHeaderRow.isDefault()) {
+                        return;
+                    }
+                    if (event.getFocusedCell().getColumn() == SelectionColumn.this) {
+                        // Send events to ensure row selection state is
+                        // updated
+                        checkBox.setValue(!checkBox.getValue(), true);
+                    }
+                }
+            });
         }
 
         @Override
@@ -4204,7 +4208,6 @@ public class Grid<T> extends ResizeComposite implements
             this.grid = grid;
             if (this.grid != null) {
                 this.grid.recalculateColumnWidths();
-                updateHeader();
             }
         }
 
@@ -4742,6 +4745,17 @@ public class Grid<T> extends ResizeComposite implements
                  */
             }
         }
+
+        /**
+         * Resets the default header cell contents to column header captions.
+         * 
+         * @since
+         * @param cell
+         *            default header cell for this column
+         */
+        protected void setDefaultHeaderContent(HeaderCell cell) {
+            cell.setText(headerCaption);
+        }
     }
 
     protected class BodyUpdater implements EscalatorUpdater {
@@ -4995,35 +5009,29 @@ public class Grid<T> extends ResizeComposite implements
                 final StaticSection.StaticCell metadata = staticRow
                         .getCell(columns.get(cell.getColumn()));
 
-                boolean updateCellData = true;
                 // Decorate default row with sorting indicators
                 if (staticRow instanceof HeaderRow) {
                     addSortingIndicatorsToHeaderRow((HeaderRow) staticRow, cell);
-
-                    if (isHeaderSelectionColumn(row, cell)) {
-                        updateCellData = false;
-                    }
                 }
 
                 // Assign colspan to cell before rendering
                 cell.setColSpan(metadata.getColspan());
 
                 TableCellElement element = cell.getElement();
-                if (updateCellData) {
-                    switch (metadata.getType()) {
-                    case TEXT:
-                        element.setInnerText(metadata.getText());
-                        break;
-                    case HTML:
-                        element.setInnerHTML(metadata.getHtml());
-                        break;
-                    case WIDGET:
-                        preDetach(row, Arrays.asList(cell));
-                        element.setInnerHTML("");
-                        postAttach(row, Arrays.asList(cell));
-                        break;
-                    }
+                switch (metadata.getType()) {
+                case TEXT:
+                    element.setInnerText(metadata.getText());
+                    break;
+                case HTML:
+                    element.setInnerHTML(metadata.getHtml());
+                    break;
+                case WIDGET:
+                    preDetach(row, Arrays.asList(cell));
+                    element.setInnerHTML("");
+                    postAttach(row, Arrays.asList(cell));
+                    break;
                 }
+
                 setCustomStyleName(element, metadata.getStyleName());
 
                 cellFocusHandler.updateFocusedCellStyle(cell, container);
@@ -5082,27 +5090,6 @@ public class Grid<T> extends ResizeComposite implements
 
         @Override
         public void preAttach(Row row, Iterable<FlyweightCell> cellsToAttach) {
-            // Add select all checkbox if needed on rebuild.
-            for (FlyweightCell cell : cellsToAttach) {
-                if (isHeaderSelectionColumn(row, cell)) {
-                    selectionColumn.addSelectAllToDefaultHeader();
-                }
-            }
-        }
-
-        /**
-         * Check if selectionColumn in the default header row
-         */
-        private boolean isHeaderSelectionColumn(Row row, FlyweightCell cell) {
-            return selectionColumn != null && isDefaultHeaderRow(row)
-                    && getColumn(cell.getColumn()).equals(selectionColumn);
-        }
-
-        /**
-         * Row is the default header row.
-         */
-        private boolean isDefaultHeaderRow(Row row) {
-            return section.getRow(row.getRow()).equals(header.getDefaultRow());
         }
 
         @Override
@@ -5731,6 +5718,9 @@ public class Grid<T> extends ResizeComposite implements
     /**
      * Sets the default row of the header. The default row is a special header
      * row providing a user interface for sorting columns.
+     * <p>
+     * Note: Setting the default header row will reset all cell contents to
+     * Column defaults.
      * 
      * @param row
      *            the new default row, or null for no default row
diff --git a/uitest/src/com/vaadin/tests/components/grid/basicfeatures/server/GridSelectionTest.java b/uitest/src/com/vaadin/tests/components/grid/basicfeatures/server/GridSelectionTest.java
index b4eb473d4b..74972fe4d2 100644
--- a/uitest/src/com/vaadin/tests/components/grid/basicfeatures/server/GridSelectionTest.java
+++ b/uitest/src/com/vaadin/tests/components/grid/basicfeatures/server/GridSelectionTest.java
@@ -271,6 +271,20 @@ public class GridSelectionTest extends GridBasicFeaturesTest {
     }
 
     @Test
+    public void testSelectAllCheckboxWithHeaderOperations() {
+        openTestURL();
+
+        setSelectionModelMulti();
+        selectMenuPath("Component", "Header", "Prepend row");
+        selectMenuPath("Component", "Header", "Append row");
+
+        GridCellElement header = getGridElement().getHeaderCell(1, 0);
+        assertTrue("Multi Selection Model should have select all checkbox",
+                header.isElementPresent(By.tagName("input")));
+
+    }
+
+    @Test
     public void testToggleDeselectAllowed() {
         openTestURL();