Hierarchical data (#8842)

* Initial HierarchicalDataProvider for TreeGrid
* Initial in-memory hierarchical data implementation
* TreeGrid declarative support

Fixes #8611, Fixes #8620 

6 files changed, 403 insertions, 217 deletions

diff --git a/uitest/src/main/java/com/vaadin/tests/components/treegrid/LazyHierarchicalDataProvider.java b/uitest/src/main/java/com/vaadin/tests/components/treegrid/LazyHierarchicalDataProvider.java
new file mode 100644
index 0000000000..33c8f15a47
--- /dev/null
+++ b/uitest/src/main/java/com/vaadin/tests/components/treegrid/LazyHierarchicalDataProvider.java
@@ -0,0 +1,63 @@
+package com.vaadin.tests.components.treegrid;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Stream;
+
+import com.vaadin.data.provider.AbstractHierarchicalDataProvider;
+import com.vaadin.data.provider.HierarchicalQuery;
+import com.vaadin.tests.components.treegrid.TreeGridBasicFeatures.HierarchicalTestBean;
+
+public class LazyHierarchicalDataProvider
+        extends AbstractHierarchicalDataProvider<HierarchicalTestBean, Void> {
+
+    private final int nodesPerLevel;
+    private final int depth;
+
+    public LazyHierarchicalDataProvider(int nodesPerLevel, int depth) {
+        this.nodesPerLevel = nodesPerLevel;
+        this.depth = depth;
+    }
+
+    @Override
+    public int getChildCount(
+            HierarchicalQuery<HierarchicalTestBean, Void> query) {
+
+        Optional<Integer> count = query.getParentOptional()
+                .flatMap(parent -> Optional.of(Integer.valueOf(
+                        (internalHasChildren(parent) ? nodesPerLevel : 0))));
+
+        return count.orElse(nodesPerLevel);
+    }
+
+    @Override
+    public Stream<HierarchicalTestBean> fetchChildren(
+            HierarchicalQuery<HierarchicalTestBean, Void> query) {
+        final int depth = query.getParentOptional().isPresent()
+                ? query.getParent().getDepth() + 1 : 0;
+        final Optional<String> parentKey = query.getParentOptional()
+                .flatMap(parent -> Optional.of(parent.getId()));
+
+        List<HierarchicalTestBean> list = new ArrayList<>();
+        for (int i = 0; i < query.getLimit(); i++) {
+            list.add(new HierarchicalTestBean(parentKey.orElse(null), depth,
+                    i + query.getOffset()));
+        }
+        return list.stream();
+    }
+
+    @Override
+    public boolean hasChildren(HierarchicalTestBean item) {
+        return internalHasChildren(item);
+    }
+
+    private boolean internalHasChildren(HierarchicalTestBean node) {
+        return node.getDepth() < depth;
+    }
+
+    @Override
+    public boolean isInMemory() {
+        return false;
+    }
+}
diff --git a/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeatures.java b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeatures.java
index 77e1e2eefa..7998537a6e 100644
--- a/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeatures.java
+++ b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeatures.java
@@ -1,28 +1,24 @@
 package com.vaadin.tests.components.treegrid;
 
-import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.LinkedHashMap;
-import java.util.LinkedHashSet;
 import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
 
-import com.vaadin.data.provider.DataProviderListener;
-import com.vaadin.data.provider.HierarchicalDataProvider;
-import com.vaadin.data.provider.Query;
-import com.vaadin.shared.Registration;
+import com.vaadin.annotations.Theme;
+import com.vaadin.annotations.Widgetset;
+import com.vaadin.data.HierarchyData;
+import com.vaadin.data.provider.DataProvider;
+import com.vaadin.data.provider.InMemoryHierarchicalDataProvider;
 import com.vaadin.tests.components.AbstractComponentTest;
-import com.vaadin.ui.MenuBar.MenuItem;
 import com.vaadin.ui.TreeGrid;
 
+@Theme("valo")
+@Widgetset("com.vaadin.DefaultWidgetSet")
 public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
 
-    private TreeGrid<TestBean> grid;
-    private TestDataProvider dataProvider = new TestDataProvider();
+    private TreeGrid<HierarchicalTestBean> grid;
+    private InMemoryHierarchicalDataProvider<HierarchicalTestBean> inMemoryDataProvider;
+    private LazyHierarchicalDataProvider lazyDataProvider;
 
     @Override
     public TreeGrid getComponent() {
@@ -36,12 +32,17 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
 
     @Override
     protected void initializeComponents() {
+        initializeDataProviders();
         grid = new TreeGrid<>();
         grid.setSizeFull();
-        grid.addColumn(TestBean::getStringValue).setId("First column");
-        grid.addColumn(TestBean::getStringValue).setId("Second column");
-        grid.setHierarchyColumn("First column");
-        grid.setDataProvider(dataProvider);
+        grid.addColumn(HierarchicalTestBean::toString).setCaption("String")
+                .setId("string");
+        grid.addColumn(HierarchicalTestBean::getDepth).setCaption("Depth")
+                .setId("depth");
+        grid.addColumn(HierarchicalTestBean::getIndex)
+                .setCaption("Index on this depth").setId("index");
+        grid.setHierarchyColumn("string");
+        grid.setDataProvider(new LazyHierarchicalDataProvider(3, 2));
 
         grid.setId("testComponent");
         addTestComponent(grid);
@@ -51,8 +52,44 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
     protected void createActions() {
         super.createActions();
 
+        createDataProviderSelect();
         createHierarchyColumnSelect();
-        createToggleCollapseSelect();
+    }
+
+    private void initializeDataProviders() {
+        HierarchyData<HierarchicalTestBean> data = new HierarchyData<>();
+
+        List<Integer> ints = Arrays.asList(0, 1, 2);
+
+        ints.stream().forEach(index -> {
+            HierarchicalTestBean bean = new HierarchicalTestBean(null, 0,
+                    index);
+            data.addItem(null, bean);
+            ints.stream().forEach(childIndex -> {
+                HierarchicalTestBean childBean = new HierarchicalTestBean(
+                        bean.getId(), 1, childIndex);
+                data.addItem(bean, childBean);
+                ints.stream()
+                        .forEach(grandChildIndex -> data.addItem(childBean,
+                                new HierarchicalTestBean(childBean.getId(), 2,
+                                        grandChildIndex)));
+            });
+        });
+
+        inMemoryDataProvider = new InMemoryHierarchicalDataProvider<>(data);
+        lazyDataProvider = new LazyHierarchicalDataProvider(3, 2);
+    }
+
+    @SuppressWarnings("unchecked")
+    private void createDataProviderSelect() {
+        @SuppressWarnings("rawtypes")
+        LinkedHashMap<String, DataProvider> options = new LinkedHashMap<>();
+        options.put("LazyHierarchicalDataProvider", lazyDataProvider);
+        options.put("InMemoryHierarchicalDataProvider", inMemoryDataProvider);
+
+        createSelectAction("Set data provider", CATEGORY_FEATURES, options,
+                "LazyHierarchicalDataProvider",
+                (treeGrid, value, data) -> treeGrid.setDataProvider(value));
     }
 
     private void createHierarchyColumnSelect() {
@@ -65,205 +102,64 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
                 (treeGrid, value, data) -> treeGrid.setHierarchyColumn(value));
     }
 
-    private void createToggleCollapseSelect() {
-        MenuItem menu = createCategory("Toggle expand", CATEGORY_FEATURES);
-        dataProvider.getAllItems().forEach(testBean -> {
-            createClickAction(testBean.getStringValue(), "Toggle expand",
-                    (grid, bean, data) -> grid.toggleCollapse(bean), testBean);
-        });
-    }
-
-    private static class TestBean {
-
-        private String stringValue;
-
-        public TestBean(String stringValue) {
-            this.stringValue = stringValue;
-        }
-
-        public String getStringValue() {
-            return stringValue;
-        }
-
-        public void setStringValue(String stringValue) {
-            this.stringValue = stringValue;
-        }
-    }
-
-    private static class TestDataProvider
-            implements HierarchicalDataProvider<TestBean, Void> {
-
-        private static class HierarchyWrapper<T> {
-            private T item;
-            private T parent;
-            private Set<T> children;
-            private boolean collapsed;
-
-            public HierarchyWrapper(T item, T parent, boolean collapsed) {
-                this.item = item;
-                this.parent = parent;
-                this.collapsed = collapsed;
-                children = new LinkedHashSet<>();
-            }
-
-            public T getItem() {
-                return item;
-            }
-
-            public void setItem(T item) {
-                this.item = item;
-            }
-
-            public T getParent() {
-                return parent;
-            }
-
-            public void setParent(T parent) {
-                this.parent = parent;
-            }
-
-            public Set<T> getChildren() {
-                return children;
-            }
-
-            public void setChildren(Set<T> children) {
-                this.children = children;
-            }
-
-            public boolean isCollapsed() {
-                return collapsed;
-            }
-
-            public void setCollapsed(boolean collapsed) {
-                this.collapsed = collapsed;
-            }
-        }
-
-        private Map<TestBean, HierarchyWrapper<TestBean>> itemToWrapperMap;
-        private Map<HierarchyWrapper<TestBean>, TestBean> wrapperToItemMap;
-        private Map<TestBean, HierarchyWrapper<TestBean>> rootNodes;
-
-        public TestDataProvider() {
-            itemToWrapperMap = new LinkedHashMap<>();
-            wrapperToItemMap = new LinkedHashMap<>();
-            rootNodes = new LinkedHashMap<>();
+    static class HierarchicalTestBean {
 
-            List<String> strings = Arrays.asList("a", "b", "c");
+        private final String id;
+        private final int depth;
+        private final int index;
 
-            strings.stream().forEach(string -> {
-                TestBean rootBean = new TestBean(string);
-
-                HierarchyWrapper<TestBean> wrappedParent = new HierarchyWrapper<>(
-                        rootBean, null, true);
-                itemToWrapperMap.put(rootBean, wrappedParent);
-                wrapperToItemMap.put(wrappedParent, rootBean);
-
-                List<TestBean> children = strings.stream().map(string2 -> {
-                    TestBean childBean = new TestBean(string + "/" + string2);
-                    HierarchyWrapper<TestBean> wrappedChild = new HierarchyWrapper<>(
-                            new TestBean(string + "/" + string2), rootBean,
-                            true);
-                    itemToWrapperMap.put(childBean, wrappedChild);
-                    wrapperToItemMap.put(wrappedChild, childBean);
-                    return childBean;
-                }).collect(Collectors.toList());
-
-                wrappedParent.setChildren(new LinkedHashSet<>(children));
-
-                rootNodes.put(rootBean, wrappedParent);
-            });
+        public HierarchicalTestBean(String parentId, int depth, int index) {
+            id = (parentId == null ? "" : parentId) + "/" + depth + "/" + index;
+            this.depth = depth;
+            this.index = index;
         }
 
-        @Override
-        public int getDepth(TestBean item) {
-            int depth = 0;
-            while (getItem(item) != null) {
-                item = getItem(item).getParent();
-                depth++;
-            }
+        public int getDepth() {
             return depth;
         }
 
-        @Override
-        public boolean isInMemory() {
-            return true;
-        }
-
-        @Override
-        public void refreshItem(TestBean item) {
-            // NO-OP
-        }
-
-        @Override
-        public void refreshAll() {
-            // NO-OP
-        }
-
-        @Override
-        public Registration addDataProviderListener(
-                DataProviderListener<TestBean> listener) {
-            return () -> {
-            };
-        }
-
-        private List<TestBean> getAllItems() {
-            return new ArrayList<>(itemToWrapperMap.keySet());
-        }
-
-        private List<TestBean> getVisibleItemsRecursive(
-                Collection<HierarchyWrapper<TestBean>> wrappedItems) {
-            List<TestBean> items = new ArrayList<>();
-
-            wrappedItems.forEach(wrappedItem -> {
-                items.add(wrapperToItemMap.get(wrappedItem));
-                if (!wrappedItem.isCollapsed()) {
-                    List<HierarchyWrapper<TestBean>> wrappedChildren = wrappedItem
-                            .getChildren().stream()
-                            .map(childItem -> getItem(childItem))
-                            .collect(Collectors.toList());
-                    items.addAll(getVisibleItemsRecursive(wrappedChildren));
-                }
-            });
-            return items;
+        public int getIndex() {
+            return index;
         }
 
-        @Override
-        public int size(Query<TestBean, Void> query) {
-            return getVisibleItemsRecursive(rootNodes.values()).size();
-        }
-
-        @Override
-        public Stream<TestBean> fetch(Query<TestBean, Void> query) {
-            return getVisibleItemsRecursive(rootNodes.values()).stream();
-        }
-
-        @Override
-        public boolean isRoot(TestBean item) {
-            return getItem(item).getParent() == null;
+        public String getId() {
+            return id;
         }
 
         @Override
-        public TestBean getParent(TestBean item) {
-            return getItem(item).getParent();
+        public String toString() {
+            return depth + " | " + index;
         }
 
         @Override
-        public boolean isCollapsed(TestBean item) {
-            return getItem(item).isCollapsed();
+        public int hashCode() {
+            final int prime = 31;
+            int result = 1;
+            result = prime * result + ((id == null) ? 0 : id.hashCode());
+            return result;
         }
 
         @Override
-        public boolean hasChildren(TestBean item) {
-            return !getItem(item).getChildren().isEmpty();
-        }
-
-        @Override
-        public void setCollapsed(TestBean item, boolean b) {
-            getItem(item).setCollapsed(b);
+        public boolean equals(Object obj) {
+            if (this == obj) {
+                return true;
+            }
+            if (obj == null) {
+                return false;
+            }
+            if (getClass() != obj.getClass()) {
+                return false;
+            }
+            HierarchicalTestBean other = (HierarchicalTestBean) obj;
+            if (id == null) {
+                if (other.id != null) {
+                    return false;
+                }
+            } else if (!id.equals(other.id)) {
+                return false;
+            }
+            return true;
         }
 
-        private HierarchyWrapper<TestBean> getItem(TestBean item) {
-            return itemToWrapperMap.get(item);
-        }
     }
 }
diff --git a/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridScrolling.java b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridScrolling.java
new file mode 100644
index 0000000000..8d859918b8
--- /dev/null
+++ b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridScrolling.java
@@ -0,0 +1,42 @@
+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.tests.components.treegrid.TreeGridBasicFeatures.HierarchicalTestBean;
+import com.vaadin.ui.TreeGrid;
+
+@Widgetset("com.vaadin.DefaultWidgetSet")
+public class TreeGridScrolling extends AbstractTestUI {
+
+    public static final int DEFAULT_NODES = 20;
+    public static final int DEFAULT_DEPTH = 3;
+    public static final String NODES_PARAMETER = "nodes";
+    public static final String DEPTH_PARAMETER = "depth";
+
+    @Override
+    protected void setup(VaadinRequest request) {
+        int depth = DEFAULT_DEPTH;
+        if (request.getParameter(DEPTH_PARAMETER) != null) {
+            depth = Integer.parseInt(request.getParameter(DEPTH_PARAMETER));
+        }
+        int nodes = DEFAULT_NODES;
+        if (request.getParameter(NODES_PARAMETER) != null) {
+            nodes = Integer.parseInt(request.getParameter(NODES_PARAMETER));
+        }
+
+        TreeGrid<HierarchicalTestBean> grid = new TreeGrid<>();
+        grid.setSizeFull();
+        grid.addColumn(HierarchicalTestBean::toString).setCaption("String")
+                .setId("string");
+        grid.addColumn(HierarchicalTestBean::getDepth).setCaption("Depth")
+                .setId(DEPTH_PARAMETER);
+        grid.addColumn(HierarchicalTestBean::getIndex)
+                .setCaption("Index on this depth").setId("index");
+        grid.setHierarchyColumn("string");
+        grid.setDataProvider(new LazyHierarchicalDataProvider(nodes, depth));
+
+        addComponent(grid);
+    }
+
+}
diff --git a/uitest/src/main/java/com/vaadin/tests/components/treetable/TreeTableCacheOnPartialUpdates.java b/uitest/src/main/java/com/vaadin/tests/components/treetable/TreeTableCacheOnPartialUpdates.java
index 3824c6c5d0..e008dde0f5 100644
--- a/uitest/src/main/java/com/vaadin/tests/components/treetable/TreeTableCacheOnPartialUpdates.java
+++ b/uitest/src/main/java/com/vaadin/tests/components/treetable/TreeTableCacheOnPartialUpdates.java
@@ -62,7 +62,7 @@ public class TreeTableCacheOnPartialUpdates extends TestBase {
 
         @Override
         public String toString() {
-            return "TestBean [col1=" + col1 + ", col2=" + col2 + "]";
+            return "HierarchicalTestBean [col1=" + col1 + ", col2=" + col2 + "]";
         }
 
     }
diff --git a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeaturesTest.java b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeaturesTest.java
index c84c96232c..b01a6273e2 100644
--- a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeaturesTest.java
+++ b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeaturesTest.java
@@ -1,19 +1,38 @@
 package com.vaadin.tests.components.treegrid;
 
+import java.util.Arrays;
+import java.util.Collection;
+
 import org.junit.Assert;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized.Parameters;
 import org.openqa.selenium.Keys;
 import org.openqa.selenium.interactions.Actions;
 
 import com.vaadin.testbench.By;
 import com.vaadin.testbench.elements.TreeGridElement;
 import com.vaadin.tests.tb3.MultiBrowserTest;
+import com.vaadin.tests.tb3.ParameterizedTB3Runner;
 
+@RunWith(ParameterizedTB3Runner.class)
 public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
 
     private TreeGridElement grid;
 
+    public void setDataProvider(String dataProviderString) {
+        selectMenuPath("Component", "Features", "Set data provider",
+                dataProviderString);
+    }
+
+    @Parameters
+    public static Collection<String> getDataProviders() {
+        return Arrays.asList("LazyHierarchicalDataProvider",
+                "InMemoryHierarchicalDataProvider");
+    }
+
     @Before
     public void before() {
         openTestURL("theme=valo");
@@ -21,56 +40,59 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
     }
 
     @Test
+    @Ignore // currently no implementation exists for toggling from the server
+            // side
     public void toggle_collapse_server_side() {
         Assert.assertEquals(3, grid.getRowCount());
-        assertCellTexts(0, 0, new String[] { "a", "b", "c" });
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
 
-        selectMenuPath("Component", "Features", "Toggle expand", "a");
+        selectMenuPath("Component", "Features", "Toggle expand", "0 | 0");
         Assert.assertEquals(6, grid.getRowCount());
-        assertCellTexts(1, 0, new String[] { "a/a", "a/b", "a/c" });
+        assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "1 | 2" });
 
-        selectMenuPath("Component", "Features", "Toggle expand", "a");
+        selectMenuPath("Component", "Features", "Toggle expand", "0 | 0");
         Assert.assertEquals(3, grid.getRowCount());
-        assertCellTexts(0, 0, new String[] { "a", "b", "c" });
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
 
         // collapsing a leaf should have no effect
-        selectMenuPath("Component", "Features", "Toggle expand", "a/a");
+        selectMenuPath("Component", "Features", "Toggle expand", "1 | 0");
         Assert.assertEquals(3, grid.getRowCount());
     }
 
     @Test
     public void non_leaf_collapse_on_click() {
         Assert.assertEquals(3, grid.getRowCount());
-        assertCellTexts(0, 0, new String[] { "a", "b", "c" });
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
 
-        // click the expander corresponding to "a"
+        // Should expand "0 | 0"
         grid.getRow(0).getCell(0)
                 .findElement(By.className("v-tree-grid-expander")).click();
         Assert.assertEquals(6, grid.getRowCount());
-        assertCellTexts(1, 0, new String[] { "a/a", "a/b", "a/c" });
+        assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "1 | 2" });
 
-        // click the expander corresponding to "a"
+        // Should collapse "0 | 0"
         grid.getRow(0).getCell(0)
                 .findElement(By.className("v-tree-grid-expander")).click();
         Assert.assertEquals(3, grid.getRowCount());
-        assertCellTexts(0, 0, new String[] { "a", "b", "c" });
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
     }
 
     @Test
+    @Ignore // FIXME: remove ignore annotation once #8758 is done
     public void keyboard_navigation() {
         grid.getRow(0).getCell(0).click();
 
-        // Should expand "a"
+        // Should expand "0 | 0"
         new Actions(getDriver()).keyDown(Keys.ALT).sendKeys(Keys.RIGHT)
                 .keyUp(Keys.ALT).perform();
         Assert.assertEquals(6, grid.getRowCount());
-        assertCellTexts(1, 0, new String[] { "a/a", "a/b", "a/c" });
+        assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "1 | 2" });
 
-        // Should collapse "a"
+        // Should collapse "0 | 0"
         new Actions(getDriver()).keyDown(Keys.ALT).sendKeys(Keys.LEFT)
                 .keyUp(Keys.ALT).perform();
         Assert.assertEquals(3, grid.getRowCount());
-        assertCellTexts(0, 0, new String[] { "a", "b", "c" });
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
     }
 
     @Test
@@ -81,7 +103,7 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
                 .isElementPresent(By.className("v-tree-grid-expander")));
 
         selectMenuPath("Component", "Features", "Set hierarchy column",
-                "Second column");
+                "depth");
 
         Assert.assertFalse(grid.getRow(0).getCell(0)
                 .isElementPresent(By.className("v-tree-grid-expander")));
@@ -89,7 +111,7 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
                 .isElementPresent(By.className("v-tree-grid-expander")));
 
         selectMenuPath("Component", "Features", "Set hierarchy column",
-                "First column");
+                "string");
 
         Assert.assertTrue(grid.getRow(0).getCell(0)
                 .isElementPresent(By.className("v-tree-grid-expander")));
diff --git a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridScrollingTest.java b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridScrollingTest.java
new file mode 100644
index 0000000000..be7953ba35
--- /dev/null
+++ b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridScrollingTest.java
@@ -0,0 +1,163 @@
+package com.vaadin.tests.components.treegrid;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import com.vaadin.testbench.elements.TreeGridElement;
+import com.vaadin.tests.tb3.SingleBrowserTest;
+
+public class TreeGridScrollingTest extends SingleBrowserTest {
+
+    @Test
+    public void testScrollingTree_expandCollapseFromBeginning_correctItemsShown() {
+        // TODO refactor this test to verify each row against a model, e.g. a
+        // InMemoryDataProvider, or the used lazy hierarchical data provider
+        openTestURL();
+
+        TreeGridElement grid = $(TreeGridElement.class).first();
+
+        Assert.assertEquals(grid.getRowCount(),
+                TreeGridScrolling.DEFAULT_NODES);
+
+        verifyRow(0, 0, 0);
+        verifyRow(10, 0, 10);
+        verifyRow(19, 0, 19);
+        verifyRow(10, 0, 10);
+        verifyRow(0, 0, 0);
+
+        grid.expandWithClick(0);
+
+        verifyRow(0, 0, 0);
+        verifyRow(1, 1, 0);
+        verifyRow(11, 1, 10);
+        verifyRow(20, 1, 19);
+        verifyRow(39, 0, 19);
+
+        // verifying in reverse order causes scrolling up
+        verifyRow(20, 1, 19);
+        verifyRow(11, 1, 10);
+        verifyRow(1, 1, 0);
+        verifyRow(0, 0, 0);
+
+        grid.expandWithClick(3);
+
+        verifyRow(0, 0, 0);
+
+        verifyRow(1, 1, 0);
+        verifyRow(2, 1, 1);
+        verifyRow(3, 1, 2);
+
+        verifyRow(4, 2, 0);
+
+        verifyRow(14, 2, 10);
+        verifyRow(23, 2, 19);
+        verifyRow(24, 1, 3);
+        verifyRow(40, 1, 19);
+        verifyRow(59, 0, 19);
+
+        // scroll back up
+
+        verifyRow(40, 1, 19);
+        verifyRow(24, 1, 3);
+        verifyRow(23, 2, 19);
+        verifyRow(14, 2, 10);
+
+        verifyRow(4, 2, 0);
+        verifyRow(2, 1, 1);
+        verifyRow(3, 1, 2);
+        verifyRow(1, 1, 0);
+        verifyRow(0, 0, 0);
+
+        grid.expandWithClick(2);
+
+        verifyRow(0, 0, 0);
+
+        verifyRow(1, 1, 0);
+        verifyRow(2, 1, 1);
+        verifyRow(3, 2, 0);
+        verifyRow(22, 2, 19);
+
+        verifyRow(23, 1, 2);
+        verifyRow(24, 2, 0);
+
+        verifyRow(43, 2, 19);
+        verifyRow(44, 1, 3);
+        verifyRow(60, 1, 19);
+        verifyRow(79, 0, 19);
+
+        // scroll back up
+        verifyRow(60, 1, 19);
+        verifyRow(44, 1, 3);
+        verifyRow(43, 2, 19);
+
+        verifyRow(24, 2, 0);
+        verifyRow(23, 1, 2);
+
+        verifyRow(22, 2, 19);
+        verifyRow(3, 2, 0);
+        verifyRow(2, 1, 1);
+        verifyRow(1, 1, 0);
+
+        verifyRow(0, 0, 0);
+
+        grid.collapseWithClick(2);
+
+        verifyRow(0, 0, 0);
+
+        verifyRow(1, 1, 0);
+        verifyRow(2, 1, 1);
+        verifyRow(3, 1, 2);
+
+        verifyRow(4, 2, 0);
+
+        verifyRow(14, 2, 10);
+        verifyRow(23, 2, 19);
+        verifyRow(24, 1, 3);
+        verifyRow(40, 1, 19);
+        verifyRow(59, 0, 19);
+
+        // scroll back up
+
+        verifyRow(40, 1, 19);
+        verifyRow(24, 1, 3);
+        verifyRow(23, 2, 19);
+        verifyRow(14, 2, 10);
+
+        verifyRow(4, 2, 0);
+        verifyRow(2, 1, 1);
+        verifyRow(3, 1, 2);
+        verifyRow(1, 1, 0);
+        verifyRow(0, 0, 0);
+
+        grid.expandWithClick(3);
+
+        verifyRow(0, 0, 0);
+        verifyRow(1, 1, 0);
+        verifyRow(11, 1, 10);
+        verifyRow(20, 1, 19);
+        verifyRow(39, 0, 19);
+
+        // scroll back up
+
+        verifyRow(20, 1, 19);
+        verifyRow(11, 1, 10);
+        verifyRow(1, 1, 0);
+        verifyRow(0, 0, 0);
+
+        grid.expandWithClick(0);
+
+        verifyRow(0, 0, 0);
+        verifyRow(10, 0, 10);
+        verifyRow(19, 0, 19);
+        verifyRow(10, 0, 10);
+        verifyRow(0, 0, 0);
+    }
+
+    private void verifyRow(int rowActualIndex, int depth, int levelIndex) {
+        TreeGridElement grid = $(TreeGridElement.class).first();
+
+        Assert.assertEquals("Invalid row at index " + rowActualIndex,
+                depth + " | " + levelIndex,
+                grid.getCell(rowActualIndex, 0).getText());
+    }
+}