Fix race condition between rendering TreeGrid and setting up headers (#9105)

* Fix race condition between rendering Grid and setting up headers
* Remove Grid default headers when initializing the widget
* Fix header clean up, add null check to select all checkbox

4 files changed, 59 insertions, 16 deletions

diff --git a/client/src/main/java/com/vaadin/client/connectors/grid/GridConnector.java b/client/src/main/java/com/vaadin/client/connectors/grid/GridConnector.java
index 0181980ef8..5ede13ca68 100644
--- a/client/src/main/java/com/vaadin/client/connectors/grid/GridConnector.java
+++ b/client/src/main/java/com/vaadin/client/connectors/grid/GridConnector.java
@@ -202,6 +202,11 @@ public class GridConnector extends AbstractListingConnector
     protected void init() {
         super.init();
 
+        // Remove default headers when initializing Grid widget
+        while (getWidget().getHeaderRowCount() > 0) {
+            getWidget().removeHeaderRow(0);
+        }
+
         registerRpc(GridClientRpc.class, new GridClientRpc() {
 
             @Override
@@ -308,24 +313,22 @@ public class GridConnector extends AbstractListingConnector
      */
     @OnStateChange("header")
     void updateHeader() {
-        Scheduler.get().scheduleFinally(() -> {
-            final Grid<JsonObject> grid = getWidget();
-            final SectionState state = getState().header;
-
-            while (grid.getHeaderRowCount() > 0) {
-                grid.removeHeaderRow(0);
-            }
+        final Grid<JsonObject> grid = getWidget();
+        final SectionState state = getState().header;
 
-            for (RowState rowState : state.rows) {
-                HeaderRow row = grid.appendHeaderRow();
+        while (grid.getHeaderRowCount() > 0) {
+            grid.removeHeaderRow(0);
+        }
 
-                if (rowState.defaultHeader) {
-                    grid.setDefaultHeaderRow(row);
-                }
+        for (RowState rowState : state.rows) {
+            HeaderRow row = grid.appendHeaderRow();
 
-                updateStaticRow(rowState, row);
+            if (rowState.defaultHeader) {
+                grid.setDefaultHeaderRow(row);
             }
-        });
+
+            updateStaticRow(rowState, row);
+        }
     }
 
     @OnStateChange("rowHeight")
diff --git a/client/src/main/java/com/vaadin/client/widgets/Grid.java b/client/src/main/java/com/vaadin/client/widgets/Grid.java
index cb50dd5977..27db38c4a1 100755
--- a/client/src/main/java/com/vaadin/client/widgets/Grid.java
+++ b/client/src/main/java/com/vaadin/client/widgets/Grid.java
@@ -3045,8 +3045,10 @@ public class Grid<T> extends ResizeComposite implements HasSelectionHandlers<T>,
          * Sets the select all checkbox visible or hidden.
          */
         protected void doSetSelectAllCheckBoxVisible() {
-            assert selectAllCheckBox != null : "Select All Checkbox has not been created for selection column.";
-            assert selectionCell != null : "Default header cell for selection column not been set.";
+            if (selectAllCheckBox == null || selectionCell == null) {
+                // There is no default header row to display select all checkbox
+                return;
+            }
 
             if (selectAllCheckBoxVisible) {
                 selectionCell.setWidget(selectAllCheckBox);
diff --git a/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridNoHeaderOnInit.java b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridNoHeaderOnInit.java
new file mode 100644
index 0000000000..fa9ff25b45
--- /dev/null
+++ b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridNoHeaderOnInit.java
@@ -0,0 +1,23 @@
+package com.vaadin.tests.components.treegrid;
+
+import com.vaadin.annotations.Widgetset;
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.tests.components.AbstractTestUI;
+import com.vaadin.ui.Label;
+import com.vaadin.ui.TreeGrid;
+import com.vaadin.ui.renderers.ComponentRenderer;
+
+@Widgetset("com.vaadin.DefaultWidgetSet")
+public class TreeGridNoHeaderOnInit extends AbstractTestUI {
+
+    @Override
+    protected void setup(VaadinRequest request) {
+        TreeGrid<String> grid = new TreeGrid<>();
+        grid.addColumn(Object::toString).setCaption("toString with Caption");
+        grid.addColumn(t -> new Label(t), new ComponentRenderer());
+        grid.setItems("Foo", "Bar", "Baz");
+        grid.removeHeaderRow(0);
+        grid.appendFooterRow();
+        addComponent(grid);
+    }
+}
diff --git a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridNoHeaderOnInitTest.java b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridNoHeaderOnInitTest.java
new file mode 100644
index 0000000000..c9463e519f
--- /dev/null
+++ b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridNoHeaderOnInitTest.java
@@ -0,0 +1,15 @@
+package com.vaadin.tests.components.treegrid;
+
+import org.junit.Test;
+
+import com.vaadin.tests.tb3.SingleBrowserTest;
+
+public class TreeGridNoHeaderOnInitTest extends SingleBrowserTest {
+
+    @Test
+    public void no_exception_thrown() {
+        setDebug(true);
+        openTestURL();
+        assertNoErrorNotifications();
+    }
+}