Fixed another memory leak in Tree #11053

Change-Id: Idf45309b83dfe718fe6e58b6cf57b211eeb9e811

2 files changed, 34 insertions, 3 deletions

diff --git a/server/src/com/vaadin/ui/Tree.java b/server/src/com/vaadin/ui/Tree.java
index 219b39020b..32c5712f0f 100644
--- a/server/src/com/vaadin/ui/Tree.java
+++ b/server/src/com/vaadin/ui/Tree.java
@@ -414,6 +414,9 @@ public class Tree extends AbstractSelect implements Container.Hierarchical,
                 final Object id = itemIdMapper.get(keys[i]);
                 if (id != null && isExpanded(id)) {
                     expanded.remove(id);
+                    if (expandedItemId == id) {
+                        expandedItemId = null;
+                    }
                     fireCollapseEvent(id);
                 }
             }
@@ -841,6 +844,10 @@ public class Tree extends AbstractSelect implements Container.Hierarchical,
             super.setContainerDataSource(new ContainerHierarchicalWrapper(
                     newDataSource));
         }
+
+        // Ensure previous expanded items are cleaned up if they don't exist in
+        // the new container
+        cleanupExpandedItems();
     }
 
     /* Expand event and listener */
@@ -1675,6 +1682,12 @@ public class Tree extends AbstractSelect implements Container.Hierarchical,
             com.vaadin.data.Container.ItemSetChangeEvent event) {
 
         // Ensure removed items are cleaned up from expanded list
+        cleanupExpandedItems();
+
+        super.containerItemSetChange(event);
+    }
+
+    private void cleanupExpandedItems() {
         for (Object expandedItemId : expanded) {
             if (getItem(expandedItemId) == null) {
                 expanded.remove(expandedItemId);
@@ -1683,8 +1696,5 @@ public class Tree extends AbstractSelect implements Container.Hierarchical,
                 }
             }
         }
-
-        super.containerItemSetChange(event);
     }
-
 }
diff --git a/server/tests/src/com/vaadin/tests/server/component/tree/TreeTest.java b/server/tests/src/com/vaadin/tests/server/component/tree/TreeTest.java
index f467a42339..c1d7653c01 100644
--- a/server/tests/src/com/vaadin/tests/server/component/tree/TreeTest.java
+++ b/server/tests/src/com/vaadin/tests/server/component/tree/TreeTest.java
@@ -112,4 +112,25 @@ public class TreeTest extends TestCase {
         assertEquals(0, expanded.size());
         assertNull(expandedItemId);
     }
+
+    public void testRemoveExpandedItemsOnContainerChange() throws Exception {
+        tree.expandItem("parent");
+        tree.expandItem("child");
+
+        tree.setContainerDataSource(new HierarchicalContainer());
+
+        Field expandedField = tree.getClass().getDeclaredField("expanded");
+        Field expandedItemIdField = tree.getClass().getDeclaredField(
+                "expandedItemId");
+
+        expandedField.setAccessible(true);
+        expandedItemIdField.setAccessible(true);
+
+        HashSet<Object> expanded = (HashSet<Object>) expandedField.get(tree);
+        assertEquals(0, expanded.size());
+
+        Object expandedItemId = expandedItemIdField.get(tree);
+        assertNull(expandedItemId);
+    }
+
 }