Add an initial version of Tree component (#9212)

This patch introduces a Tree component, implemented as a Composite
using TreeGrid. As initial version of the component, some functionality
of the old Tree component has been left out for now.

Partly addresses #8617

9 files changed, 692 insertions, 179 deletions

diff --git a/server/src/main/java/com/vaadin/event/CollapseEvent.java b/server/src/main/java/com/vaadin/event/CollapseEvent.java
new file mode 100644
index 0000000000..ba740142e3
--- /dev/null
+++ b/server/src/main/java/com/vaadin/event/CollapseEvent.java
@@ -0,0 +1,97 @@
+/*
+ * 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.event;
+
+import java.io.Serializable;
+import java.lang.reflect.Method;
+
+import com.vaadin.ui.Component;
+import com.vaadin.util.ReflectTools;
+
+/**
+ * An event that is fired when an item is collapsed. Note that expanded subtrees
+ * of the collapsed item will not trigger collapse events.
+ *
+ * @author Vaadin Ltd
+ * @since 8.1
+ * @param <T>
+ *            collapsed item type
+ */
+public class CollapseEvent<T> extends Component.Event {
+
+    private final T collapsedItem;
+
+    private final boolean userOriginated;
+
+    /**
+     * Construct a collapse event.
+     *
+     * @param source
+     *            the hierarchical 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(Component source, T collapsedItem,
+            boolean userOriginated) {
+        super(source);
+        this.collapsedItem = collapsedItem;
+        this.userOriginated = userOriginated;
+    }
+
+    /**
+     * Get the collapsed item that triggered this event.
+     *
+     * @return the collapsed item
+     */
+    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;
+    }
+
+    /**
+     * Item collapse event listener.
+     *
+     * @param <T>
+     *            the collapsed item's type
+     */
+    @FunctionalInterface
+    public interface CollapseListener<T> extends Serializable {
+
+        public static final Method COLLAPSE_METHOD = ReflectTools.findMethod(
+                CollapseListener.class, "itemCollapse", CollapseEvent.class);
+
+        /**
+         * Callback method for when an item has been collapsed.
+         *
+         * @param event
+         *            the collapse event
+         */
+        public void itemCollapse(CollapseEvent<T> event);
+    }
+}
diff --git a/server/src/main/java/com/vaadin/event/ExpandEvent.java b/server/src/main/java/com/vaadin/event/ExpandEvent.java
new file mode 100644
index 0000000000..775c368a12
--- /dev/null
+++ b/server/src/main/java/com/vaadin/event/ExpandEvent.java
@@ -0,0 +1,96 @@
+/*
+ * 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.event;
+
+import java.io.Serializable;
+import java.lang.reflect.Method;
+
+import com.vaadin.ui.Component;
+import com.vaadin.util.ReflectTools;
+
+/**
+ * An event that is fired when an item is expanded.
+ *
+ * @author Vaadin Ltd
+ * @since 8.1
+ * @param <T>
+ *            the expanded item's type
+ */
+public class ExpandEvent<T> extends Component.Event {
+
+    private final T expandedItem;
+
+    private final boolean userOriginated;
+
+    /**
+     * Construct an expand event.
+     *
+     * @param source
+     *            the hierarchical component 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(Component source, T expandedItem,
+            boolean userOriginated) {
+        super(source);
+        this.expandedItem = expandedItem;
+        this.userOriginated = userOriginated;
+    }
+
+    /**
+     * Get the expanded item that triggered this event.
+     *
+     * @return the expanded item
+     */
+    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;
+    }
+
+    /**
+     * Item expand event listener.
+     *
+     * @param <T>
+     *            the expanded item's type
+     */
+    @FunctionalInterface
+    public interface ExpandListener<T> extends Serializable {
+
+        public static final Method EXPAND_METHOD = ReflectTools.findMethod(
+                ExpandListener.class, "itemExpand", ExpandEvent.class);
+
+        /**
+         * Callback method for when an item has been expanded.
+         *
+         * @param event
+         *            the expand event
+         */
+        public void itemExpand(ExpandEvent<T> event);
+    }
+}
diff --git a/server/src/main/java/com/vaadin/ui/Tree.java b/server/src/main/java/com/vaadin/ui/Tree.java
new file mode 100644
index 0000000000..8b34fdb06e
--- /dev/null
+++ b/server/src/main/java/com/vaadin/ui/Tree.java
@@ -0,0 +1,364 @@
+/*
+ * 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.ui;
+
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+import com.vaadin.data.Binder;
+import com.vaadin.data.HasDataProvider;
+import com.vaadin.data.SelectionModel;
+import com.vaadin.data.provider.DataGenerator;
+import com.vaadin.data.provider.DataProvider;
+import com.vaadin.data.provider.HierarchicalDataProvider;
+import com.vaadin.event.CollapseEvent;
+import com.vaadin.event.CollapseEvent.CollapseListener;
+import com.vaadin.event.ExpandEvent;
+import com.vaadin.event.ExpandEvent.ExpandListener;
+import com.vaadin.event.selection.SelectionListener;
+import com.vaadin.server.Resource;
+import com.vaadin.shared.Registration;
+import com.vaadin.shared.ui.grid.HeightMode;
+import com.vaadin.shared.ui.tree.TreeRendererState;
+import com.vaadin.ui.Grid.SelectionMode;
+import com.vaadin.ui.renderers.AbstractRenderer;
+
+import elemental.json.JsonObject;
+
+/**
+ * Tree component. A Tree can be used to select an item from a hierarchical set
+ * of items.
+ *
+ * @author Vaadin Ltd
+ * @since 8.1
+ *
+ * @param <T>
+ *            the data type
+ */
+public class Tree<T> extends Composite implements HasDataProvider<T> {
+
+    /**
+     * String renderer that handles icon resources and stores their identifiers
+     * into data objects.
+     */
+    public final class TreeRenderer extends AbstractRenderer<T, String>
+            implements DataGenerator<T> {
+
+        /**
+         * Constructs a new TreeRenderer.
+         */
+        protected TreeRenderer() {
+            super(String.class);
+        }
+
+        private Map<T, String> resourceKeyMap = new HashMap<>();
+        private int counter = 0;
+
+        @Override
+        public void generateData(T item, JsonObject jsonObject) {
+            Resource resource = iconProvider.apply(item);
+            if (resource == null) {
+                destroyData(item);
+                return;
+            }
+
+            if (!resourceKeyMap.containsKey(item)) {
+                resourceKeyMap.put(item, "icon" + (counter++));
+            }
+            setResource(resourceKeyMap.get(item), resource);
+            jsonObject.put("itemIcon", resourceKeyMap.get(item));
+        }
+
+        @Override
+        public void destroyData(T item) {
+            if (resourceKeyMap.containsKey(item)) {
+                setResource(resourceKeyMap.get(item), null);
+                resourceKeyMap.remove(item);
+            }
+        }
+
+        @Override
+        public void destroyAllData() {
+            Set<T> keys = new HashSet<>(resourceKeyMap.keySet());
+            for (T key : keys) {
+                destroyData(key);
+            }
+        }
+
+        @Override
+        protected TreeRendererState getState() {
+            return (TreeRendererState) super.getState();
+        }
+
+        @Override
+        protected TreeRendererState getState(boolean markAsDirty) {
+            return (TreeRendererState) super.getState(markAsDirty);
+        }
+    }
+
+    private TreeGrid<T> treeGrid = new TreeGrid<>();
+    private ItemCaptionGenerator<T> captionGenerator = String::valueOf;
+    private IconGenerator<T> iconProvider = t -> null;
+
+    /**
+     * Constructs a new Tree Component.
+     */
+    public Tree() {
+        setCompositionRoot(treeGrid);
+        TreeRenderer renderer = new TreeRenderer();
+        treeGrid.getDataCommunicator().addDataGenerator(renderer);
+        treeGrid.addColumn(i -> captionGenerator.apply(i), renderer)
+                .setId("column");
+        treeGrid.setHierarchyColumn("column");
+        while (treeGrid.getHeaderRowCount() > 0) {
+            treeGrid.removeHeaderRow(0);
+        }
+        treeGrid.setPrimaryStyleName("v-newtree");
+
+        setWidth("100%");
+        setHeightUndefined();
+        treeGrid.setHeightMode(HeightMode.UNDEFINED);
+
+        treeGrid.addExpandListener(e -> fireExpandEvent(e.getExpandedItem(),
+                e.isUserOriginated()));
+        treeGrid.addCollapseListener(e -> fireCollapseEvent(
+                e.getCollapsedItem(), e.isUserOriginated()));
+    }
+
+    /**
+     * Constructs a new Tree Component.
+     *
+     * @param caption
+     *            the caption for component
+     */
+    public Tree(String caption) {
+        this();
+
+        setCaption(caption);
+    }
+
+    @Override
+    public HierarchicalDataProvider<T, ?> getDataProvider() {
+        return treeGrid.getDataProvider();
+    }
+
+    @Override
+    public void setDataProvider(DataProvider<T, ?> dataProvider) {
+        treeGrid.setDataProvider(dataProvider);
+    }
+
+    /**
+     * Adds an ExpandListener to this Tree.
+     *
+     * @see ExpandEvent
+     *
+     * @param listener
+     *            the listener to add
+     * @return a registration for the listener
+     */
+    public Registration addExpandListener(ExpandListener<T> listener) {
+        return addListener(ExpandEvent.class, listener,
+                ExpandListener.EXPAND_METHOD);
+    }
+
+    /**
+     * Adds a CollapseListener to this Tree.
+     *
+     * @see CollapseEvent
+     *
+     * @param listener
+     *            the listener to add
+     * @return a registration for the listener
+     */
+    public Registration addCollapseListener(CollapseListener<T> listener) {
+        return addListener(CollapseEvent.class, listener,
+                CollapseListener.COLLAPSE_METHOD);
+    }
+
+    /**
+     * Fires an expand event with given item.
+     *
+     * @param item
+     *            the expanded item
+     * @param userOriginated
+     *            whether the expand was triggered by a user interaction or the
+     *            server
+     */
+    protected void fireExpandEvent(T item, boolean userOriginated) {
+        fireEvent(new ExpandEvent<>(this, item, userOriginated));
+    }
+
+    /**
+     * Fires a collapse event with given item.
+     *
+     * @param item
+     *            the collapsed item
+     * @param userOriginated
+     *            whether the collapse was triggered by a user interaction or
+     *            the server
+     */
+    protected void fireCollapseEvent(T item, boolean userOriginated) {
+        fireEvent(new CollapseEvent<>(this, item, userOriginated));
+    }
+
+    /**
+     * 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) {
+        treeGrid.expand(item);
+    }
+
+    /**
+     * Collapses the given item.
+     * <p>
+     * If the item is already collapsed, does nothing.
+     *
+     * @param item
+     *            the item to collapse
+     */
+    public void collapse(T item) {
+        treeGrid.collapse(item);
+    }
+
+    /**
+     * This method is a shorthand that delegates to the currently set selection
+     * model.
+     *
+     * @see #getSelectionModel()
+     *
+     * @return set of selected items
+     */
+    public Set<T> getSelectedItems() {
+        return treeGrid.getSelectedItems();
+    }
+
+    /**
+     * This method is a shorthand that delegates to the currently set selection
+     * model.
+     *
+     * @param item
+     *            item to select
+     *
+     * @see SelectionModel#select(Object)
+     * @see #getSelectionModel()
+     */
+    public void select(T item) {
+        treeGrid.select(item);
+    }
+
+    /**
+     * This method is a shorthand that delegates to the currently set selection
+     * model.
+     *
+     * @param item
+     *            item to deselect
+     *
+     * @see SelectionModel#deselect(Object)
+     * @see #getSelectionModel()
+     */
+    public void deselect(T item) {
+        treeGrid.deselect(item);
+    }
+
+    /**
+     * Adds a selection listener to the current selection model.
+     * <p>
+     * <strong>NOTE:</strong> If selection mode is switched with
+     * {@link setSelectionMode(SelectionMode)}, then this listener is not
+     * triggered anymore when selection changes!
+     *
+     * @param listener
+     *            the listener to add
+     * @return a registration handle to remove the listener
+     *
+     * @throws UnsupportedOperationException
+     *             if selection has been disabled with
+     *             {@link SelectionMode.NONE}
+     */
+    public Registration addSelectionListener(SelectionListener<T> listener) {
+        return treeGrid.addSelectionListener(listener);
+    }
+
+    /**
+     * Use this tree as a single select in {@link Binder}. Throws
+     * {@link IllegalStateException} if the tree is not using
+     * {@link SelectionMode#SINGLE}.
+     *
+     * @return the single select wrapper that can be used in binder
+     */
+    public SingleSelect<T> asSingleSelect() {
+        return treeGrid.asSingleSelect();
+    }
+
+    /**
+     * Returns the selection model for this Tree.
+     *
+     * @return the selection model, not <code>null</code>
+     */
+    public SelectionModel<T> getSelectionModel() {
+        return treeGrid.getSelectionModel();
+    }
+
+    @Override
+    public void setItems(Collection<T> items) {
+        treeGrid.setItems(items);
+    }
+
+    /**
+     * Sets the item caption generator that is used to produce the strings shown
+     * as the text for each item. By default, {@link String#valueOf(Object)} is
+     * used.
+     *
+     * @param captionGenerator
+     *            the item caption provider to use, not <code>null</code>
+     */
+    public void setItemCaptionGenerator(
+            ItemCaptionGenerator<T> captionGenerator) {
+        Objects.requireNonNull(captionGenerator,
+                "Caption generator must not be null");
+        this.captionGenerator = captionGenerator;
+        treeGrid.getDataCommunicator().reset();
+    }
+
+    /**
+     * Sets the item icon generator that is used to produce custom icons for
+     * items. The generator can return <code>null</code> for items with no icon.
+     *
+     * @see IconGenerator
+     *
+     * @param iconGenerator
+     *            the item icon generator to set, not <code>null</code>
+     * @throws NullPointerException
+     *             if {@code itemIconGenerator} is {@code null}
+     */
+    public void setItemIconGenerator(IconGenerator<T> iconGenerator) {
+        Objects.requireNonNull(iconGenerator,
+                "Item icon generator must not be null");
+        this.iconProvider = iconGenerator;
+        treeGrid.getDataCommunicator().reset();
+    }
+
+}
diff --git a/server/src/main/java/com/vaadin/ui/TreeGrid.java b/server/src/main/java/com/vaadin/ui/TreeGrid.java
index 4abad5e478..720bffaf3e 100644
--- a/server/src/main/java/com/vaadin/ui/TreeGrid.java
+++ b/server/src/main/java/com/vaadin/ui/TreeGrid.java
@@ -15,8 +15,6 @@
  */
 package com.vaadin.ui;
 
-import java.io.Serializable;
-import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -36,6 +34,10 @@ import com.vaadin.data.provider.HierarchicalDataCommunicator;
 import com.vaadin.data.provider.HierarchicalDataProvider;
 import com.vaadin.data.provider.HierarchicalQuery;
 import com.vaadin.data.provider.InMemoryHierarchicalDataProvider;
+import com.vaadin.event.CollapseEvent;
+import com.vaadin.event.CollapseEvent.CollapseListener;
+import com.vaadin.event.ExpandEvent;
+import com.vaadin.event.ExpandEvent.ExpandListener;
 import com.vaadin.server.SerializablePredicate;
 import com.vaadin.shared.Registration;
 import com.vaadin.shared.ui.treegrid.FocusParentRpc;
@@ -48,7 +50,6 @@ import com.vaadin.ui.declarative.DesignContext;
 import com.vaadin.ui.declarative.DesignFormatter;
 import com.vaadin.ui.renderers.AbstractRenderer;
 import com.vaadin.ui.renderers.Renderer;
-import com.vaadin.util.ReflectTools;
 
 /**
  * A grid component for displaying hierarchical tabular data.
@@ -61,159 +62,6 @@ import com.vaadin.util.ReflectTools;
  */
 public class TreeGrid<T> extends Grid<T> {
 
-    /**
-     * Item expand event listener.
-     *
-     * @author Vaadin Ltd
-     * @since 8.1
-     * @param <T>
-     *            the expanded item's type
-     */
-    @FunctionalInterface
-    public interface ExpandListener<T> extends Serializable {
-
-        public static final Method EXPAND_METHOD = ReflectTools.findMethod(
-                ExpandListener.class, "itemExpand", ExpandEvent.class);
-
-        /**
-         * Callback method for when an item has been expanded.
-         *
-         * @param event
-         *            the expand event
-         */
-        public void itemExpand(ExpandEvent<T> event);
-    }
-
-    /**
-     * Item collapse event listener.
-     *
-     * @author Vaadin Ltd
-     * @since 8.1
-     * @param <T>
-     *            the collapsed item's type
-     */
-    @FunctionalInterface
-    public interface CollapseListener<T> extends Serializable {
-
-        public static final Method COLLAPSE_METHOD = ReflectTools.findMethod(
-                CollapseListener.class, "itemCollapse", CollapseEvent.class);
-
-        /**
-         * Callback method for when an item has been collapsed.
-         *
-         * @param event
-         *            the collapse event
-         */
-        public void itemCollapse(CollapseEvent<T> event);
-    }
-
-    /**
-     * An event that is fired when an item is expanded.
-     *
-     * @author Vaadin Ltd
-     * @since 8.1
-     * @param <T>
-     *            the expanded item's type
-     */
-    public static class ExpandEvent<T> extends Component.Event {
-
-        private final T expandedItem;
-
-        private final boolean userOriginated;
-
-        /**
-         * Construct an expand event.
-         *
-         * @param source
-         *            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,
-                boolean userOriginated) {
-            super(source);
-            this.expandedItem = expandedItem;
-            this.userOriginated = userOriginated;
-        }
-
-        /**
-         * Get the expanded item that triggered this event.
-         *
-         * @return the expanded item
-         */
-        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;
-        }
-    }
-
-    /**
-     * An event that is fired when an item is collapsed. Note that expanded
-     * subtrees of the collapsed item will not trigger collapse events.
-     *
-     * @author Vaadin Ltd
-     * @since 8.1
-     * @param <T>
-     *            collapsed item type
-     */
-    public static class CollapseEvent<T> extends Component.Event {
-
-        private final T collapsedItem;
-
-        private final boolean userOriginated;
-
-        /**
-         * Construct a collapse event.
-         *
-         * @param source
-         *            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,
-                boolean userOriginated) {
-            super(source);
-            this.collapsedItem = collapsedItem;
-            this.userOriginated = userOriginated;
-        }
-
-        /**
-         * Get the collapsed item that triggered this event.
-         *
-         * @return the collapsed item
-         */
-        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() {
         super(new HierarchicalDataCommunicator<>());
 
@@ -239,9 +87,11 @@ public class TreeGrid<T> extends Grid<T> {
         registerRpc(new FocusParentRpc() {
             @Override
             public void focusParent(int rowIndex, int cellIndex) {
-                Integer parentIndex = getDataCommunicator().getParentIndex(rowIndex);
+                Integer parentIndex = getDataCommunicator()
+                        .getParentIndex(rowIndex);
                 if (parentIndex != null) {
-                    getRpcProxy(FocusRpc.class).focusCell(parentIndex, cellIndex);
+                    getRpcProxy(FocusRpc.class).focusCell(parentIndex,
+                            cellIndex);
                 }
             }
         });
diff --git a/server/src/test/java/com/vaadin/server/DedicatedStateTest.java b/server/src/test/java/com/vaadin/server/DedicatedStateTest.java
index 6d8fa485c6..3a21d136f7 100644
--- a/server/src/test/java/com/vaadin/server/DedicatedStateTest.java
+++ b/server/src/test/java/com/vaadin/server/DedicatedStateTest.java
@@ -24,6 +24,7 @@ import org.junit.Test;
 
 import com.vaadin.navigator.Navigator;
 import com.vaadin.tests.VaadinClasses;
+import com.vaadin.ui.Composite;
 import com.vaadin.ui.LegacyWindow;
 import com.vaadin.ui.components.colorpicker.ColorPickerHistory;
 import com.vaadin.ui.components.colorpicker.ColorPickerPopup;
@@ -47,7 +48,8 @@ public class DedicatedStateTest {
     }
 
     private void checkState(Class<?> clazz) {
-        if (WHITE_LIST.contains(clazz.getCanonicalName())) {
+        if (WHITE_LIST.contains(clazz.getCanonicalName())
+                || Composite.class.isAssignableFrom(clazz)) {
             return;
         }
         Method getStateNoArg = getStateNoArg(clazz);
diff --git a/server/src/test/java/com/vaadin/tests/components/TreeTest.java b/server/src/test/java/com/vaadin/tests/components/TreeTest.java
new file mode 100644
index 0000000000..13ef0acc49
--- /dev/null
+++ b/server/src/test/java/com/vaadin/tests/components/TreeTest.java
@@ -0,0 +1,75 @@
+package com.vaadin.tests.components;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import com.vaadin.data.HierarchyData;
+import com.vaadin.data.provider.InMemoryHierarchicalDataProvider;
+import com.vaadin.event.CollapseEvent;
+import com.vaadin.event.CollapseEvent.CollapseListener;
+import com.vaadin.event.ExpandEvent;
+import com.vaadin.event.ExpandEvent.ExpandListener;
+import com.vaadin.ui.Tree;
+
+public class TreeTest {
+
+    private static class TreeCollapseExpandListener
+            implements ExpandListener<String>, CollapseListener<String> {
+
+        private boolean collapsed = false;
+        private boolean expanded = false;
+        private final Tree<String> tree;
+
+        public TreeCollapseExpandListener(Tree<String> tree) {
+            this.tree = tree;
+        }
+
+        @Override
+        public void itemCollapse(CollapseEvent<String> event) {
+            Assert.assertEquals("Source component was incorrect", tree,
+                    event.getComponent());
+            Assert.assertFalse("Multiple collapse events", collapsed);
+            collapsed = true;
+        }
+
+        @Override
+        public void itemExpand(ExpandEvent<String> event) {
+            Assert.assertEquals("Source component was incorrect", tree,
+                    event.getComponent());
+            Assert.assertFalse("Multiple expand events", expanded);
+            expanded = true;
+        }
+
+        public boolean isExpanded() {
+            return expanded;
+        }
+
+        public boolean isCollapsed() {
+            return collapsed;
+        }
+    }
+
+    @Test
+    public void event_source_is_tree() {
+        Tree<String> tree = new Tree<>();
+        HierarchyData<String> hierarchyData = new HierarchyData<>();
+        hierarchyData.addItem(null, "Foo");
+        hierarchyData.addItem("Foo", "Bar");
+        hierarchyData.addItem("Foo", "Baz");
+        tree.setDataProvider(
+                new InMemoryHierarchicalDataProvider<>(hierarchyData));
+
+        TreeCollapseExpandListener listener = new TreeCollapseExpandListener(
+                tree);
+        tree.addExpandListener(listener);
+        tree.addCollapseListener(listener);
+
+        Assert.assertFalse(listener.isExpanded());
+        tree.expand("Foo");
+        Assert.assertTrue("Expand event not fired", listener.isExpanded());
+        Assert.assertFalse(listener.isCollapsed());
+        tree.collapse("Foo");
+        Assert.assertTrue("Collapse event not fired", listener.isCollapsed());
+    }
+
+}
diff --git a/server/src/test/java/com/vaadin/tests/components/treegrid/TreeGridColumnTest.java b/server/src/test/java/com/vaadin/tests/components/treegrid/TreeGridColumnTest.java
deleted file mode 100644
index a928a48a08..0000000000
--- a/server/src/test/java/com/vaadin/tests/components/treegrid/TreeGridColumnTest.java
+++ /dev/null
@@ -1,19 +0,0 @@
-package com.vaadin.tests.components.treegrid;
-
-import org.junit.Test;
-
-import com.vaadin.ui.TreeGrid;
-import com.vaadin.ui.renderers.TextRenderer;
-
-public class TreeGridColumnTest {
-
-    private TreeGrid<String> treeGrid = new TreeGrid<>();
-
-    @Test(expected = IllegalStateException.class)
-    public void testChangeRendererOfHierarchyColumn() {
-        treeGrid.addColumn(Object::toString).setId("foo");
-        treeGrid.setHierarchyColumn("foo");
-        // This should not be allowed.
-        treeGrid.getColumn("foo").setRenderer(new TextRenderer());
-    }
-}
diff --git a/server/src/test/java/com/vaadin/tests/components/treegrid/TreeGridTest.java b/server/src/test/java/com/vaadin/tests/components/treegrid/TreeGridTest.java
new file mode 100644
index 0000000000..2e335f865e
--- /dev/null
+++ b/server/src/test/java/com/vaadin/tests/components/treegrid/TreeGridTest.java
@@ -0,0 +1,47 @@
+package com.vaadin.tests.components.treegrid;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import com.vaadin.data.HierarchyData;
+import com.vaadin.data.provider.InMemoryHierarchicalDataProvider;
+import com.vaadin.ui.TreeGrid;
+import com.vaadin.ui.renderers.TextRenderer;
+
+public class TreeGridTest {
+
+    private TreeGrid<String> treeGrid = new TreeGrid<>();
+    private boolean expandEventFired = false;
+    private boolean collapseEventFired = false;
+
+    @Test(expected = IllegalStateException.class)
+    public void testChangeRendererOfHierarchyColumn() {
+        treeGrid.addColumn(Object::toString).setId("foo");
+        treeGrid.setHierarchyColumn("foo");
+        // This should not be allowed.
+        treeGrid.getColumn("foo").setRenderer(new TextRenderer());
+    }
+
+    @Test
+    public void testExpandAndCollapseEvents() {
+        HierarchyData<String> hierarchyData = new HierarchyData<>();
+        hierarchyData.addItem(null, "Foo");
+        hierarchyData.addItem("Foo", "Bar");
+        hierarchyData.addItem("Foo", "Baz");
+        treeGrid.setDataProvider(
+                new InMemoryHierarchicalDataProvider<>(hierarchyData));
+
+        treeGrid.addExpandListener(e -> expandEventFired = true);
+        treeGrid.addCollapseListener(e -> collapseEventFired = true);
+
+        // Test expand event
+        Assert.assertFalse(expandEventFired);
+        treeGrid.expand("Foo");
+        Assert.assertTrue("Expand event not fired", expandEventFired);
+
+        // Test collapse event
+        Assert.assertFalse(collapseEventFired);
+        treeGrid.collapse("Foo");
+        Assert.assertTrue("Collapse event not fired", collapseEventFired);
+    }
+}
diff --git a/server/src/test/java/com/vaadin/tests/server/component/StateGetDoesNotMarkDirtyTest.java b/server/src/test/java/com/vaadin/tests/server/component/StateGetDoesNotMarkDirtyTest.java
index 29dd6638fa..99c686743c 100644
--- a/server/src/test/java/com/vaadin/tests/server/component/StateGetDoesNotMarkDirtyTest.java
+++ b/server/src/test/java/com/vaadin/tests/server/component/StateGetDoesNotMarkDirtyTest.java
@@ -127,7 +127,8 @@ public class StateGetDoesNotMarkDirtyTest {
             if (component instanceof UI) {
                 return component;
             }
-            if (component instanceof Composite) {
+            if (component.getClass().equals(Composite.class)) {
+                // Plain Composite needs a root.
                 ComponentRootSetter.setRoot(component, new Label());
             }
             emulateAttach(component);