Fix initial expand of TreeGrid items

Fixes #9050

5 files changed, 72 insertions, 10 deletions

diff --git a/client/src/main/java/com/vaadin/client/connectors/treegrid/TreeGridConnector.java b/client/src/main/java/com/vaadin/client/connectors/treegrid/TreeGridConnector.java
index 10935264c7..72c4487aaf 100644
--- a/client/src/main/java/com/vaadin/client/connectors/treegrid/TreeGridConnector.java
+++ b/client/src/main/java/com/vaadin/client/connectors/treegrid/TreeGridConnector.java
@@ -19,7 +19,6 @@ import java.util.Collection;
 import java.util.HashSet;
 import java.util.Set;
 import java.util.logging.Logger;
-import java.util.stream.IntStream;
 
 import com.google.gwt.core.client.Scheduler;
 import com.google.gwt.dom.client.BrowserEvents;
@@ -171,6 +170,11 @@ public class TreeGridConnector extends GridConnector {
             public void setCollapsed(String key) {
                 rowKeysPendingExpand.remove(key);
             }
+
+            @Override
+            public void clearPendingExpands() {
+                rowKeysPendingExpand.clear();
+            }
         });
     }
 
@@ -201,7 +205,7 @@ public class TreeGridConnector extends GridConnector {
 
             @Override
             public void resetDataAndSize(int estimatedNewDataSize) {
-                rowKeysPendingExpand.clear();
+                // NO-OP
             }
         });
     }
@@ -348,14 +352,15 @@ public class TreeGridConnector extends GridConnector {
         if (rowKeysPendingExpand.isEmpty()) {
             return;
         }
-        IntStream.range(firstRowIndex, firstRowIndex + numberOfRows)
-                .forEach(rowIndex -> {
-                    String rowKey = getDataSource().getRow(rowIndex)
-                            .getString(DataCommunicatorConstants.KEY);
-                    if (rowKeysPendingExpand.remove(rowKey)) {
-                        setCollapsedServerInitiated(rowIndex, false);
-                    }
-                });
+        for (int rowIndex = firstRowIndex; rowIndex < firstRowIndex
+                + numberOfRows; rowIndex++) {
+            String rowKey = getDataSource().getRow(rowIndex)
+                    .getString(DataCommunicatorConstants.KEY);
+            if (rowKeysPendingExpand.remove(rowKey)) {
+                setCollapsedServerInitiated(rowIndex, false);
+                return;
+            }
+        }
     }
 
     private static boolean isCollapsed(JsonObject rowData) {
diff --git a/server/src/main/java/com/vaadin/ui/TreeGrid.java b/server/src/main/java/com/vaadin/ui/TreeGrid.java
index b5da3f8645..6d5b9bb191 100644
--- a/server/src/main/java/com/vaadin/ui/TreeGrid.java
+++ b/server/src/main/java/com/vaadin/ui/TreeGrid.java
@@ -391,6 +391,7 @@ public class TreeGrid<T> extends Grid<T> {
             throw new IllegalArgumentException(
                     "TreeGrid only accepts hierarchical data providers");
         }
+        getRpcProxy(TreeGridClientRpc.class).clearPendingExpands();
         super.setDataProvider(dataProvider);
     }
 
diff --git a/shared/src/main/java/com/vaadin/shared/ui/treegrid/TreeGridClientRpc.java b/shared/src/main/java/com/vaadin/shared/ui/treegrid/TreeGridClientRpc.java
index a07a2d4885..16d2179741 100644
--- a/shared/src/main/java/com/vaadin/shared/ui/treegrid/TreeGridClientRpc.java
+++ b/shared/src/main/java/com/vaadin/shared/ui/treegrid/TreeGridClientRpc.java
@@ -42,4 +42,9 @@ public interface TreeGridClientRpc extends ClientRpc {
      *            the communication key of the collapsed item
      */
     public void setCollapsed(String key);
+
+    /**
+     * Clear all pending expands from the client.
+     */
+    public void clearPendingExpands();
 }
diff --git a/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridInitialExpand.java b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridInitialExpand.java
new file mode 100644
index 0000000000..db5b8c54b6
--- /dev/null
+++ b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridInitialExpand.java
@@ -0,0 +1,28 @@
+package com.vaadin.tests.components.treegrid;
+
+import com.vaadin.data.HierarchyData;
+import com.vaadin.data.provider.InMemoryHierarchicalDataProvider;
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.tests.components.AbstractTestUI;
+import com.vaadin.ui.TreeGrid;
+
+public class TreeGridInitialExpand extends AbstractTestUI {
+
+    @Override
+    protected void setup(VaadinRequest request) {
+        TreeGrid<String> treeGrid = new TreeGrid<>();
+        treeGrid.setCaption("Test");
+        treeGrid.addColumn(String::toString).setCaption("String");
+        HierarchyData<String> data = new HierarchyData<>();
+        data.addItem(null, "parent1");
+        data.addItem("parent1", "parent1-child1");
+        data.addItem("parent1", "parent1-child2");
+        data.addItem(null, "parent2");
+        data.addItem("parent2", "parent2-child2");
+        treeGrid.setDataProvider(new InMemoryHierarchicalDataProvider<>(data));
+        treeGrid.setHeightByRows(5);
+        treeGrid.expand("parent1");
+        treeGrid.expand("parent2");
+        addComponent(treeGrid);
+    }
+}
diff --git a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridInitialExpandTest.java b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridInitialExpandTest.java
new file mode 100644
index 0000000000..d82155da6b
--- /dev/null
+++ b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridInitialExpandTest.java
@@ -0,0 +1,23 @@
+package com.vaadin.tests.components.treegrid;
+
+import static org.junit.Assert.assertEquals;
+
+import org.junit.Test;
+
+import com.vaadin.testbench.elements.TreeGridElement;
+import com.vaadin.tests.tb3.SingleBrowserTest;
+
+public class TreeGridInitialExpandTest extends SingleBrowserTest {
+
+    @Test
+    public void initial_expand_of_items() {
+        openTestURL();
+        TreeGridElement grid = $(TreeGridElement.class).first();
+
+        assertEquals("parent1", grid.getCell(0, 0).getText());
+        assertEquals("parent1-child1", grid.getCell(1, 0).getText());
+        assertEquals("parent1-child2", grid.getCell(2, 0).getText());
+        assertEquals("parent2", grid.getCell(3, 0).getText());
+        assertEquals("parent2-child2", grid.getCell(4, 0).getText());
+    }
+}