Fix AbstractSingleSelect selection and state updates (#10796)

3 files changed, 118 insertions, 125 deletions

diff --git a/server/src/main/java/com/vaadin/ui/AbstractSingleSelect.java b/server/src/main/java/com/vaadin/ui/AbstractSingleSelect.java
index 541480ebd6..a4dad34830 100644
--- a/server/src/main/java/com/vaadin/ui/AbstractSingleSelect.java
+++ b/server/src/main/java/com/vaadin/ui/AbstractSingleSelect.java
@@ -28,6 +28,7 @@ import org.jsoup.nodes.Element;
 import com.vaadin.data.HasValue;
 import com.vaadin.data.SelectionModel.Single;
 import com.vaadin.data.provider.DataCommunicator;
+import com.vaadin.data.provider.DataGenerator;
 import com.vaadin.event.selection.SingleSelectionEvent;
 import com.vaadin.event.selection.SingleSelectionListener;
 import com.vaadin.shared.Registration;
@@ -37,6 +38,7 @@ import com.vaadin.ui.declarative.DesignContext;
 import com.vaadin.ui.declarative.DesignException;
 
 import elemental.json.Json;
+import elemental.json.JsonObject;
 
 /**
  * An abstract base class for listing components that only support single
@@ -54,6 +56,8 @@ import elemental.json.Json;
 public abstract class AbstractSingleSelect<T> extends AbstractListing<T>
         implements SingleSelect<T> {
 
+    private T selectedItem = null;
+
     /**
      * Creates a new {@code AbstractListing} with a default data communicator.
      * <p>
@@ -101,7 +105,7 @@ public abstract class AbstractSingleSelect<T> extends AbstractListing<T>
      *         otherwise
      */
     public Optional<T> getSelectedItem() {
-        return Optional.ofNullable(keyToItem(getSelectedKey()));
+        return Optional.ofNullable(selectedItem);
     }
 
     /**
@@ -112,7 +116,7 @@ public abstract class AbstractSingleSelect<T> extends AbstractListing<T>
      *            the item to select or {@code null} to clear selection
      */
     public void setSelectedItem(T item) {
-        setSelectedFromServer(item);
+        setSelectedItem(item, false);
     }
 
     /**
@@ -188,108 +192,6 @@ public abstract class AbstractSingleSelect<T> extends AbstractListing<T>
     }
 
     /**
-     * Returns the communication key of the selected item or {@code null} if no
-     * item is selected.
-     *
-     * @return the key of the selected item if any, {@code null} otherwise.
-     */
-    protected String getSelectedKey() {
-        return getState(false).selectedItemKey;
-    }
-
-    /**
-     * Sets the selected item based on the given communication key. If the key
-     * is {@code null}, clears the current selection if any.
-     *
-     * @param key
-     *            the key of the selected item or {@code null} to clear
-     *            selection
-     */
-    protected void doSetSelectedKey(String key) {
-        getState().selectedItemKey = key;
-    }
-
-    /**
-     * Sets the selection based on a client request. Does nothing if the select
-     * component is {@linkplain #isReadOnly()} or if the selection would not
-     * change. Otherwise updates the selection and fires a selection change
-     * event with {@code isUserOriginated == true}.
-     *
-     * @param key
-     *            the key of the item to select or {@code null} to clear
-     *            selection
-     */
-    protected void setSelectedFromClient(String key) {
-        if (isReadOnly()) {
-            return;
-        }
-        if (isKeySelected(key)) {
-            return;
-        }
-
-        T oldSelection = getSelectedItem().orElse(getEmptyValue());
-        doSetSelectedKey(key);
-
-        // Set diffstate to something that will always send selection to client
-        updateDiffstate("selectedItemKey", Json.createObject());
-
-        fireEvent(new SingleSelectionEvent<>(AbstractSingleSelect.this,
-                oldSelection, true));
-    }
-
-    /**
-     * Sets the selection based on server API call. Does nothing if the
-     * selection would not change; otherwise updates the selection and fires a
-     * selection change event with {@code isUserOriginated == false}.
-     *
-     * @param item
-     *            the item to select or {@code null} to clear selection
-     */
-    protected void setSelectedFromServer(T item) {
-        // TODO creates a key if item not in data provider
-        String key = itemToKey(item);
-
-        if (isKeySelected(key) || isSelected(item)) {
-            return;
-        }
-
-        T oldSelection = getSelectedItem().orElse(getEmptyValue());
-        doSetSelectedKey(key);
-
-        fireEvent(new SingleSelectionEvent<>(AbstractSingleSelect.this,
-                oldSelection, false));
-    }
-
-    /**
-     * Returns whether the given key maps to the currently selected item.
-     *
-     * @param key
-     *            the key to test or {@code null} to test whether nothing is
-     *            selected
-     * @return {@code true} if the key equals the key of the currently selected
-     *         item (or {@code null} if no selection), {@code false} otherwise.
-     */
-    protected boolean isKeySelected(String key) {
-        return Objects.equals(key, getSelectedKey());
-    }
-
-    /**
-     * Returns the communication key assigned to the given item.
-     *
-     * @param item
-     *            the item whose key to return
-     * @return the assigned key
-     */
-    protected String itemToKey(T item) {
-        if (item == null) {
-            return null;
-        } else {
-            // TODO creates a key if item not in data provider
-            return getDataCommunicator().getKeyMapper().key(item);
-        }
-    }
-
-    /**
      * Returns the item that the given key is assigned to, or {@code null} if
      * there is no such item.
      *
@@ -309,7 +211,16 @@ public abstract class AbstractSingleSelect<T> extends AbstractListing<T>
      * @return {@code true} if the item is selected, {@code false} otherwise
      */
     public boolean isSelected(T item) {
-        return Objects.equals(getValue(), item);
+        if (Objects.equals(selectedItem, item)) {
+            return true;
+        }
+
+        if (item == null || selectedItem == null) {
+            return false;
+        }
+
+        return Objects.equals(getDataProvider().getId(selectedItem),
+                getDataProvider().getId(item));
     }
 
     @Override
@@ -379,16 +290,91 @@ public abstract class AbstractSingleSelect<T> extends AbstractListing<T>
 
             @Override
             public void select(String key) {
-                setSelectedFromClient(key);
+                setSelectedItem(keyToItem(key), true);
             }
 
             @Override
             public void deselect(String key) {
-                if (isKeySelected(key)) {
-                    setSelectedFromClient(null);
+                T item = keyToItem(key);
+                if (isSelected(item)) {
+                    setSelectedItem(null, true);
+                }
+            }
+        });
+        addDataGenerator(new DataGenerator<T>() {
+
+            @Override
+            public void generateData(T item, JsonObject jsonObject) {
+                if (isSelected(item)) {
+                    // Deferred update of state.
+                    updateSelectedItemState(item);
+                }
+            }
+
+            @Override
+            public void refreshData(T item) {
+                if (isSelected(item)) {
+                    selectedItem = item;
+                    // Invalidate old data
+                    updateSelectedItemState(null);
                 }
             }
         });
     }
 
+    /**
+     * This method updates the internal selection state of the server-side of
+     * {@code AbstractSingleSelect}.
+     *
+     * @param value
+     *            the value that should be selected
+     * @param userOriginated
+     *            {@code true} if selection was done by user, {@code false} if
+     *            not
+     *
+     * @since
+     */
+    protected void setSelectedItem(T value, boolean userOriginated) {
+        if (isSelected(value)) {
+            return;
+        }
+
+        // Update selection
+        T oldValue = selectedItem;
+        selectedItem = value;
+
+        // Re-generate selected item data
+        if (oldValue != null) {
+            getDataCommunicator().refresh(oldValue);
+        }
+        if (value != null) {
+            getDataCommunicator().refresh(value);
+        }
+
+        // Deselection can be handled immediately
+        updateSelectedItemState(value);
+
+        // Update diffstate to make sure null can be selected later.
+        updateDiffstate("selectedItemKey", Json.createObject());
+
+        fireEvent(new SingleSelectionEvent<>(AbstractSingleSelect.this,
+                oldValue, userOriginated));
+    }
+
+    /**
+     * This method updates the shared selection state of the
+     * {@code AbstractSingleSelect}.
+     *
+     * @param value
+     *            the value that is selected; may be {@code null}
+     *
+     * @since
+     */
+    protected void updateSelectedItemState(T value) {
+        // FIXME: If selecting a value that does not exist, this will leave and
+        // extra object in the key mapper that will not be dropped any time.
+        getState().selectedItemKey = value != null
+                ? getDataCommunicator().getKeyMapper().key(value)
+                : null;
+    }
 }
diff --git a/server/src/main/java/com/vaadin/ui/ComboBox.java b/server/src/main/java/com/vaadin/ui/ComboBox.java
index 107e503256..494871bb2c 100644
--- a/server/src/main/java/com/vaadin/ui/ComboBox.java
+++ b/server/src/main/java/com/vaadin/ui/ComboBox.java
@@ -680,9 +680,6 @@ public class ComboBox<T> extends AbstractSingleSelect<T>
     public void setItemCaptionGenerator(
             ItemCaptionGenerator<T> itemCaptionGenerator) {
         super.setItemCaptionGenerator(itemCaptionGenerator);
-        if (getSelectedItem().isPresent()) {
-            updateSelectedItemCaption();
-        }
     }
 
     /**
@@ -724,10 +721,6 @@ public class ComboBox<T> extends AbstractSingleSelect<T>
     @Override
     public void setItemIconGenerator(IconGenerator<T> itemIconGenerator) {
         super.setItemIconGenerator(itemIconGenerator);
-
-        if (getSelectedItem().isPresent()) {
-            updateSelectedItemIcon();
-        }
     }
 
     @Override
@@ -819,25 +812,23 @@ public class ComboBox<T> extends AbstractSingleSelect<T>
     }
 
     @Override
-    protected void doSetSelectedKey(String key) {
-        super.doSetSelectedKey(key);
+    protected void updateSelectedItemState(T value) {
+        super.updateSelectedItemState(value);
 
-        updateSelectedItemCaption();
-        updateSelectedItemIcon();
+        updateSelectedItemCaption(value);
+        updateSelectedItemIcon(value);
     }
 
-    private void updateSelectedItemCaption() {
+    private void updateSelectedItemCaption(T value) {
         String selectedCaption = null;
-        T value = keyToItem(getSelectedKey());
         if (value != null) {
             selectedCaption = getItemCaptionGenerator().apply(value);
         }
         getState().selectedItemCaption = selectedCaption;
     }
 
-    private void updateSelectedItemIcon() {
+    private void updateSelectedItemIcon(T value) {
         String selectedItemIcon = null;
-        T value = keyToItem(getSelectedKey());
         if (value != null) {
             Resource icon = getItemIconGenerator().apply(value);
             if (icon != null) {
@@ -859,7 +850,8 @@ public class ComboBox<T> extends AbstractSingleSelect<T>
     public void attach() {
         super.attach();
 
-        updateSelectedItemIcon();
+        // Update icon for ConnectorResource
+        updateSelectedItemIcon(getValue());
     }
 
     @Override
diff --git a/server/src/test/java/com/vaadin/ui/ComboBoxTest.java b/server/src/test/java/com/vaadin/ui/ComboBoxTest.java
index ea431ec614..74a194e96e 100644
--- a/server/src/test/java/com/vaadin/ui/ComboBoxTest.java
+++ b/server/src/test/java/com/vaadin/ui/ComboBoxTest.java
@@ -1,5 +1,10 @@
 package com.vaadin.ui;
 
+import static org.junit.Assert.assertArrayEquals;
+
+import java.util.ArrayList;
+import java.util.List;
+
 import org.junit.Test;
 
 import com.vaadin.server.ServerRpcManager;
@@ -8,6 +13,8 @@ import com.vaadin.tests.util.MockUI;
 
 public class ComboBoxTest {
 
+    private List<String> eventValues = new ArrayList<>();
+
     @Test
     public void testResetValue() {
         ComboBox<String> comboBox = new ComboBox<>();
@@ -15,10 +22,15 @@ public class ComboBoxTest {
 
         // Reset value whenever it changes (in a real case, this listener would
         // do something with the selected value before discarding it)
-        comboBox.addValueChangeListener(event -> comboBox.setValue(null));
+        comboBox.addValueChangeListener(event -> {
+            eventValues.add(event.getValue());
+            comboBox.setValue(null);
+        });
 
         // "Attach" the component and initialize diffstate
         new MockUI().setContent(comboBox);
+        // Generate initial data
+        comboBox.getDataCommunicator().beforeClientResponse(true);
         ComponentTest.syncToClient(comboBox);
 
         // Emulate selection of "one"
@@ -27,6 +39,9 @@ public class ComboBoxTest {
         ServerRpcManager.getRpcProxy(comboBox, SelectionServerRpc.class)
                 .select(oneKey);
 
+        assertArrayEquals("Unexpected values from selection events",
+                new Object[] { "one", null }, eventValues.toArray());
+
         ComponentTest.assertEncodedStateProperties(comboBox,
                 "Selection change done by the listener should be sent to the client",
                 "selectedItemKey");