Fix TreeData item removal, small improvements (#9621)

Fixes #9614

6 files changed, 136 insertions, 17 deletions

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 054e86ca8d..2855d21cbe 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
@@ -187,6 +187,7 @@ public class DataCommunicator<T> extends AbstractExtension {
         public void destroyAllData() {
             droppedData.clear();
             activeData.clear();
+            updatedData.clear();
             getKeyMapper().removeAll();
         }
     }
@@ -524,7 +525,6 @@ public class DataCommunicator<T> extends AbstractExtension {
             if (updatedData.isEmpty()) {
                 markAsDirty();
             }
-
             updatedData.add(activeData.get(id));
         }
     }
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 f53287b1df..c4f99e2b5b 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
@@ -274,7 +274,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
 
     /**
      * Gets the item collapse allowed provider.
-     * 
+     *
      * @return the item collapse allowed provider
      */
     public ItemCollapseAllowedProvider<T> getItemCollapseAllowedProvider() {
@@ -309,4 +309,5 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
         }
         super.setFilter(filter);
     }
+
 }
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 ebed464fde..73343d8312 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
@@ -92,12 +92,10 @@ public class HierarchyMapper<T, F> implements DataGenerator<T> {
      * 
      * @param item
      *            the item to get the parent of
-     * @return the parent index
+     * @return the parent index or a negative value if the parent is not found
      * 
-     * @throws IllegalArgumentException
-     *             if given target index is not found
      */
-    public Integer getParentIndex(T item) throws IllegalArgumentException {
+    public Integer getParentIndex(T item) {
         // TODO: This can be optimised.
         List<T> flatHierarchy = getHierarchy(null).collect(Collectors.toList());
         return flatHierarchy.indexOf(getParentOfItem(item));
@@ -509,4 +507,9 @@ public class HierarchyMapper<T, F> implements DataGenerator<T> {
                 : Stream.empty();
         return Stream.concat(parentStream, children);
     }
+
+    @Override
+    public void destroyAllData() {
+        childMap.clear();
+    }
 }
diff --git a/server/src/main/java/com/vaadin/data/provider/TreeDataProvider.java b/server/src/main/java/com/vaadin/data/provider/TreeDataProvider.java
index ed5e710cc8..e12372a625 100644
--- a/server/src/main/java/com/vaadin/data/provider/TreeDataProvider.java
+++ b/server/src/main/java/com/vaadin/data/provider/TreeDataProvider.java
@@ -72,11 +72,9 @@ public class TreeDataProvider<T>
     @Override
     public boolean hasChildren(T item) {
         if (!treeData.contains(item)) {
-            throw new IllegalArgumentException("Item " + item
-                    + " could not be found in the backing TreeData. "
-                    + "Did you forget to refresh this data provider after item removal?");
+            // The item might be dropped from the tree already
+            return false;
         }
-
         return !treeData.getChildren(item).isEmpty();
     }
 
diff --git a/server/src/test/java/com/vaadin/data/provider/DataCommunicatorTest.java b/server/src/test/java/com/vaadin/data/provider/DataCommunicatorTest.java
index c6d80e6ab0..53913325a5 100644
--- a/server/src/test/java/com/vaadin/data/provider/DataCommunicatorTest.java
+++ b/server/src/test/java/com/vaadin/data/provider/DataCommunicatorTest.java
@@ -33,6 +33,8 @@ import com.vaadin.ui.UI;
 
 import elemental.json.JsonObject;
 
+import static org.junit.Assert.assertFalse;
+
 /**
  * @author Vaadin Ltd
  *
@@ -41,11 +43,11 @@ public class DataCommunicatorTest {
 
     private static final Object TEST_OBJECT = new Object();
 
-    private static class TestUI extends UI {
+    public static class TestUI extends UI {
 
         private final VaadinSession session;
 
-        TestUI(VaadinSession session) {
+        public TestUI(VaadinSession session) {
             this.session = session;
         }
 
@@ -128,7 +130,7 @@ public class DataCommunicatorTest {
         TestDataProvider dataProvider = new TestDataProvider();
         communicator.setDataProvider(dataProvider, null);
 
-        Assert.assertFalse(dataProvider.isListenerAdded());
+        assertFalse(dataProvider.isListenerAdded());
 
         communicator.extend(ui);
 
@@ -152,7 +154,7 @@ public class DataCommunicatorTest {
 
         communicator.detach();
 
-        Assert.assertFalse(dataProvider.isListenerAdded());
+        assertFalse(dataProvider.isListenerAdded());
     }
 
     @Test
@@ -203,15 +205,19 @@ public class DataCommunicatorTest {
         communicator.extend(ui);
         // Put a test object into a cache
         communicator.pushData(1, Collections.singletonList(TEST_OBJECT));
+        // Put the test object into an update queue
+        communicator.refresh(TEST_OBJECT);
         // Drop the test object from the cache
         String key = communicator.getKeyMapper().key(TEST_OBJECT);
         JsonArray keys = Json.createArray();
-        keys.set(0,key);
+        keys.set(0, key);
         communicator.onDropRows(keys);
-        // Drop everything
-        communicator.dropAllData();
+        // Replace everything
+        communicator.setDataProvider(new ListDataProvider<>(Collections.singleton(new Object())));
         // The communicator does not have to throw exceptions during
         // request finalization
         communicator.beforeClientResponse(false);
+        assertFalse("Stalled object in KeyMapper",
+                communicator.getKeyMapper().has(TEST_OBJECT));
     }
 }
diff --git a/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchicalCommunicatorTest.java b/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchicalCommunicatorTest.java
new file mode 100644
index 0000000000..cdf71a21f1
--- /dev/null
+++ b/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchicalCommunicatorTest.java
@@ -0,0 +1,111 @@
+package com.vaadin.data.provider.hierarchical;
+
+import com.vaadin.data.TreeData;
+import com.vaadin.data.provider.DataCommunicatorTest;
+import com.vaadin.data.provider.HierarchicalDataCommunicator;
+import com.vaadin.data.provider.TreeDataProvider;
+import com.vaadin.server.AbstractClientConnector;
+import com.vaadin.server.MockVaadinSession;
+import com.vaadin.server.VaadinService;
+import com.vaadin.server.VaadinSession;
+import com.vaadin.ui.UI;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.util.Arrays;
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+
+public class HierarchicalCommunicatorTest {
+
+    private static final String ROOT = "ROOT";
+    private static final String FOLDER = "FOLDER";
+    private static final String LEAF = "LEAF";
+    private TreeDataProvider<String> dataProvider;
+    private TestHierarchicalDataCommunicator<String> communicator;
+    private TreeData<String> treeData;
+
+    @Before
+    public void setUp() {
+        VaadinSession session = new MockVaadinSession(
+                Mockito.mock(VaadinService.class));
+        session.lock();
+        UI ui = new DataCommunicatorTest.TestUI(session);
+        treeData = new TreeData<>();
+        treeData.addItems(null, ROOT);
+        treeData.addItems(ROOT, FOLDER);
+        treeData.addItems(FOLDER, LEAF);
+        dataProvider = new TreeDataProvider<>(treeData);
+        communicator = new TestHierarchicalDataCommunicator<>();
+        communicator.extend(ui);
+        communicator.setDataProvider(dataProvider, null);
+        communicator.attach();
+    }
+
+    @Test
+    public void testFolderRemoveRefreshAll() {
+        testItemRemove(FOLDER, true);
+    }
+
+    @Test
+    public void testLeafRemoveRefreshAll() {
+        testItemRemove(LEAF, true);
+    }
+
+    @Test
+    public void testFolderRemove() {
+        testItemRemove(FOLDER, false);
+    }
+
+    @Test
+    public void testLeafRemove() {
+        testItemRemove(LEAF, false);
+    }
+
+    private void testItemRemove(String item, boolean refreshAll) {
+        communicator.pushData(1, Arrays.asList(ROOT, FOLDER, LEAF));
+        communicator.expand(ROOT);
+        communicator.expand(FOLDER);
+        //Put the item into client queue
+        communicator.refresh(item);
+        treeData.removeItem(item);
+        if (refreshAll) {
+            dataProvider.refreshAll();
+        } else {
+            dataProvider.refreshItem(item);
+        }
+        communicator.beforeClientResponse(false);
+    }
+
+    @Test
+    public void testReplaceAll() {
+        communicator.pushData(1, Arrays.asList(ROOT, FOLDER, LEAF));
+        //Some modifications
+        communicator.expand(ROOT);
+        communicator.expand(FOLDER);
+        communicator.refresh(LEAF);
+        //Replace dataprovider
+        communicator.setDataProvider(new TreeDataProvider<>(new TreeData<>()), null);
+        dataProvider.refreshAll();
+        communicator.beforeClientResponse(false);
+        assertFalse("Stalled object in KeyMapper",
+                communicator.getKeyMapper().has(ROOT));
+        assertEquals(-1,communicator.getParentIndex(FOLDER).longValue());
+    }
+
+    private static class TestHierarchicalDataCommunicator<T> extends HierarchicalDataCommunicator<T> {
+        @Override
+        public void extend(AbstractClientConnector target) {
+            super.extend(target);
+        }
+
+        @Override
+        public void pushData(int firstIndex, List<T> data) {
+            super.pushData(firstIndex, data);
+        }
+    }
+}