Make sure expanding/collapsing updates expected cache size. (#11753)

* Make sure expanding/collapsing updates expected cache size.

Fixes #11749

* Merge branch 'master' into issue11749

client/src/main/java/com/vaadin/client/ui/treegrid/TreeGridConnector.java

@@ -203,6 +203,12 @@ public class TreeGridConnector extends GridConnector {
             public void dataRemoved(int firstRowIndex, int numberOfRows) {
                 if (awaitingRowsState == AwaitingRowsState.COLLAPSE) {
                     awaitingRowsState = AwaitingRowsState.NONE;
+                    // make sure the cache stays up to date with the collapsing
+                    Range visibleRowRange = getWidget().getEscalator()
+                            .getVisibleRowRange();
+                    getDataSource().ensureAvailability(
+                            visibleRowRange.getStart(),
+                            visibleRowRange.getEnd());
                 }
                 checkExpand();
             }
@@ -211,6 +217,12 @@ public class TreeGridConnector extends GridConnector {
             public void dataAdded(int firstRowIndex, int numberOfRows) {
                 if (awaitingRowsState == AwaitingRowsState.EXPAND) {
                     awaitingRowsState = AwaitingRowsState.NONE;
+                    // make sure the cache stays up to date with the expanding
+                    Range visibleRowRange = getWidget().getEscalator()
+                            .getVisibleRowRange();
+                    getDataSource().ensureAvailability(
+                            visibleRowRange.getStart(),
+                            visibleRowRange.getEnd());
                 }
                 checkExpand();
             }

uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridCache.java

@@ -0,0 +1,112 @@
+package com.vaadin.tests.components.treegrid;
+
+import com.vaadin.data.TreeData;
+import com.vaadin.data.provider.TreeDataProvider;
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.shared.ui.grid.HeightMode;
+import com.vaadin.tests.components.AbstractTestUI;
+import com.vaadin.ui.Button;
+import com.vaadin.ui.Label;
+import com.vaadin.ui.TreeGrid;
+import com.vaadin.ui.VerticalLayout;
+
+public class TreeGridCache extends AbstractTestUI {
+
+    @Override
+    protected void setup(VaadinRequest request) {
+        Button button1 = new Button("Add TreeGrid1",
+                e -> addComponent(addTreeGrid1()));
+        button1.setId("button1");
+        addComponent(button1);
+        Button button2 = new Button("Add TreeGrid2",
+                e -> addComponent(addTreeGrid2()));
+        button2.setId("button2");
+        addComponent(button2);
+    }
+
+    private TreeGrid<Integer> addTreeGrid1() {
+        TreeData<Integer> treeData = new TreeData<>();
+        treeData.addItems(null, 10, 5);
+        treeData.addItems(10, 9, 8, 7, 6);
+        treeData.addItems(5, 4, 3, 2, 1);
+        treeData.addItems(1, 0, -1, -2, -3, -4);
+
+        TreeDataProvider<Integer> treeDataProvider = new TreeDataProvider<>(
+                treeData);
+
+        TreeGrid<Integer> tree = new TreeGrid<>();
+        tree.setCaption("TreeGrid1");
+        tree.setHeightMode(HeightMode.ROW);
+        tree.setHeightByRows(11);
+        tree.setDataProvider(treeDataProvider);
+        tree.addComponentColumn(integer -> {
+            VerticalLayout verticalLayout = new VerticalLayout();
+            verticalLayout.addComponent(new Label(integer.toString()));
+            return verticalLayout;
+        });
+        tree.addColumn(i -> String.valueOf(i));
+        treeDataProvider.refreshAll();
+        return tree;
+    }
+
+    private TreeGrid<TestBean> addTreeGrid2() {
+        TestBean parent = null;
+        TreeGrid<TestBean> tree = new TreeGrid<>(TestBean.class);
+        tree.setCaption("TreeGrid2");
+        tree.setColumns("key", "value");
+        tree.setHierarchyColumn("key");
+        TreeData<TestBean> treeData = new TreeData<>();
+
+        for (int i = 0; i < 2; i++) {
+            TestBean val = new TestBean("root" + i, " root value " + i);
+            treeData.addItem(parent, val);
+            for (int j = 0; j < 10; j++) {
+                TestBean v = new TestBean("leaf" + j, "leaf value" + j);
+                treeData.addItem(val, v);
+            }
+        }
+        TreeDataProvider<TestBean> dataProvider = new TreeDataProvider<>(
+                treeData);
+        tree.setDataProvider(dataProvider);
+        tree.setHeight("500px");
+
+        return tree;
+    }
+
+    @Override
+    protected String getTestDescription() {
+        return "Cache should cover more than 4 x the initial row amount after rows have been expanded. "
+                + "Clicking on a row beyond that should not cause an exception.";
+    }
+
+    @Override
+    protected Integer getTicketNumber() {
+        return 11749;
+    }
+
+    public class TestBean {
+        private String key;
+        private String value;
+
+        public TestBean(String key, String value) {
+            this.key = key;
+            this.value = value;
+        }
+
+        public String getKey() {
+            return key;
+        }
+
+        public void setKey(String key) {
+            this.key = key;
+        }
+
+        public String getValue() {
+            return value;
+        }
+
+        public void setValue(String value) {
+            this.value = value;
+        }
+    }
+}

uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridCacheTest.java

@@ -0,0 +1,41 @@
+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.GridElement.GridCellElement;
+import com.vaadin.testbench.elements.NotificationElement;
+import com.vaadin.testbench.elements.TreeGridElement;
+import com.vaadin.tests.tb3.MultiBrowserTest;
+
+public class TreeGridCacheTest extends MultiBrowserTest {
+
+    @Test
+    public void testGrid1() {
+        openTestURL();
+        $(ButtonElement.class).id("button1").click();
+        TreeGridElement grid = $(TreeGridElement.class).first();
+        grid.expandWithClick(1);
+        grid.expandWithClick(5);
+        GridCellElement cell = grid.getCell(10, 1);
+        assertEquals("-4", cell.getText());
+
+        cell.click();
+        assertEquals(0, $(NotificationElement.class).all().size());
+    }
+
+    @Test
+    public void testGrid2() {
+        openTestURL();
+        $(ButtonElement.class).id("button2").click();
+        TreeGridElement grid = $(TreeGridElement.class).first();
+        grid.expandWithClick(1);
+        GridCellElement cell = grid.getCell(10, 1);
+        assertEquals("leaf value8", cell.getText());
+
+        cell.click();
+        assertEquals(0, $(NotificationElement.class).all().size());
+    }
+}