Add recursive expand and collapse method to TreeGrid and Tree (#10283)

12 files changed, 599 insertions, 51 deletions

diff --git a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
index ffdcf042c9..e09681978d 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
@@ -156,33 +156,83 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
     }
 
     /**
-     * Collapses given item, removing all its subtrees. Calling this method will
-     * have no effect if the row is already collapsed.
+     * Collapses the given item and removes its sub-hierarchy. Calling this
+     * method will have no effect if the row is already collapsed.
      *
      * @param item
      *            the item to collapse
      */
     public void collapse(T item) {
-        if (mapper.isExpanded(item)) {
-            doCollapse(item, mapper.getIndexOf(item));
-        }
+        collapse(item, true);
+    }
+
+    /**
+     * Collapses the given item and removes its sub-hierarchy. Calling this
+     * method will have no effect if the row is already collapsed.
+     * {@code syncAndRefresh} indicates whether the changes should be
+     * synchronised to the client and the data provider be notified.
+     * 
+     * @param item
+     *            the item to collapse
+     * @param syncAndRefresh
+     *            {@code true} if the changes should be synchronised to the
+     *            client and the data provider should be notified of the
+     *            changes, {@code false} otherwise.
+     */
+    public void collapse(T item, boolean syncAndRefresh) {
+        Integer index = syncAndRefresh ? mapper.getIndexOf(item).orElse(null) : null;
+        doCollapse(item, index, syncAndRefresh);
     }
 
     /**
-     * Collapses given item, removing all its subtrees. Calling this method will
-     * have no effect if the row is already collapsed. The index is provided by
-     * the client-side or calculated from a full data request.
+     * Collapses the given item and removes its sub-hierarchy. Calling this
+     * method will have no effect if the row is already collapsed.
      *
-     * @see #collapse(Object)
+     * @param item
+     *            the item to collapse
+     * @param index
+     *            the index of the item
+     */
+    public void collapse(T item, Integer index) {
+        doCollapse(item, index, true);
+    }
+
+    /**
+     * Collapses given item and removes its sub-hierarchy. Calling this method
+     * will have no effect if the row is already collapsed. The index is
+     * provided by the client-side or calculated from a full data request.
      *
+     * 
      * @param item
      *            the item to collapse
      * @param index
      *            the index of the item
+     * @deprecated Use {@link #collapse(Object, Integer)} instead.
      */
+    @Deprecated
     public void doCollapse(T item, Optional<Integer> index) {
-        if (mapper.isExpanded(item)) {
-            Range removedRows = mapper.doCollapse(item, index);
+        doCollapse(item, index.orElse(null), true);
+    }
+
+    /**
+     * Collapses the given item and removes its sub-hierarchy. Calling this
+     * method will have no effect if the row is already collapsed. The index is
+     * provided by the client-side or calculated from a full data request.
+     * {@code syncAndRefresh} indicates whether the changes should be
+     * synchronised to the client and the data provider be notified.
+     * 
+     * @param item
+     *            the item to collapse
+     * @param index
+     *            the index of the item
+     * @param syncAndRefresh
+     *            {@code true} if the changes should be synchronised to the
+     *            client and the data provider should be notified of the
+     *            changes, {@code false} otherwise.
+     */
+    private void doCollapse(T item, Integer index, boolean syncAndRefresh) {
+        Range removedRows = mapper.collapse(item, index);
+        if (syncAndRefresh) {
             if (!reset && !removedRows.isEmpty()) {
                 getClientRpc().removeRows(removedRows.getStart(),
                         removedRows.length());
@@ -193,44 +243,97 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
 
     /**
      * Expands the given item. Calling this method will have no effect if the
-     * row is already expanded.
+     * item is already expanded or if it has no children.
      *
      * @param item
      *            the item to expand
      */
     public void expand(T item) {
-        if (!mapper.isExpanded(item) && mapper.hasChildren(item)) {
-            doExpand(item, mapper.getIndexOf(item));
-        }
+        expand(item, true);
     }
 
     /**
-     * Expands the given item at given index. Calling this method will have no
-     * effect if the row is already expanded. The index is provided by the
-     * client-side or calculated from a full data request.
+     * Expands the given item. Calling this method will have no effect if the
+     * item is already expanded or if it has no children. {@code syncAndRefresh}
+     * indicates whether the changes should be synchronised to the client and
+     * the data provider be notified.
      *
-     * @see #expand(Object)
+     * @param item
+     *            the item to expand
+     * @param syncAndRefresh
+     *            {@code true} if the changes should be synchronised to the
+     *            client and the data provider should be notified of the
+     *            changes, {@code
+     *         false} otherwise.
+     */
+    public void expand(T item, boolean syncAndRefresh) {
+        Integer index = syncAndRefresh ? mapper.getIndexOf(item).orElse(null) : null;
+        doExpand(item, index, syncAndRefresh);
+    }
+
+    /**
+     * Expands the given item at the given index. Calling this method will have
+     * no effect if the item is already expanded.
      *
      * @param item
      *            the item to expand
      * @param index
      *            the index of the item
      */
-    public void doExpand(T item, Optional<Integer> index) {
-        if (!mapper.isExpanded(item)) {
-            Range addedRows = mapper.doExpand(item, index);
+    public void expand(T item, Integer index) {
+        doExpand(item, index, true);
+    }
+
+    /**
+     * Expands the given item. Calling this method will have no effect if the
+     * item is already expanded or if it has no children. The index is provided
+     * by the client-side or calculated from a full data request.
+     * {@code syncAndRefresh} indicates whether the changes should be
+     * synchronised to the client and the data provider be notified.
+     *
+     * @param item
+     *            the item to expand
+     * @param index
+     *            the index of the item
+     * @param syncAndRefresh
+     *            {@code true} if the changes should be synchronised to the
+     *            client and the data provider should be notified of the
+     *            changes, {@code false} otherwise.
+     */
+    private void doExpand(T item, Integer index, boolean syncAndRefresh) {
+        Range addedRows = mapper.expand(item, index);
+        if (syncAndRefresh) {
             if (!reset && !addedRows.isEmpty()) {
-                int start = addedRows.getStart();
-                getClientRpc().insertRows(start, addedRows.length());
-                Stream<T> children = mapper.fetchItems(item,
-                        Range.withLength(0, addedRows.length()));
-                pushData(start, children.collect(Collectors.toList()));
+                getClientRpc()
+                        .insertRows(addedRows.getStart(), addedRows.length());
+                Stream<T> children = mapper
+                        .fetchItems(item,
+                                Range.withLength(0, addedRows.length()));
+                pushData(addedRows.getStart(),
+                        children.collect(Collectors.toList()));
             }
             refresh(item);
         }
     }
 
     /**
+     * Expands the given item at given index. Calling this method will have no
+     * effect if the row is already expanded. The index is provided by the
+     * client-side or calculated from a full data request.
+     *
+     * @param item
+     *            the item to expand
+     * @param index
+     *            the index of the item
+     * @see #expand(Object)
+     * @deprecated use {@link #expand(Object, Integer)} instead
+     */
+    @Deprecated
+    public void doExpand(T item, Optional<Integer> index) {
+        expand(item, index.orElse(null));
+    }
+
+    /**
      * Returns whether given item has children.
      *
      * @param item
diff --git a/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java b/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
index e918c4a209..f6f9f8b9a4 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
@@ -123,46 +123,89 @@ public class HierarchyMapper<T, F> implements DataGenerator<T> {
      * @param item
      *            the item to expand
      * @param position
+     *            the index of the item
+     * @return range of rows added by expanding the item
+     */
+    public Range expand(T item, Integer position) {
+        if (doExpand(item) && position != null) {
+            return Range.withLength(position + 1,
+                    (int) getHierarchy(item, false).count());
+        }
+
+        return Range.emptyRange();
+    }
+
+    /**
+     * Expands the given item.
+     *
+     * @param item
+     *            the item to expand
+     * @param position
      *            the index of item
      * @return range of rows added by expanding the item
+     * @deprecated Use {@link #expand(Object, Integer)} instead.
      */
+    @Deprecated
     public Range doExpand(T item, Optional<Integer> position) {
-        Range rows = Range.withLength(0, 0);
+        return expand(item, position.orElse(null));
+    }
+
+    /**
+     * Expands the given item if it is collapsed and has children, and returns
+     * whether this method expanded the item.
+     *
+     * @param item
+     *            the item to expand
+     * @return {@code true} if this method expanded the item, {@code false}
+     *         otherwise
+     */
+    private boolean doExpand(T item) {
+        boolean expanded = false;
         if (!isExpanded(item) && hasChildren(item)) {
-            Object id = getDataProvider().getId(item);
-            expandedItemIds.add(id);
-            if (position.isPresent()) {
-                rows = Range.withLength(position.get() + 1,
-                        (int) getHierarchy(item, false).count());
-            }
+            expandedItemIds.add(getDataProvider().getId(item));
+            expanded = true;
         }
-        return rows;
+        return expanded;
     }
 
     /**
      * Collapses the given item.
      *
      * @param item
-     *            the item to expand
+     *            the item to collapse
      * @param position
-     *            the index of item
+     *            the index of the item
      *
      * @return range of rows removed by collapsing the item
      */
-    public Range doCollapse(T item, Optional<Integer> position) {
-        Range removedRows = Range.withLength(0, 0);
+    public Range collapse(T item, Integer position) {
+        Range removedRows = Range.emptyRange();
         if (isExpanded(item)) {
-            Object id = getDataProvider().getId(item);
-            if (position.isPresent()) {
-                long childCount = getHierarchy(item, false).count();
-                removedRows = Range.withLength(position.get() + 1,
-                        (int) childCount);
+            if (position != null) {
+                removedRows = Range.withLength(position + 1,
+                        (int) getHierarchy(item, false).count());
             }
-            expandedItemIds.remove(id);
+            expandedItemIds.remove(getDataProvider().getId(item));
         }
         return removedRows;
     }
 
+    /**
+     * Collapses the given item.
+     *
+     * @param item
+     *            the item to collapse
+     * @param position
+     *            the index of item
+     *
+     * @return range of rows removed by collapsing the item
+     * @deprecated Use {@link #collapse(Object, Integer)} instead.
+     */
+    @Deprecated
+    public Range doCollapse(T item, Optional<Integer> position) {
+        return collapse(item, position.orElse(null));
+    }
+
     @Override
     public void generateData(T item, JsonObject jsonObject) {
         JsonObject hierarchyData = Json.createObject();
diff --git a/server/src/main/java/com/vaadin/ui/Tree.java b/server/src/main/java/com/vaadin/ui/Tree.java
index 4e96f4a3e4..bebe6972b5 100644
--- a/server/src/main/java/com/vaadin/ui/Tree.java
+++ b/server/src/main/java/com/vaadin/ui/Tree.java
@@ -442,6 +442,25 @@ public class Tree<T> extends Composite
     }
 
     /**
+     * Expands the given items and their children recursively until the given
+     * depth.
+     * <p>
+     * {@code depth} describes the maximum distance between a given item and its
+     * descendant, meaning that {@code expandRecursively(items, 0)} expands only
+     * the given items while {@code expandRecursively(items, 2)} expands the
+     * given items as well as their children and grandchildren.
+     *
+     * @param items
+     *            the items to expand recursively
+     * @param depth
+     *            the maximum depth of recursion
+     * @since
+     */
+    public void expandRecursively(Collection<T> items, int depth) {
+        treeGrid.expandRecursively(items, depth);
+    }
+
+    /**
      * Collapse the given items.
      * <p>
      * For items that are already collapsed, does nothing.
@@ -466,6 +485,25 @@ public class Tree<T> extends Composite
     }
 
     /**
+     * Collapse the given items and their children recursively until the given
+     * depth.
+     * <p>
+     * {@code depth} describes the maximum distance between a given item and its
+     * descendant, meaning that {@code collapseRecursively(items, 0)} collapses
+     * only the given items while {@code collapseRecursively(items, 2)}
+     * collapses the given items as well as their children and grandchildren.
+     *
+     * @param items
+     *            the items to expand recursively
+     * @param depth
+     *            the maximum depth of recursion
+     * @since
+     */
+    public void collapseRecursively(Collection<T> items, int depth) {
+        treeGrid.collapseRecursively(items, depth);
+    }
+
+    /**
      * Returns whether a given item is expanded or collapsed.
      *
      * @param item
diff --git a/server/src/main/java/com/vaadin/ui/TreeGrid.java b/server/src/main/java/com/vaadin/ui/TreeGrid.java
index 65329d24c8..1c5716cc61 100644
--- a/server/src/main/java/com/vaadin/ui/TreeGrid.java
+++ b/server/src/main/java/com/vaadin/ui/TreeGrid.java
@@ -175,12 +175,12 @@ public class TreeGrid<T> extends Grid<T>
                 userOriginated) -> {
             T item = getDataCommunicator().getKeyMapper().get(rowKey);
             if (collapse && getDataCommunicator().isExpanded(item)) {
-                getDataCommunicator().doCollapse(item, Optional.of(rowIndex));
+                getDataCommunicator().collapse(item, rowIndex);
                 fireCollapseEvent(
                         getDataCommunicator().getKeyMapper().get(rowKey),
                         userOriginated);
             } else if (!collapse && !getDataCommunicator().isExpanded(item)) {
-                getDataCommunicator().doExpand(item, Optional.of(rowIndex));
+                getDataCommunicator().expand(item, rowIndex);
                 fireExpandEvent(
                         getDataCommunicator().getKeyMapper().get(rowKey),
                         userOriginated);
@@ -345,6 +345,64 @@ public class TreeGrid<T> extends Grid<T>
     }
 
     /**
+     * Expands the given items and their children recursively until the given
+     * depth.
+     * <p>
+     * {@code depth} describes the maximum distance between a given item and its
+     * descendant, meaning that {@code expandRecursively(items, 0)} expands only
+     * the given items while {@code expandRecursively(items, 2)} expands the
+     * given items as well as their children and grandchildren.
+     * <p>
+     * This method will <i>not</i> fire events for expanded nodes.
+     *
+     * @param items
+     *            the items to expand recursively
+     * @param depth
+     *            the maximum depth of recursion
+     * @since
+     */
+    public void expandRecursively(Collection<T> items, int depth) {
+        expandRecursively(items.stream(), depth);
+    }
+
+    /**
+     * Expands the given items and their children recursively until the given
+     * depth.
+     * <p>
+     * {@code depth} describes the maximum distance between a given item and its
+     * descendant, meaning that {@code expandRecursively(items, 0)} expands only
+     * the given items while {@code expandRecursively(items, 2)} expands the
+     * given items as well as their children and grandchildren.
+     * <p>
+     * This method will <i>not</i> fire events for expanded nodes.
+     *
+     * @param items
+     *            the items to expand recursively
+     * @param depth
+     *            the maximum depth of recursion
+     * @since
+     */
+    public void expandRecursively(Stream<T> items, int depth) {
+        if (depth < 0) {
+            return;
+        }
+
+        HierarchicalDataCommunicator<T> communicator = getDataCommunicator();
+        items.forEach(item -> {
+            if (communicator.hasChildren(item)) {
+                communicator.expand(item, false);
+
+                expandRecursively(
+                        getDataProvider().fetchChildren(
+                                new HierarchicalQuery<>(null, item)),
+                        depth - 1);
+            }
+        });
+
+        getDataProvider().refreshAll();
+    }
+
+    /**
      * Collapse the given items.
      * <p>
      * For items that are already collapsed, does nothing.
@@ -375,6 +433,64 @@ public class TreeGrid<T> extends Grid<T>
     }
 
     /**
+     * Collapse the given items and their children recursively until the given
+     * depth.
+     * <p>
+     * {@code depth} describes the maximum distance between a given item and its
+     * descendant, meaning that {@code collapseRecursively(items, 0)} collapses
+     * only the given items while {@code collapseRecursively(items, 2)}
+     * collapses the given items as well as their children and grandchildren.
+     * <p>
+     * This method will <i>not</i> fire events for collapsed nodes.
+     *
+     * @param items
+     *            the items to collapse recursively
+     * @param depth
+     *            the maximum depth of recursion
+     * @since
+     */
+    public void collapseRecursively(Collection<T> items, int depth) {
+        collapseRecursively(items.stream(), depth);
+    }
+
+    /**
+     * Collapse the given items and their children recursively until the given
+     * depth.
+     * <p>
+     * {@code depth} describes the maximum distance between a given item and its
+     * descendant, meaning that {@code collapseRecursively(items, 0)} collapses
+     * only the given items while {@code collapseRecursively(items, 2)}
+     * collapses the given items as well as their children and grandchildren.
+     * <p>
+     * This method will <i>not</i> fire events for collapsed nodes.
+     *
+     * @param items
+     *            the items to collapse recursively
+     * @param depth
+     *            the maximum depth of recursion
+     * @since
+     */
+    public void collapseRecursively(Stream<T> items, int depth) {
+        if (depth < 0) {
+            return;
+        }
+
+        HierarchicalDataCommunicator<T> communicator = getDataCommunicator();
+        items.forEach(item -> {
+            if (communicator.hasChildren(item)) {
+                collapseRecursively(
+                        getDataProvider().fetchChildren(
+                                new HierarchicalQuery<>(null, item)),
+                        depth - 1);
+
+                communicator.collapse(item, false);
+            }
+        });
+
+        getDataProvider().refreshAll();
+    }
+
+    /**
      * Returns whether a given item is expanded or collapsed.
      *
      * @param item
diff --git a/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithDataTest.java b/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithDataTest.java
index 80d6a70cfd..81ee908a52 100644
--- a/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithDataTest.java
+++ b/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithDataTest.java
@@ -209,11 +209,11 @@ public class HierarchyMapperWithDataTest {
     }
 
     private void expand(Node node) {
-        insertRows(mapper.doExpand(node, mapper.getIndexOf(node)));
+        insertRows(mapper.expand(node, mapper.getIndexOf(node).orElse(null)));
     }
 
     private void collapse(Node node) {
-        removeRows(mapper.doCollapse(node, mapper.getIndexOf(node)));
+        removeRows(mapper.collapse(node, mapper.getIndexOf(node).orElse(null)));
     }
 
     private void verifyFetchIsCorrect(List<Node> expectedResult, Range range) {
diff --git a/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithNumerousDataTest.java b/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithNumerousDataTest.java
index 6c0091f8bf..58208a6cc3 100644
--- a/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithNumerousDataTest.java
+++ b/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithNumerousDataTest.java
@@ -63,7 +63,7 @@ public class HierarchyMapperWithNumerousDataTest {
     }
 
     private void expand(Node node) {
-        insertRows(mapper.doExpand(node, mapper.getIndexOf(node)));
+        insertRows(mapper.expand(node, mapper.getIndexOf(node).orElse(null)));
     }
 
     public void insertRows(Range range) {
diff --git a/shared/src/main/java/com/vaadin/shared/Range.java b/shared/src/main/java/com/vaadin/shared/Range.java
index 042cc527a0..b47963fd50 100644
--- a/shared/src/main/java/com/vaadin/shared/Range.java
+++ b/shared/src/main/java/com/vaadin/shared/Range.java
@@ -31,6 +31,9 @@ import java.io.Serializable;
  * @author Vaadin Ltd
  */
 public final class Range implements Serializable {
+
+    private static final Range EMPTY = Range.withLength(0, 0);
+
     private final int start;
     private final int end;
 
@@ -90,6 +93,10 @@ public final class Range implements Serializable {
         return new Range(start, start + length);
     }
 
+    public static Range emptyRange() {
+        return EMPTY;
+    }
+
     /**
      * Creates a new range between two numbers: <code>[start..end[</code>.
      *
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
index a468f34d29..014ee7d1f0 100644
--- a/uitest/src/main/java/com/vaadin/tests/components/treegrid/LazyHierarchicalDataProvider.java
+++ b/uitest/src/main/java/com/vaadin/tests/components/treegrid/LazyHierarchicalDataProvider.java
@@ -49,7 +49,8 @@ public class LazyHierarchicalDataProvider extends
                 .flatMap(parent -> Optional.of(parent.getId()));
 
         List<HierarchicalTestBean> list = new ArrayList<>();
-        for (int i = 0; i < query.getLimit(); i++) {
+        int limit = Math.min(query.getLimit(), nodesPerLevel);
+        for (int i = 0; i < limit; i++) {
             list.add(new HierarchicalTestBean(parentKey.orElse(null), depth,
                     i + query.getOffset()));
         }
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 2ac0c0673d..c12baf3900 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
@@ -180,6 +180,11 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
         createClickAction("Expand 2 | 1", "Server-side expand",
                 (treeGrid, value, data) -> treeGrid.expand(value),
                 new HierarchicalTestBean("/0/0/1/1", 2, 1));
+
+        createClickAction("Expand 0 | 0 recursively", "Server-side expand",
+                (treeGrid, value, data) -> treeGrid
+                        .expandRecursively(Arrays.asList(value), 1),
+                new HierarchicalTestBean(null, 0, 0));
     }
 
     @SuppressWarnings("unchecked")
@@ -194,6 +199,11 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
         createClickAction("Collapse 2 | 1", "Server-side collapse",
                 (treeGrid, value, data) -> treeGrid.collapse(value),
                 new HierarchicalTestBean("/0/0/1/1", 2, 1));
+
+        createClickAction("Collapse 0 | 0 recursively", "Server-side collapse",
+                (treeGrid, value, data) -> treeGrid
+                        .collapseRecursively(Arrays.asList(value), 2),
+                new HierarchicalTestBean(null, 0, 0));
     }
 
     @SuppressWarnings("unchecked")
diff --git a/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridExpandCollapseRecursively.java b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridExpandCollapseRecursively.java
new file mode 100644
index 0000000000..f85eba5815
--- /dev/null
+++ b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridExpandCollapseRecursively.java
@@ -0,0 +1,106 @@
+package com.vaadin.tests.components.treegrid;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
+
+import com.vaadin.annotations.Widgetset;
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.tests.components.AbstractTestUI;
+import com.vaadin.ui.Button;
+import com.vaadin.ui.HorizontalLayout;
+import com.vaadin.ui.RadioButtonGroup;
+import com.vaadin.ui.TreeGrid;
+
+@Widgetset("com.vaadin.DefaultWidgetSet")
+public class TreeGridExpandCollapseRecursively extends AbstractTestUI {
+
+    private static class Directory {
+
+        private String name;
+        private Directory parent;
+        private List<Directory> subDirectories = new ArrayList<>();
+
+        public Directory(String name, Directory parent) {
+            this.name = name;
+            this.parent = parent;
+        }
+
+        public String getName() {
+            return name;
+        }
+
+        public void setName(String name) {
+            this.name = name;
+        }
+
+        public Directory getParent() {
+            return parent;
+        }
+
+        public void setParent(
+                Directory parent) {
+            this.parent = parent;
+        }
+
+        public List<Directory> getSubDirectories() {
+            return subDirectories;
+        }
+
+        public void setSubDirectories(List<Directory> subDirectories) {
+            this.subDirectories = subDirectories;
+        }
+    }
+
+    private static final int DEPTH = 4;
+    private static final int CHILDREN = 5;
+
+    @Override
+    protected void setup(VaadinRequest request) {
+
+        Collection<Directory> roots = generateDirectoryStructure(DEPTH);
+
+        TreeGrid<Directory> grid = new TreeGrid<>();
+        grid.addColumn(item -> "Item" + item.getName());
+
+        grid.setItems(roots, Directory::getSubDirectories);
+
+        RadioButtonGroup<Integer> depthSelector = new RadioButtonGroup<>(
+                "Depth", Arrays.asList(0, 1, 2, 3));
+        depthSelector.addStyleName("horizontal");
+        depthSelector.setValue(3);
+
+        HorizontalLayout buttons = new HorizontalLayout();
+        buttons.addComponent(new Button("Expand recursively", e -> grid
+                .expandRecursively(roots, depthSelector.getValue())));
+        buttons.addComponent(new Button("Collapse recursively", e -> grid
+                .collapseRecursively(roots, depthSelector.getValue())));
+
+        addComponents(depthSelector, buttons, grid);
+    }
+
+    private Collection<Directory> generateDirectoryStructure(int depth) {
+        return generateDirectories(depth, null, CHILDREN);
+    }
+
+    private Collection<Directory> generateDirectories(int depth,
+            Directory parent, int childCount) {
+        Collection<Directory> dirs = new ArrayList<>();
+        if (depth >= 0) {
+            for (int i = 0; i < childCount; i++) {
+                String name = parent != null
+                        ? parent.getName() + "-" + i
+                        : "-" + i;
+                Directory dir = new Directory(name, parent);
+                if (parent != null) {
+                    parent.getSubDirectories().add(dir);
+                }
+                dirs.add(dir);
+
+                generateDirectories(depth - 1, dir, childCount);
+            }
+        }
+        return dirs;
+    }
+}
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 43c08415df..a96e2d3c1d 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
@@ -95,6 +95,25 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
         assertEquals(3, grid.getRowCount());
         assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
 
+        // expand 0 | 0 recursively
+        selectMenuPath("Component", "Features", "Server-side expand",
+                "Expand 0 | 0 recursively");
+        assertEquals(15, grid.getRowCount());
+        assertCellTexts(0, 0, new String[] { "0 | 0", "1 | 0", "2 | 0" });
+
+        // collapse 0 | 0 recursively
+        selectMenuPath("Component", "Features", "Server-side collapse",
+                "Collapse 0 | 0 recursively");
+        assertEquals(3, grid.getRowCount());
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
+
+        // expanding 0 | 0 should result in 3 additional nodes after recursive
+        // collapse
+        selectMenuPath("Component", "Features", "Server-side expand",
+                "Expand 0 | 0");
+        assertEquals(6, grid.getRowCount());
+        assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "1 | 2" });
+
         assertNoSystemNotifications();
         assertNoErrorNotifications();
     }
diff --git a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridExpandCollapseRecursivelyTest.java b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridExpandCollapseRecursivelyTest.java
new file mode 100644
index 0000000000..6e2451b96a
--- /dev/null
+++ b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridExpandCollapseRecursivelyTest.java
@@ -0,0 +1,105 @@
+package com.vaadin.tests.components.treegrid;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import com.vaadin.testbench.elements.ButtonElement;
+import com.vaadin.testbench.elements.RadioButtonGroupElement;
+import com.vaadin.testbench.elements.TreeGridElement;
+import com.vaadin.tests.tb3.SingleBrowserTest;
+
+import static org.junit.Assert.assertEquals;
+
+public class TreeGridExpandCollapseRecursivelyTest extends SingleBrowserTest {
+
+    private static final int rowCount0 = 5;
+    private static final int rowCount1 = rowCount0 + rowCount0 * 5;
+    private static final int rowCount2 =
+            rowCount1 + (rowCount1 - rowCount0) * 5;
+    private static final int rowCount3 =
+            rowCount2 + (rowCount2 - rowCount1) * 5;
+    private static final int rowCount4 =
+            rowCount3 + (rowCount3 - rowCount2) * 5;
+
+    private TreeGridElement grid;
+    private RadioButtonGroupElement depthSelector;
+    private ButtonElement expandButton;
+    private ButtonElement collapseButton;
+
+    @Before
+    public void before() {
+        openTestURL();
+        grid = $(TreeGridElement.class).first();
+        depthSelector = $(RadioButtonGroupElement.class).first();
+        expandButton = $(ButtonElement.class).get(0);
+        collapseButton = $(ButtonElement.class).get(1);
+    }
+
+    @Test
+    public void expandVariousDepth() {
+        assertEquals(rowCount0, grid.getRowCount());
+
+        selectDepth(0);
+        expandButton.click();
+
+        assertEquals(rowCount1, grid.getRowCount());
+
+        selectDepth(1);
+        expandButton.click();
+
+        assertEquals(rowCount2, grid.getRowCount());
+
+        selectDepth(2);
+        expandButton.click();
+
+        assertEquals(rowCount3, grid.getRowCount());
+
+        selectDepth(3);
+        expandButton.click();
+
+        assertEquals(rowCount4, grid.getRowCount());
+    }
+
+    @Test(timeout = 5000)
+    public void expandAndCollapseAllItems() {
+        assertEquals(rowCount0, grid.getRowCount());
+
+        selectDepth(3);
+        expandButton.click();
+
+        assertEquals(rowCount4, grid.getRowCount());
+
+        collapseButton.click();
+
+        assertEquals(rowCount0, grid.getRowCount());
+    }
+
+    @Test
+    public void partialCollapse() {
+        assertEquals(rowCount0, grid.getRowCount());
+
+        selectDepth(3);
+        expandButton.click();
+
+        assertEquals(rowCount4, grid.getRowCount());
+
+        selectDepth(1);
+        collapseButton.click();
+
+        assertEquals(rowCount0, grid.getRowCount());
+
+        selectDepth(0);
+        expandButton.click();
+
+        assertEquals(rowCount1, grid.getRowCount());
+
+        // Open just one subtree to see if it is still fully expanded
+        grid.getExpandElement(2, 0).click();
+
+        assertEquals(rowCount1 + rowCount2, grid.getRowCount());
+    }
+
+    private void selectDepth(int depth) {
+        depthSelector.setValue(String.valueOf(depth));
+    }
+}