Fix expanding an added TreeGrid node (#11018)

Fixes #9997

3 files changed, 76 insertions, 7 deletions

diff --git a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
index 3df69d45f3..f5d2d087e6 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
@@ -665,7 +665,8 @@ public class DataCommunicator<T> extends AbstractExtension {
         if (isAttached()) {
             attachDataProviderListener();
         }
-        hardReset();
+        reset = true;
+        markAsDirty();
 
         return filter -> {
             if (this.dataProvider != dataProvider) {
@@ -764,16 +765,12 @@ public class DataCommunicator<T> extends AbstractExtension {
                         generators.forEach(g -> g.refreshData(item));
                         getUI().access(() -> refresh(item));
                     } else {
-                        getUI().access(this::hardReset);
+                        reset = true;
+                        getUI().access(() -> markAsDirty());
                     }
                 });
     }
 
-    private void hardReset() {
-        reset = true;
-        markAsDirty();
-    }
-
     private void detachDataProviderListener() {
         if (dataProviderUpdateRegistration != null) {
             dataProviderUpdateRegistration.remove();
diff --git a/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridExpandOnAdd.java b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridExpandOnAdd.java
new file mode 100644
index 0000000000..a86e8cbc6c
--- /dev/null
+++ b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridExpandOnAdd.java
@@ -0,0 +1,43 @@
+package com.vaadin.tests.components.treegrid;
+
+import com.vaadin.data.TreeData;
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.tests.components.AbstractTestUI;
+import com.vaadin.ui.Button;
+import com.vaadin.ui.TreeGrid;
+
+public class TreeGridExpandOnAdd extends AbstractTestUI {
+
+    @Override
+    protected void setup(VaadinRequest request) {
+        TreeGrid<Item> tree = new TreeGrid<>(Item.class);
+        Button button = new Button("Add new item", e -> addItem(tree));
+        addComponents(tree, button);
+    }
+
+    private void addItem(TreeGrid<Item> tree) {
+        Item parent = new Item("Parent");
+        Item child = new Item("Child");
+
+        TreeData<Item> treeData = tree.getTreeData();
+        treeData.addItem(null, parent);
+        treeData.addItem(parent, child);
+
+        tree.getDataProvider().refreshAll();
+
+        tree.expand(parent);
+    }
+
+    public static class Item {
+        private String name;
+
+        public Item(String name) {
+            this.name = name;
+        }
+
+        public String getName() {
+            return name;
+        }
+    }
+
+}
diff --git a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridExpandOnAddTest.java b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridExpandOnAddTest.java
new file mode 100644
index 0000000000..00908a71e6
--- /dev/null
+++ b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridExpandOnAddTest.java
@@ -0,0 +1,29 @@
+package com.vaadin.tests.components.treegrid;
+
+import static org.junit.Assert.assertEquals;
+
+import org.junit.Test;
+
+import com.vaadin.testbench.elements.ButtonElement;
+import com.vaadin.testbench.elements.TreeGridElement;
+import com.vaadin.tests.tb3.SingleBrowserTest;
+
+public class TreeGridExpandOnAddTest extends SingleBrowserTest {
+
+    @Test
+    public void testNoException() {
+        setDebug(true);
+        openTestURL();
+
+        $(ButtonElement.class).first().click();
+
+        TreeGridElement treeGrid = $(TreeGridElement.class).first();
+        assertEquals("Parent node not added", "Parent",
+                treeGrid.getCell(0, 0).getText());
+        assertEquals("Child node not added", "Child",
+                treeGrid.getCell(1, 0).getText());
+
+        assertNoErrorNotifications();
+    }
+
+}