fix: Sync HierarchicalDataCommunicator's expand state with client side (

#9275)

HierarchicalDataCommunicator's expanded state were not being synchronised with client side and TreeGrid's expanded nodes were being collapsed after re-attaching.

Warranty: Sync TreeGrid expanded items state with the client side when detaching and reattaching

Fixes: #9175

(cherry picked from commit bb51a2b)

flow-data/src/main/java/com/vaadin/flow/data/provider/hierarchy/HierarchicalDataCommunicator.java

@@ -158,6 +158,19 @@ public void reset() {
             HierarchicalUpdate update = arrayUpdater
                     .startUpdate(getHierarchyMapper().getRootSize());
             update.enqueue("$connector.ensureHierarchy");
+
+            Collection<T> expandedItems = getHierarchyMapper().getExpandedItems();
+            update.enqueue("$connector.expandItems",
+                    expandedItems
+                            .stream()
+                            .map(getKeyMapper()::key)
+                            .map(key -> {
+                                JsonObject json = Json.createObject();
+                                json.put("key", key);
+                                return json;
+                            }).collect(
+                            JsonUtils.asArray()));
+
             requestFlush(update);
         }
     }

flow-data/src/main/java/com/vaadin/flow/data/provider/hierarchy/HierarchyMapper.java

@@ -16,6 +16,7 @@
 package com.vaadin.flow.data.provider.hierarchy;
 
 import java.io.Serializable;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.HashMap;
@@ -63,7 +64,7 @@ public class HierarchyMapper<T, F> implements Serializable {
     private List<QuerySortOrder> backEndSorting;
     private Comparator<T> inMemorySorting;
 
-    private Set<Object> expandedItemIds = new HashSet<>();
+    private Map<Object, T> expandedItems = new HashMap<>();
 
     /**
      * Constructs a new HierarchyMapper.
@@ -132,7 +133,7 @@ public boolean isExpanded(T item) {
             // Root nodes are always visible.
             return true;
         }
-        return expandedItemIds.contains(getDataProvider().getId(item));
+        return expandedItems.containsKey(getDataProvider().getId(item));
     }
 
     /**
@@ -177,7 +178,7 @@ public Range expand(T item, Integer position) {
     private boolean doExpand(T item) {
         boolean expanded = false;
         if (!isExpanded(item) && hasChildren(item)) {
-            expandedItemIds.add(getDataProvider().getId(item));
+            expandedItems.put(getDataProvider().getId(item), item);
             expanded = true;
         }
         return expanded;
@@ -194,7 +195,7 @@ public boolean collapse(T item) {
             return false;
         }
         if (isExpanded(item)) {
-            expandedItemIds.remove(getDataProvider().getId(item));
+            expandedItems.remove(getDataProvider().getId(item));
             return true;
         }
         return false;
@@ -217,7 +218,7 @@ public Range collapse(T item, Integer position) {
                 removedRows = Range.withLength(position + 1,
                         (int) getHierarchy(item, false).count());
             }
-            expandedItemIds.remove(getDataProvider().getId(item));
+            expandedItems.remove(getDataProvider().getId(item));
         }
         return removedRows;
     }
@@ -439,7 +440,7 @@ protected void removeChildren(Object id) {
                 iterator.remove();
             }
         }
-        expandedItemIds.remove(id);
+        expandedItems.remove(id);
         invalidatedChildren.stream().map(getDataProvider()::getId)
                 .forEach(x -> {
                     removeChildren(x);
@@ -616,7 +617,7 @@ private Stream<T> combineParentAndChildStreams(T parent, Stream<T> children,
     public void destroyAllData() {
         childMap.clear();
         parentIdMap.clear();
-        expandedItemIds.clear();
+        expandedItems.clear();
     }
 
     /**
@@ -625,6 +626,16 @@ public void destroyAllData() {
      * @return {@code true} if there is any expanded items.
      */
     public boolean hasExpandedItems() {
-        return !expandedItemIds.isEmpty();
+        return !expandedItems.isEmpty();
+    }
+
+    /**
+     * Returns the expanded items in form of an unmodifiable collection.
+     *
+     * @return an unmodifiable {@code Collection<T>} containing the expanded
+     * items.
+     */
+    public Collection<T> getExpandedItems() {
+        return Collections.unmodifiableCollection(expandedItems.values());
     }
 }

flow-data/src/test/java/com/vaadin/flow/data/provider/hierarchy/HierarchicalCommunicatorTest.java

@@ -15,12 +15,11 @@
  */
 package com.vaadin.flow.data.provider.hierarchy;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import org.junit.Assert;
 import org.junit.Before;
@@ -36,7 +35,11 @@
 import com.vaadin.flow.internal.StateNode;
 import com.vaadin.flow.internal.StateTree;
 
+import elemental.json.JsonArray;
+import elemental.json.JsonObject;
 import elemental.json.JsonValue;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 
 public class HierarchicalCommunicatorTest {
 
@@ -54,6 +57,9 @@ public class HierarchicalCommunicatorTest {
 
     private List<String> enqueueFunctions = new ArrayList<>();
 
+    private Map<String, Serializable[]> enqueueFunctionsWithParams =
+            new HashMap<>();
+
     private class UpdateQueue implements HierarchicalUpdate {
         @Override
         public void clear(int start, int length) {
@@ -89,6 +95,13 @@ public void commit() {
         }
     }
 
+    private class UpdateQueueWithArguments extends UpdateQueue {
+        @Override
+        public void enqueue(String name, Serializable... arguments) {
+            enqueueFunctionsWithParams.put(name, arguments);
+        }
+    }
+
     private final HierarchicalArrayUpdater arrayUpdater = new HierarchicalArrayUpdater() {
         @Override
         public HierarchicalUpdate startUpdate(int sizeChange) {
@@ -100,6 +113,18 @@ public void initialize() {
         }
     };
 
+    private final HierarchicalArrayUpdater arrayUpdaterWithArguments =
+            new HierarchicalArrayUpdater() {
+        @Override
+        public HierarchicalUpdate startUpdate(int sizeChange) {
+            return new UpdateQueueWithArguments();
+        }
+
+        @Override
+        public void initialize() {
+        }
+    };
+
     @Before
     public void setUp() {
         ui = Mockito.mock(UI.class);
@@ -188,9 +213,61 @@ public void reset_noDataControllers_hierarchicalUpdateIsCalled() {
         // any data controllers
         communicator.reset();
 
-        Assert.assertEquals(1, enqueueFunctions.size());
+        Assert.assertEquals(2, enqueueFunctions.size());
         Assert.assertEquals("$connector.ensureHierarchy",
                 enqueueFunctions.get(0));
+        Assert.assertEquals("$connector.expandItems",
+                enqueueFunctions.get(1));
+    }
+
+    @Test
+    public void reset_expandSomeItems_updateContainsProperJsonObjectsToExpand() {
+        enqueueFunctionsWithParams = new HashMap<>();
+
+        TreeData<String> hierarchyTreeData = new TreeData<>();
+        hierarchyTreeData.addItem(null, "root");
+        hierarchyTreeData.addItem("root", "first-1");
+        hierarchyTreeData.addItem("root", "first-2");
+        hierarchyTreeData.addItem("first-1", "second-1-1");
+        hierarchyTreeData.addItem("first-2", "second-2-1");
+
+        TreeDataProvider<String> treeDataProvider =
+                new TreeDataProvider<>(hierarchyTreeData);
+
+        HierarchicalDataCommunicator<String> dataCommunicator =
+                new HierarchicalDataCommunicator<String>(
+                        Mockito.mock(CompositeDataGenerator.class),
+                        arrayUpdaterWithArguments,
+                        json -> {}, Mockito.mock(StateNode.class), () -> null);
+
+        dataCommunicator.setDataProvider(treeDataProvider, null);
+
+        dataCommunicator.expand("root");
+        dataCommunicator.expand("first-1");
+
+        dataCommunicator.reset();
+
+
+        Assert.assertTrue(enqueueFunctionsWithParams
+                .containsKey("$connector.expandItems"));
+        JsonArray arguments = (JsonArray) enqueueFunctionsWithParams
+                .get("$connector.expandItems")[0];
+        Assert.assertNotNull(arguments);
+        Assert.assertEquals(2, arguments.length());
+
+        JsonObject first1 = arguments.getObject(0);
+        JsonObject root = arguments.getObject(1);
+
+        Assert.assertNotNull(first1);
+        Assert.assertNotNull(root);
+
+        Assert.assertTrue(first1.hasKey("key"));
+        Assert.assertTrue(root.hasKey("key"));
+
+        Assert.assertEquals(dataCommunicator.getKeyMapper().key("first-1"),
+                first1.getString("key"));
+        Assert.assertEquals(dataCommunicator.getKeyMapper().key("root"),
+                root.getString("key"));
     }
 
 }