fix: Use KeyMapper by default to generate unique keys (#13883) (#13928)

Delegates to regular KeyMapper whenever the unique key provider is not present and a unique key is needed for a row.

Part-of: vaadin/flow-components#3156

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

@@ -65,28 +65,6 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
 
     private final Map<String, HierarchicalCommunicationController<T>> dataControllers = new HashMap<>();
 
-    private KeyMapper<T> uniqueKeyMapper = new KeyMapper<T>() {
-
-        private T object;
-
-        @Override
-        public String key(T o) {
-            this.object = o;
-            try {
-                return super.key(o);
-            } finally {
-                this.object = null;
-            }
-        }
-
-        @Override
-        protected String createKey() {
-            return Optional.ofNullable(uniqueKeyProviderSupplier.get())
-                    .map(provider -> provider.apply(object))
-                    .orElse(super.createKey());
-        }
-    };
-
     /**
      * Construct a new hierarchical data communicator backed by a
      * {@link TreeDataProvider}.
@@ -113,16 +91,16 @@ public HierarchicalDataCommunicator(CompositeDataGenerator<T> dataGenerator,
         this.stateNode = stateNode;
         this.uniqueKeyProviderSupplier = uniqueKeyProviderSupplier;
 
-        setKeyMapper(uniqueKeyMapper);
+        KeyMapperWrapper<T> keyMapperWrapper = new KeyMapperWrapper<>();
+        setKeyMapper(keyMapperWrapper);
 
         dataGenerator.addDataGenerator(this::generateTreeData);
         setDataProvider(new TreeDataProvider<>(new TreeData<>()), null);
     }
 
     private void generateTreeData(T item, JsonObject jsonObject) {
-        Optional.ofNullable(getParentItem(item))
-                .ifPresent(parent -> jsonObject.put("parentUniqueKey",
-                        uniqueKeyProviderSupplier.get().apply(parent)));
+        Optional.ofNullable(getParentItem(item)).ifPresent(parent -> jsonObject
+                .put("parentUniqueKey", getKeyMapper().key(parent)));
     }
 
     private void requestFlush(HierarchicalUpdate update) {
@@ -180,7 +158,7 @@ protected void handleDataRefreshEvent(
         if (event.isRefreshChildren()) {
             T item = event.getItem();
             if (isExpanded(item)) {
-                String parentKey = uniqueKeyProviderSupplier.get().apply(item);
+                String parentKey = getKeyMapper().key(item);
 
                 if (!dataControllers.containsKey(parentKey)) {
                     setParentRequestedRange(0, mapper.countChildItems(item),
@@ -205,7 +183,7 @@ public Stream<T> fetchFromProvider(int offset, int limit) {
     }
 
     public void setParentRequestedRange(int start, int length, T parentItem) {
-        String parentKey = uniqueKeyProviderSupplier.get().apply(parentItem);
+        String parentKey = getKeyMapper().key(parentItem);
 
         HierarchicalCommunicationController<T> controller = dataControllers
                 .computeIfAbsent(parentKey,
@@ -561,4 +539,37 @@ private JsonValue generateJsonForExpandedOrCollapsedItem(T item) {
         return json;
     }
 
+    /**
+     * KeyMapper extension delegating row key creation to the
+     * <code>uniqueKeyProviderSupplier</code> passed to the hierarchical data
+     * communicator constructor from the component.
+     * <p>
+     * If <code>uniqueKeyProviderSupplier</code> is not present, this class uses
+     * {@link KeyMapper#createKey()} for key creation.
+     *
+     * @param <V>
+     *            the bean type
+     */
+    private class KeyMapperWrapper<V> extends KeyMapper<T> {
+
+        private T object;
+
+        @Override
+        public String key(T o) {
+            this.object = o;
+            try {
+                return super.key(o);
+            } finally {
+                this.object = null;
+            }
+        }
+
+        @Override
+        protected String createKey() {
+            return Optional.ofNullable(uniqueKeyProviderSupplier.get())
+                    .map(provider -> provider.apply(object))
+                    .orElse(super.createKey());
+        }
+    }
+
 }

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

@@ -15,6 +15,17 @@
  */
 package com.vaadin.flow.data.provider.hierarchy;
 
+import java.io.Serializable;
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.IntStream;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+
 import com.vaadin.flow.component.UI;
 import com.vaadin.flow.data.provider.CompositeDataGenerator;
 import com.vaadin.flow.data.provider.DataCommunicatorTest;
@@ -24,15 +35,8 @@
 import com.vaadin.flow.server.VaadinRequest;
 import com.vaadin.flow.server.VaadinService;
 import com.vaadin.flow.server.VaadinSession;
-import elemental.json.JsonValue;
-import org.junit.Assert;
-import org.junit.Before;
-import org.junit.Test;
-import org.mockito.Mockito;
-import org.mockito.MockitoAnnotations;
 
-import java.io.Serializable;
-import java.util.List;
+import elemental.json.JsonValue;
 
 public class HierarchicalCommunicatorDataTest {
     /**
@@ -75,6 +79,7 @@ public int hashCode() {
     private HierarchicalDataCommunicator<Item> communicator;
     private TreeData<Item> treeData;
     private MockUI ui;
+    private Element element;
 
     private static class UpdateQueue implements HierarchicalUpdate {
         @Override
@@ -125,7 +130,7 @@ public void initialize() {
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         ui = new MockUI();
-        Element element = new Element("div");
+        element = new Element("div");
         ui.getElement().appendChild(element);
 
         treeData = new TreeData<>();
@@ -167,6 +172,47 @@ public void sameKeyDifferentInstance_latestInstanceUsed() {
         Assert.assertSame(updatedLeaf, communicator.getKeyMapper().get(key));
     }
 
+    @Test
+    public void uniqueKeyProviderIsSet_keysGeneratedByProvider() {
+        communicator = new HierarchicalDataCommunicator<Item>(
+                Mockito.mock(CompositeDataGenerator.class), arrayUpdater,
+                json -> {
+                }, element.getNode(), () -> item -> item.value);
+        communicator.setDataProvider(dataProvider, null);
+        // expand test items to force key calculation
+        communicator.expand(ROOT);
+        communicator.expand(FOLDER);
+        // this is needed to force key calculation for leaf item
+        communicator.setParentRequestedRange(0, 50, LEAF);
+        fakeClientCommunication();
+
+        Arrays.asList("ROOT", "FOLDER", "LEAF")
+                .forEach(key -> Assert.assertNotNull("Expected key '" + key
+                        + "' to be generated when unique key provider used",
+                        communicator.getKeyMapper().get(key)));
+    }
+
+    @Test
+    public void uniqueKeyProviderIsNotSet_keysGeneratedByKeyMapper() {
+        communicator = new HierarchicalDataCommunicator<Item>(
+                Mockito.mock(CompositeDataGenerator.class), arrayUpdater,
+                json -> {
+                }, element.getNode(), () -> null);
+        communicator.setDataProvider(dataProvider, null);
+        // expand test items to force key calculation
+        communicator.expand(ROOT);
+        communicator.expand(FOLDER);
+        // this is needed to force key calculation for leaf item
+        communicator.setParentRequestedRange(0, 50, LEAF);
+        fakeClientCommunication();
+
+        // key mapper should generate keys 1,2,3
+        IntStream.range(1, 4).mapToObj(String::valueOf)
+                .forEach(i -> Assert.assertNotNull("Expected key '" + i
+                        + "' to be generated when unique key provider is not set",
+                        communicator.getKeyMapper().get(i)));
+    }
+
     private void fakeClientCommunication() {
         ui.getInternals().getStateTree().runExecutionsBeforeClientResponse();
         ui.getInternals().getStateTree().collectChanges(ignore -> {