Add server-side expand and collapse to TreeGrid (#9021)

* Add server-side expand and collapse to TreeGrid

* Add javadocs

* Fix variable naming in TreeGridHugeTreeTest

* Fix review comments

* Merge remote-tracking branch 'github/master' into 8759-server-expand

* Clear pending expands when all data is dropped

* Add documentation

13 files changed, 588 insertions, 43 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 d8e91ceaa0..10935264c7 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
@@ -16,7 +16,10 @@
 package com.vaadin.client.connectors.treegrid;
 
 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;
@@ -25,6 +28,9 @@ import com.google.gwt.event.dom.client.KeyCodes;
 import com.google.gwt.user.client.Event;
 import com.vaadin.client.annotations.OnStateChange;
 import com.vaadin.client.connectors.grid.GridConnector;
+import com.vaadin.client.data.AbstractRemoteDataSource;
+import com.vaadin.client.data.DataChangeHandler;
+import com.vaadin.client.data.DataSource;
 import com.vaadin.client.renderers.HierarchyRenderer;
 import com.vaadin.client.widget.grid.EventCellReference;
 import com.vaadin.client.widget.grid.GridEventHandler;
@@ -32,10 +38,13 @@ import com.vaadin.client.widget.grid.events.GridClickEvent;
 import com.vaadin.client.widget.treegrid.TreeGrid;
 import com.vaadin.client.widget.treegrid.events.TreeGridClickEvent;
 import com.vaadin.client.widgets.Grid;
+import com.vaadin.shared.Range;
+import com.vaadin.shared.data.DataCommunicatorConstants;
 import com.vaadin.shared.ui.Connect;
 import com.vaadin.shared.ui.treegrid.FocusParentRpc;
 import com.vaadin.shared.ui.treegrid.FocusRpc;
 import com.vaadin.shared.ui.treegrid.NodeCollapseRpc;
+import com.vaadin.shared.ui.treegrid.TreeGridClientRpc;
 import com.vaadin.shared.ui.treegrid.TreeGridCommunicationConstants;
 import com.vaadin.shared.ui.treegrid.TreeGridState;
 
@@ -60,6 +69,8 @@ public class TreeGridConnector extends GridConnector {
 
     private HierarchyRenderer hierarchyRenderer;
 
+    private Set<String> rowKeysPendingExpand = new HashSet<>();
+
     @Override
     public TreeGrid getWidget() {
         return (TreeGrid) super.getWidget();
@@ -145,6 +156,54 @@ public class TreeGridConnector extends GridConnector {
         // widget check
         replaceClickEvent(getWidget(),
                 new TreeGridClickEvent(getWidget(), getEventCell(getWidget())));
+
+        registerRpc(TreeGridClientRpc.class, new TreeGridClientRpc() {
+
+            @Override
+            public void setExpanded(String key) {
+                rowKeysPendingExpand.add(key);
+                Range cache = ((AbstractRemoteDataSource) getDataSource())
+                        .getCachedRange();
+                checkExpand(cache.getStart(), cache.length());
+            }
+
+            @Override
+            public void setCollapsed(String key) {
+                rowKeysPendingExpand.remove(key);
+            }
+        });
+    }
+
+    @Override
+    public void setDataSource(DataSource<JsonObject> dataSource) {
+        super.setDataSource(dataSource);
+        dataSource.addDataChangeHandler(new DataChangeHandler() {
+
+            @Override
+            public void dataUpdated(int firstRowIndex, int numberOfRows) {
+                checkExpand(firstRowIndex, numberOfRows);
+            }
+
+            @Override
+            public void dataRemoved(int firstRowIndex, int numberOfRows) {
+                // NO-OP
+            }
+
+            @Override
+            public void dataAdded(int firstRowIndex, int numberOfRows) {
+                // NO-OP
+            }
+
+            @Override
+            public void dataAvailable(int firstRowIndex, int numberOfRows) {
+                // NO-OP
+            }
+
+            @Override
+            public void resetDataAndSize(int estimatedNewDataSize) {
+                rowKeysPendingExpand.clear();
+            }
+        });
     }
 
     private native void replaceCellFocusEventHandler(Grid<?> grid,
@@ -173,7 +232,13 @@ public class TreeGridConnector extends GridConnector {
     private void setCollapsed(int rowIndex, boolean collapsed) {
         String rowKey = getRowKey(getDataSource().getRow(rowIndex));
         getRpcProxy(NodeCollapseRpc.class).setNodeCollapsed(rowKey, rowIndex,
-                collapsed);
+                collapsed, true);
+    }
+
+    private void setCollapsedServerInitiated(int rowIndex, boolean collapsed) {
+        String rowKey = getRowKey(getDataSource().getRow(rowIndex));
+        getRpcProxy(NodeCollapseRpc.class).setNodeCollapsed(rowKey, rowIndex,
+                collapsed, false);
     }
 
     /**
@@ -279,6 +344,20 @@ public class TreeGridConnector extends GridConnector {
         }
     }
 
+    private void checkExpand(int firstRowIndex, int numberOfRows) {
+        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);
+                    }
+                });
+    }
+
     private static boolean isCollapsed(JsonObject rowData) {
         assert rowData
                 .hasKey(TreeGridCommunicationConstants.ROW_HIERARCHY_DESCRIPTION) : "missing hierarchy data for row "
diff --git a/documentation/components/components-treegrid.asciidoc b/documentation/components/components-treegrid.asciidoc
index 98d3d17839..77f82d1c15 100644
--- a/documentation/components/components-treegrid.asciidoc
+++ b/documentation/components/components-treegrid.asciidoc
@@ -53,7 +53,7 @@ treeGrid.setDataProvider(new InMemoryHierarchicalDataProvider<>(data));
 // the first column gets the hierarchy indicator by default
 treeGrid.addColumn(Project::getName).setCaption("Project Name");
 treeGrid.addColumn(Project::getHoursDone).setCaption("Hours Done");
-treeGrid.addColumn(Project::getdLastModified).setCaption("Last Modified");
+treeGrid.addColumn(Project::getLastModified).setCaption("Last Modified");
 ----
 
 The [classname]#HierarchyData# class can be used to build the hierarchical data structure,
@@ -78,6 +78,26 @@ dataProvider.refreshAll();
 
 Note that for adding or removing nodes, you always need to call the [methodname]#refreshAll# method in the data provider you are using. The [methodname]#refreshItem# method can only be used when just the data for that item is updated, but not for updates that add or remove items.
 
+[[components.treegrid.expandcollapse]]
+== Expanding and Collapsing Nodes
+
+[classname]#TreeGrid# nodes that have children can be expanded and collapsed by either user interaction or through the server-side API:
+
+[source, java]
+----
+// Expands a child project. If the child project is not yet
+// in the visible hierarchy, nothing will be shown.
+treeGrid.expand(childProject);
+// Expands the root project. If child project now becomes
+// visible it is also expanded into view.
+treeGrid.expand(rootProject);
+// Collapses the child project.
+treeGrid.collapse(childProject);
+----
+
+To use the server-side API with a backend data provider the [methodname]#hashCode# method for the node's type must be implemented so that when the desired node is  retrieved from the backend it has the same hash as the object passed to either [methodname]#expand# or [methodname]#collapse#.
+
+[[components.treegrid.hierarchycolumn]]
 == Changing the Hierarchy Column
 
 By default, the [classname]#TreeGrid# shows the hierarchy indicator by default in the first column of the grid.
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 ae3c6f4cc1..ffdd79f36e 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
@@ -463,7 +463,10 @@ public class DataCommunicator<T> extends AbstractExtension {
         }
     }
 
-    private void dropAllData() {
+    /**
+     * Drops all data associated with this data communicator.
+     */
+    protected void dropAllData() {
         for (DataGenerator<T> g : generators) {
             g.destroyAllData();
         }
diff --git a/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java b/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java
index 0ec6748a85..f5be454f86 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java
+++ b/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java
@@ -39,6 +39,16 @@ public interface DataKeyMapper<T> extends Serializable {
     String key(T dataObject);
 
     /**
+     * Check whether this key mapper contains the given data object.
+     *
+     * @param dataObject
+     *            the data object to check
+     * @return {@code true} if the given data object is contained in this key
+     *         mapper, {@code false} otherwise
+     */
+    boolean has(T dataObject);
+
+    /**
      * Gets the data object identified by given key.
      *
      * @param key
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 b986f74ea3..fedad9dd55 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
@@ -17,8 +17,11 @@ package com.vaadin.data.provider;
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
 import java.util.logging.Logger;
@@ -27,6 +30,7 @@ import java.util.stream.Stream;
 
 import com.vaadin.data.HierarchyData;
 import com.vaadin.data.provider.HierarchyMapper.TreeLevelQuery;
+import com.vaadin.data.provider.HierarchyMapper.TreeNode;
 import com.vaadin.server.SerializableConsumer;
 import com.vaadin.server.SerializablePredicate;
 import com.vaadin.shared.Range;
@@ -58,6 +62,8 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
 
     private HierarchyMapper mapper = new HierarchyMapper();
 
+    private Set<String> rowKeysPendingExpand = new HashSet<>();
+
     /**
      * Collapse allowed provider used to allow/disallow collapsing nodes.
      */
@@ -278,13 +284,19 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
             // cannot drop expanded rows since the parent item is needed always
             // when fetching more rows
             String itemKey = keys.getString(i);
-            if (mapper.isCollapsed(itemKey)) {
+            if (mapper.isCollapsed(itemKey) && !rowKeysPendingExpand.contains(itemKey)) {
                 getActiveDataHandler().dropActiveData(itemKey);
             }
         }
     }
 
     @Override
+    protected void dropAllData() {
+        super.dropAllData();
+        rowKeysPendingExpand.clear();
+    }
+
+    @Override
     public HierarchicalDataProvider<T, ?> getDataProvider() {
         return (HierarchicalDataProvider<T, ?>) super.getDataProvider();
     }
@@ -400,6 +412,8 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
             return false;
         }
         mapper.expand(expandedRowKey, expandedRowIndex, expandedNodeSize);
+        rowKeysPendingExpand.remove(expandedRowKey);
+
         getClientRpc().insertRows(expandedRowIndex + 1, expandedNodeSize);
         // TODO optimize by sending "just enough" of the expanded items
         // directly
@@ -414,6 +428,68 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
     }
 
     /**
+     * Set an item as pending expansion.
+     * <p>
+     * Calling this method reserves a communication key for the item that is
+     * guaranteed to not be invalidated until the item is expanded. Has no
+     * effect and returns an empty optional if the given item is already
+     * expanded or has no children.
+     *
+     * @param item
+     *            the item to set as pending expansion
+     * @return an optional of the communication key used for the item, empty if
+     *         the item cannot be expanded
+     */
+    public Optional<String> setPendingExpand(T item) {
+        Objects.requireNonNull(item, "Item cannot be null");
+        if (getKeyMapper().has(item)
+                && !mapper.isCollapsed(getKeyMapper().key(item))) {
+            // item is already expanded
+            return Optional.empty();
+        }
+        if (!getDataProvider().hasChildren(item)) {
+            // ignore item with no children
+            return Optional.empty();
+        }
+        String key = getKeyMapper().key(item);
+        rowKeysPendingExpand.add(key);
+        return Optional.of(key);
+    }
+
+    /**
+     * Collapse an item.
+     * <p>
+     * This method will either collapse an item directly, or remove its pending
+     * expand status. If the item is not expanded or pending expansion, calling
+     * this method has no effect.
+     *
+     * @param item
+     *            the item to collapse
+     * @return an optional of the communication key used for the item, empty if
+     *         the item cannot be collapsed
+     */
+    public Optional<String> collapseItem(T item) {
+        Objects.requireNonNull(item, "Item cannot be null");
+        if (!getKeyMapper().has(item)) {
+            // keymapper should always have items that are expanded or pending
+            // expand
+            return Optional.empty();
+        }
+        String nodeKey = getKeyMapper().key(item);
+        Optional<TreeNode> node = mapper.getNodeForKey(nodeKey);
+        if (node.isPresent()) {
+            rowKeysPendingExpand.remove(nodeKey);
+            doCollapse(nodeKey, node.get().getStartIndex() - 1);
+            return Optional.of(nodeKey);
+        }
+        if (rowKeysPendingExpand.contains(nodeKey)) {
+            rowKeysPendingExpand.remove(nodeKey);
+            return Optional.of(nodeKey);
+        }
+        return Optional.empty();
+    }
+
+    /**
      * Sets the item collapse allowed provider for this
      * HierarchicalDataCommunicator. The provider should return {@code true} for
      * any item that the user can collapse.
diff --git a/server/src/main/java/com/vaadin/server/KeyMapper.java b/server/src/main/java/com/vaadin/server/KeyMapper.java
index 7ed723be66..58058c6f69 100644
--- a/server/src/main/java/com/vaadin/server/KeyMapper.java
+++ b/server/src/main/java/com/vaadin/server/KeyMapper.java
@@ -64,6 +64,11 @@ public class KeyMapper<V> implements DataKeyMapper<V>, Serializable {
         return key;
     }
 
+    @Override
+    public boolean has(V o) {
+        return objectKeyMap.containsKey(o);
+    }
+
     /**
      * Retrieves object with the key.
      *
diff --git a/server/src/main/java/com/vaadin/ui/TreeGrid.java b/server/src/main/java/com/vaadin/ui/TreeGrid.java
index d91a0f9af1..b5da3f8645 100644
--- a/server/src/main/java/com/vaadin/ui/TreeGrid.java
+++ b/server/src/main/java/com/vaadin/ui/TreeGrid.java
@@ -40,6 +40,7 @@ import com.vaadin.shared.Registration;
 import com.vaadin.shared.ui.treegrid.FocusParentRpc;
 import com.vaadin.shared.ui.treegrid.FocusRpc;
 import com.vaadin.shared.ui.treegrid.NodeCollapseRpc;
+import com.vaadin.shared.ui.treegrid.TreeGridClientRpc;
 import com.vaadin.shared.ui.treegrid.TreeGridState;
 import com.vaadin.ui.declarative.DesignAttributeHandler;
 import com.vaadin.ui.declarative.DesignContext;
@@ -117,6 +118,8 @@ public class TreeGrid<T> extends Grid<T> {
 
         private final T expandedItem;
 
+        private final boolean userOriginated;
+
         /**
          * Construct an expand event.
          *
@@ -124,10 +127,15 @@ public class TreeGrid<T> extends Grid<T> {
          *            the tree grid this event originated from
          * @param expandedItem
          *            the item that was expanded
+         * @param userOriginated
+         *            whether the expand was triggered by a user interaction or
+         *            the server
          */
-        public ExpandEvent(TreeGrid<T> source, T expandedItem) {
+        public ExpandEvent(TreeGrid<T> source, T expandedItem,
+                boolean userOriginated) {
             super(source);
             this.expandedItem = expandedItem;
+            this.userOriginated = userOriginated;
         }
 
         /**
@@ -138,6 +146,17 @@ public class TreeGrid<T> extends Grid<T> {
         public T getExpandedItem() {
             return expandedItem;
         }
+
+        /**
+         * Returns whether this event was triggered by user interaction, on the
+         * client side, or programmatically, on the server side.
+         *
+         * @return {@code true} if this event originates from the client,
+         *         {@code false} otherwise.
+         */
+        public boolean isUserOriginated() {
+            return userOriginated;
+        }
     }
 
     /**
@@ -153,6 +172,8 @@ public class TreeGrid<T> extends Grid<T> {
 
         private final T collapsedItem;
 
+        private final boolean userOriginated;
+
         /**
          * Construct a collapse event.
          *
@@ -160,10 +181,15 @@ public class TreeGrid<T> extends Grid<T> {
          *            the tree grid this event originated from
          * @param collapsedItem
          *            the item that was collapsed
+         * @param userOriginated
+         *            whether the collapse was triggered by a user interaction
+         *            or the server
          */
-        public CollapseEvent(TreeGrid<T> source, T collapsedItem) {
+        public CollapseEvent(TreeGrid<T> source, T collapsedItem,
+                boolean userOriginated) {
             super(source);
             this.collapsedItem = collapsedItem;
+            this.userOriginated = userOriginated;
         }
 
         /**
@@ -174,6 +200,17 @@ public class TreeGrid<T> extends Grid<T> {
         public T getCollapsedItem() {
             return collapsedItem;
         }
+
+        /**
+         * Returns whether this event was triggered by user interaction, on the
+         * client side, or programmatically, on the server side.
+         *
+         * @return {@code true} if this event originates from the client,
+         *         {@code false} otherwise.
+         */
+        public boolean isUserOriginated() {
+            return userOriginated;
+        }
     }
 
     public TreeGrid() {
@@ -182,16 +219,18 @@ public class TreeGrid<T> extends Grid<T> {
         registerRpc(new NodeCollapseRpc() {
             @Override
             public void setNodeCollapsed(String rowKey, int rowIndex,
-                    boolean collapse) {
+                    boolean collapse, boolean userOriginated) {
                 if (collapse) {
-                    if (getDataCommunicator().doCollapse(rowKey, rowIndex)) {
+                    if (getDataCommunicator().doCollapse(rowKey, rowIndex)
+                            && userOriginated) {
                         fireCollapseEvent(getDataCommunicator().getKeyMapper()
-                                .get(rowKey));
+                                .get(rowKey), true);
                     }
                 } else {
-                    if (getDataCommunicator().doExpand(rowKey, rowIndex)) {
+                    if (getDataCommunicator().doExpand(rowKey, rowIndex)
+                            && userOriginated) {
                         fireExpandEvent(getDataCommunicator().getKeyMapper()
-                                .get(rowKey));
+                                .get(rowKey), true);
                     }
                 }
             }
@@ -400,6 +439,37 @@ public class TreeGrid<T> extends Grid<T> {
         getDataCommunicator().setItemCollapseAllowedProvider(provider);
     }
 
+    /**
+     * Expands the given item.
+     * <p>
+     * If the item is currently expanded, does nothing. If the item does not
+     * have any children, does nothing.
+     *
+     * @param item
+     *            the item to expand
+     */
+    public void expand(T item) {
+        getDataCommunicator().setPendingExpand(item).ifPresent(key -> {
+            getRpcProxy(TreeGridClientRpc.class).setExpanded(key);
+            fireExpandEvent(item, false);
+        });
+    }
+
+    /**
+     * Collapses the given item.
+     * <p>
+     * If the item is already collapsed, does nothing.
+     *
+     * @param item
+     *            the item to collapse
+     */
+    public void collapse(T item) {
+        getDataCommunicator().collapseItem(item).ifPresent(key -> {
+            getRpcProxy(TreeGridClientRpc.class).setCollapsed(key);
+            fireCollapseEvent(item, false);
+        });
+    }
+
     @Override
     protected TreeGridState getState() {
         return (TreeGridState) super.getState();
@@ -527,9 +597,12 @@ public class TreeGrid<T> extends Grid<T> {
      *
      * @param item
      *            the item that was expanded
+     * @param userOriginated
+     *            whether the expand was triggered by a user interaction or the
+     *            server
      */
-    private void fireExpandEvent(T item) {
-        fireEvent(new ExpandEvent<>(this, item));
+    private void fireExpandEvent(T item, boolean userOriginated) {
+        fireEvent(new ExpandEvent<>(this, item, userOriginated));
     }
 
     /**
@@ -537,8 +610,11 @@ public class TreeGrid<T> extends Grid<T> {
      *
      * @param item
      *            the item that was collapsed
+     * @param userOriginated
+     *            whether the collapse was triggered by a user interaction or
+     *            the server
      */
-    private void fireCollapseEvent(T item) {
-        fireEvent(new CollapseEvent<>(this, item));
+    private void fireCollapseEvent(T item, boolean userOriginated) {
+        fireEvent(new CollapseEvent<>(this, item, userOriginated));
     }
 }
diff --git a/shared/src/main/java/com/vaadin/shared/ui/treegrid/NodeCollapseRpc.java b/shared/src/main/java/com/vaadin/shared/ui/treegrid/NodeCollapseRpc.java
index 1b16203648..f0748a9b88 100644
--- a/shared/src/main/java/com/vaadin/shared/ui/treegrid/NodeCollapseRpc.java
+++ b/shared/src/main/java/com/vaadin/shared/ui/treegrid/NodeCollapseRpc.java
@@ -36,6 +36,10 @@ public interface NodeCollapseRpc extends ServerRpc {
      *            index where the row is in grid (all rows)
      * @param collapse
      *            {@code true} to collapse, {@code false} to expand
+     * @param userOriginated
+     *            {@code true} if this RPC was triggered by a user interaction,
+     *            {@code false} otherwise
      */
-    void setNodeCollapsed(String rowKey, int rowIndex, boolean collapse);
+    void setNodeCollapsed(String rowKey, int rowIndex, boolean collapse,
+            boolean userOriginated);
 }
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
new file mode 100644
index 0000000000..a07a2d4885
--- /dev/null
+++ b/shared/src/main/java/com/vaadin/shared/ui/treegrid/TreeGridClientRpc.java
@@ -0,0 +1,45 @@
+/*
+ * Copyright 2000-2016 Vaadin Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package com.vaadin.shared.ui.treegrid;
+
+import com.vaadin.shared.communication.ClientRpc;
+
+/**
+ * Server-to-client RPC interface for the TreeGrid component.
+ *
+ * @since 8.1
+ * @author Vaadin Ltd
+ */
+public interface TreeGridClientRpc extends ClientRpc {
+
+    /**
+     * Inform the client that an item with the given key has been expanded by
+     * the server.
+     *
+     * @param key
+     *            the communication key of the expanded item
+     */
+    public void setExpanded(String key);
+
+    /**
+     * Inform the client that an item with the given key has been collapsed by
+     * the server.
+     *
+     * @param key
+     *            the communication key of the collapsed item
+     */
+    public void setCollapsed(String key);
+}
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 ae011355c9..a8ce0a7489 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
@@ -63,6 +63,8 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
         createDataProviderSelect();
         createHierarchyColumnSelect();
         createCollapseAllowedSelect();
+        createExpandMenu();
+        createCollapseMenu();
         createListenerMenu();
     }
 
@@ -146,13 +148,45 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
     }
 
     @SuppressWarnings("unchecked")
+    private void createExpandMenu() {
+        createCategory("Server-side expand", CATEGORY_FEATURES);
+        createClickAction("Expand 0 | 0", "Server-side expand",
+                (treeGrid, value, data) -> treeGrid.expand(value),
+                new HierarchicalTestBean(null, 0, 0));
+        createClickAction("Expand 1 | 1", "Server-side expand",
+                (treeGrid, value, data) -> treeGrid.expand(value),
+                new HierarchicalTestBean("/0/0", 1, 1));
+        createClickAction("Expand 2 | 1", "Server-side expand",
+                (treeGrid, value, data) -> treeGrid.expand(value),
+                new HierarchicalTestBean("/0/0/1/1", 2, 1));
+    }
+
+    @SuppressWarnings("unchecked")
+    private void createCollapseMenu() {
+        createCategory("Server-side collapse", CATEGORY_FEATURES);
+        createClickAction("Collapse 0 | 0", "Server-side collapse",
+                (treeGrid, value, data) -> treeGrid.collapse(value),
+                new HierarchicalTestBean(null, 0, 0));
+        createClickAction("Collapse 1 | 1", "Server-side collapse",
+                (treeGrid, value, data) -> treeGrid.collapse(value),
+                new HierarchicalTestBean("/0/0", 1, 1));
+        createClickAction("Collapse 2 | 1", "Server-side collapse",
+                (treeGrid, value, data) -> treeGrid.collapse(value),
+                new HierarchicalTestBean("/0/0/1/1", 2, 1));
+    }
+
+    @SuppressWarnings("unchecked")
     private void createListenerMenu() {
         createListenerAction("Collapse listener", "State",
                 treeGrid -> treeGrid.addCollapseListener(event -> log(
-                        "Item collapsed: " + event.getCollapsedItem())));
+                        "Item collapsed (user originated: "
+                                + event.isUserOriginated() + "): "
+                                + event.getCollapsedItem())));
         createListenerAction("Expand listener", "State",
                 treeGrid -> treeGrid.addExpandListener(event -> log(
-                        "Item expanded: " + event.getExpandedItem())));
+                        "Item expanded (user originated: "
+                                + event.isUserOriginated() + "): "
+                                + event.getExpandedItem())));
     }
 
     static class HierarchicalTestBean {
diff --git a/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridHugeTree.java b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridHugeTree.java
new file mode 100644
index 0000000000..7cfb926849
--- /dev/null
+++ b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridHugeTree.java
@@ -0,0 +1,57 @@
+package com.vaadin.tests.components.treegrid;
+
+import com.vaadin.annotations.Theme;
+import com.vaadin.annotations.Widgetset;
+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.Button;
+import com.vaadin.ui.TreeGrid;
+
+@Theme("valo")
+@Widgetset("com.vaadin.DefaultWidgetSet")
+public class TreeGridHugeTree
+        extends AbstractTestUI {
+
+    private TreeGrid<String> treeGrid;
+    private InMemoryHierarchicalDataProvider<String> inMemoryDataProvider;
+
+    private void initializeDataProvider() {
+        HierarchyData<String> data = new HierarchyData<>();
+        for (int i = 0; i < 3; i++) {
+            String granddad = "Granddad " + i;
+            data.addItem(null, granddad);
+            for (int j = 0; j < 3; j++) {
+                String dad = "Dad " + i + "/" + j;
+                data.addItem(granddad, dad);
+                for (int k = 0; k < 300; k++) {
+                    String son = "Son " + i + "/" + j + "/" + k;
+                    data.addItem(dad, son);
+                }
+            }
+        }
+        inMemoryDataProvider = new InMemoryHierarchicalDataProvider<>(data);
+    }
+
+    @Override
+    protected void setup(VaadinRequest request) {
+        initializeDataProvider();
+        treeGrid = new TreeGrid<>();
+        treeGrid.setDataProvider(inMemoryDataProvider);
+        treeGrid.setSizeFull();
+        treeGrid.addColumn(String::toString).setCaption("String")
+                .setId("string");
+        treeGrid.addColumn((i) -> "--").setCaption("Nothing");
+        treeGrid.setHierarchyColumn("string");
+        treeGrid.setId("testComponent");
+
+        Button expand = new Button("Expand Granddad 1");
+        expand.addClickListener(event -> treeGrid.expand("Granddad 1"));
+        Button collapse = new Button("Collapse Granddad 1");
+        collapse.addClickListener(event -> treeGrid.collapse("Granddad 1"));
+
+        addComponents(treeGrid, expand, collapse);
+    }
+
+}
\ No newline at end of file
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 7fb7ddf3f3..dd4546c6c0 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,12 +1,13 @@
 package com.vaadin.tests.components.treegrid;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.List;
 
-import com.vaadin.testbench.parallel.Browser;
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized.Parameters;
@@ -18,10 +19,6 @@ import com.vaadin.testbench.elements.TreeGridElement;
 import com.vaadin.tests.tb3.MultiBrowserTest;
 import com.vaadin.tests.tb3.ParameterizedTB3Runner;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-
 @RunWith(ParameterizedTB3Runner.class)
 public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
 
@@ -41,28 +38,82 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
     @Before
     public void before() {
         setDebug(true);
-        openTestURL("theme=valo");
+        openTestURL();
         grid = $(TreeGridElement.class).first();
     }
 
     @Test
-    @Ignore // currently no implementation exists for toggling from the server
-            // side
     public void toggle_collapse_server_side() {
         assertEquals(3, grid.getRowCount());
         assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
 
-        selectMenuPath("Component", "Features", "Toggle expand", "0 | 0");
+        selectMenuPath("Component", "Features", "Server-side expand",
+                "Expand 0 | 0");
+        assertEquals(6, grid.getRowCount());
+        assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "1 | 2" });
+
+        // expanding already expanded item should have no effect
+        selectMenuPath("Component", "Features", "Server-side expand",
+                "Expand 0 | 0");
         assertEquals(6, grid.getRowCount());
         assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "1 | 2" });
 
-        selectMenuPath("Component", "Features", "Toggle expand", "0 | 0");
+        selectMenuPath("Component", "Features", "Server-side collapse",
+                "Collapse 0 | 0");
+        assertEquals(3, grid.getRowCount());
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
+
+        // collapsing the same item twice should have no effect
+        selectMenuPath("Component", "Features", "Server-side collapse",
+                "Collapse 0 | 0");
+        assertEquals(3, grid.getRowCount());
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
+
+        selectMenuPath("Component", "Features", "Server-side expand",
+                "Expand 1 | 1");
+        // 1 | 1 not yet visible, shouldn't immediately expand anything
         assertEquals(3, grid.getRowCount());
         assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
 
+        selectMenuPath("Component", "Features", "Server-side expand",
+                "Expand 0 | 0");
+        // 1 | 1 becomes visible and is also expanded
+        assertEquals(9, grid.getRowCount());
+        assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "2 | 0", "2 | 1",
+                "2 | 2", "1 | 2" });
+
         // collapsing a leaf should have no effect
-        selectMenuPath("Component", "Features", "Toggle expand", "1 | 0");
+        selectMenuPath("Component", "Features", "Server-side collapse",
+                "Collapse 2 | 1");
+        assertEquals(9, grid.getRowCount());
+        assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "2 | 0", "2 | 1",
+                "2 | 2", "1 | 2" });
+
+        // collapsing 0 | 0 should collapse the expanded 1 | 1
+        selectMenuPath("Component", "Features", "Server-side collapse",
+                "Collapse 0 | 0");
         assertEquals(3, grid.getRowCount());
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
+
+        // 1 | 1 should not be expanded this time
+        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();
+    }
+
+    @Test
+    public void pending_expands_cleared_when_data_provider_set() {
+        selectMenuPath("Component", "Features", "Server-side expand",
+                "Expand 1 | 1");
+        selectMenuPath("Component", "Features", "Set data provider",
+                "LazyHierarchicalDataProvider");
+        grid.expandWithClick(0);
+        assertEquals(6, grid.getRowCount());
+        assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "1 | 2" });
     }
 
     @Test
@@ -105,20 +156,21 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
         // Should expand "1 | 1" without moving focus
         new Actions(getDriver()).sendKeys(Keys.RIGHT).perform();
         assertEquals(9, grid.getRowCount());
-        assertCellTexts(2, 0, new String[] { "1 | 1", "2 | 0",  "2 | 1",  "2 | 2",  "1 | 2"});
+        assertCellTexts(2, 0,
+                new String[] { "1 | 1", "2 | 0", "2 | 1", "2 | 2", "1 | 2" });
         assertTrue(grid.getRow(2).hasClassName("v-grid-rowmode-row-focused"));
 
         // Should collapse "1 | 1"
         new Actions(getDriver()).sendKeys(Keys.LEFT).perform();
         assertEquals(6, grid.getRowCount());
-        assertCellTexts(2, 0, new String[] { "1 | 1", "1 | 2",  "0 | 1"});
+        assertCellTexts(2, 0, new String[] { "1 | 1", "1 | 2", "0 | 1" });
         assertTrue(grid.getRow(2).hasClassName("v-grid-rowmode-row-focused"));
 
-
         // Should navigate to "0 | 0"
         new Actions(getDriver()).sendKeys(Keys.LEFT).perform();
         assertEquals(6, grid.getRowCount());
-        assertCellTexts(0, 0, new String[] { "0 | 0", "1 | 0", "1 | 1", "1 | 2" , "0 | 1" });
+        assertCellTexts(0, 0,
+                new String[] { "0 | 0", "1 | 0", "1 | 1", "1 | 2", "0 | 1" });
         assertTrue(grid.getRow(0).hasClassName("v-grid-rowmode-row-focused"));
 
         // Should collapse "0 | 0"
@@ -171,18 +223,40 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
         selectMenuPath("Component", "State", "Expand listener");
         selectMenuPath("Component", "State", "Collapse listener");
 
-        assertFalse(logContainsText("Item expanded: 0 | 0"));
-        assertFalse(logContainsText("Item collapsed: 0 | 0"));
+        assertFalse(logContainsText(
+                "Item expanded (user originated: true): 0 | 0"));
+        assertFalse(logContainsText(
+                "Item collapsed (user originated: true): 0 | 0"));
 
         grid.expandWithClick(0);
 
-        assertTrue(logContainsText("Item expanded: 0 | 0"));
-        assertFalse(logContainsText("Item collapsed: 0 | 0"));
+        assertTrue(logContainsText(
+                "Item expanded (user originated: true): 0 | 0"));
+        assertFalse(logContainsText(
+                "Item collapsed (user originated: true): 0 | 0"));
 
         grid.collapseWithClick(0);
 
-        assertTrue(logContainsText("Item expanded: 0 | 0"));
-        assertTrue(logContainsText("Item collapsed: 0 | 0"));
+        assertTrue(logContainsText(
+                "Item expanded (user originated: true): 0 | 0"));
+        assertTrue(logContainsText(
+                "Item collapsed (user originated: true): 0 | 0"));
+
+        selectMenuPath("Component", "Features", "Server-side expand",
+                "Expand 0 | 0");
+
+        assertTrue(logContainsText(
+                "Item expanded (user originated: false): 0 | 0"));
+        assertFalse(logContainsText(
+                "Item collapsed (user originated: false): 0 | 0"));
+
+        selectMenuPath("Component", "Features", "Server-side collapse",
+                "Collapse 0 | 0");
+
+        assertTrue(logContainsText(
+                "Item expanded (user originated: false): 0 | 0"));
+        assertTrue(logContainsText(
+                "Item collapsed (user originated: false): 0 | 0"));
 
         selectMenuPath("Component", "State", "Expand listener");
         selectMenuPath("Component", "State", "Collapse listener");
@@ -190,8 +264,10 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
         grid.expandWithClick(1);
         grid.collapseWithClick(1);
 
-        assertFalse(logContainsText("Item expanded: 0 | 1"));
-        assertFalse(logContainsText("Item collapsed: 0 | 1"));
+        assertFalse(logContainsText(
+                "Item expanded (user originated: true): 0 | 1"));
+        assertFalse(logContainsText(
+                "Item collapsed (user originated: true): 0 | 1"));
     }
 
     private void assertCellTexts(int startRowIndex, int cellIndex,
diff --git a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridHugeTreeTest.java b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridHugeTreeTest.java
new file mode 100644
index 0000000000..1d5d9b309d
--- /dev/null
+++ b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridHugeTreeTest.java
@@ -0,0 +1,60 @@
+package com.vaadin.tests.components.treegrid;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import com.vaadin.testbench.elements.ButtonElement;
+import com.vaadin.testbench.elements.TreeGridElement;
+import com.vaadin.tests.tb3.SingleBrowserTest;
+
+public class TreeGridHugeTreeTest extends SingleBrowserTest {
+
+    private TreeGridElement grid;
+
+    @Test
+    public void toggle_expand_when_row_out_of_cache() {
+        openTestURL();
+
+        grid = $(TreeGridElement.class).first();
+        ButtonElement expandSecondRowButton = $(ButtonElement.class).get(0);
+        ButtonElement collapseSecondRowButton = $(ButtonElement.class).get(1);
+
+        grid.expandWithClick(2);
+        grid.expandWithClick(3);
+        grid.scrollToRow(300);
+        grid.waitForVaadin();
+
+        expandSecondRowButton.click();
+
+        grid.scrollToRow(0);
+        assertCellTexts(0, 0, new String[] { "Granddad 0", "Granddad 1",
+                "Dad 1/0", "Dad 1/1", "Dad 1/2", "Granddad 2", "Dad 2/0" });
+
+        grid.scrollToRow(300);
+        grid.waitForVaadin();
+        collapseSecondRowButton.click();
+        grid.scrollToRow(0);
+        grid.waitForVaadin();
+        assertCellTexts(0, 0, new String[] { "Granddad 0", "Granddad 1",
+                "Granddad 2", "Dad 2/0" });
+
+        grid.scrollToRow(300);
+        grid.waitForVaadin();
+        expandSecondRowButton.click();
+        collapseSecondRowButton.click();
+        grid.scrollToRow(0);
+        grid.waitForVaadin();
+        assertCellTexts(0, 0, new String[] { "Granddad 0", "Granddad 1",
+                "Granddad 2", "Dad 2/0" });
+    }
+
+    private void assertCellTexts(int startRowIndex, int cellIndex,
+            String[] cellTexts) {
+        int index = startRowIndex;
+        for (String cellText : cellTexts) {
+            Assert.assertEquals(cellText,
+                    grid.getRow(index).getCell(cellIndex).getText());
+            index++;
+        }
+    }
+}