Fix Grid not using getId on select all

server/src/main/java/com/vaadin/ui/components/grid/MultiSelectionModelImpl.java

@@ -15,30 +15,19 @@
  */
 package com.vaadin.ui.components.grid;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.LinkedHashSet;
-import java.util.List;
-import java.util.Objects;
-import java.util.Set;
-import java.util.function.Consumer;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
-
-import com.vaadin.data.provider.DataCommunicator;
-import com.vaadin.data.provider.DataProvider;
-import com.vaadin.data.provider.HierarchicalDataProvider;
-import com.vaadin.data.provider.HierarchicalQuery;
-import com.vaadin.data.provider.Query;
+import com.vaadin.data.provider.*;
 import com.vaadin.event.selection.MultiSelectionEvent;
 import com.vaadin.event.selection.MultiSelectionListener;
 import com.vaadin.shared.Registration;
 import com.vaadin.shared.data.selection.GridMultiSelectServerRpc;
 import com.vaadin.shared.ui.grid.MultiSelectionModelState;
 import com.vaadin.ui.MultiSelect;
 
+import java.util.*;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
 /**
  * Multiselection model for grid.
  * <p>
@@ -429,12 +418,20 @@ protected void updateSelection(Set<T> addedItems, Set<T> removedItems,
                     + " although user selection is disallowed");
         }
 
-        // if there are duplicates, some item is both added & removed, just
-        // discard that and leave things as was before
-        addedItems.removeIf(item -> removedItems.remove(item));
+        DataProvider<T, ?> dataProvider = getGrid().getDataProvider();
 
-        if (selection.containsAll(addedItems)
-                && Collections.disjoint(selection, removedItems)) {
+        addedItems.removeIf(item -> {
+            Object id = dataProvider.getId(item);
+            Optional<T> toRemove = removedItems.stream()
+                    .filter(i -> dataProvider.getId(i).equals(id)).findFirst();
+            toRemove.ifPresent(i -> removedItems.remove(i));
+            return toRemove.isPresent();
+        });
+
+        if (addedItems.stream().map(dataProvider::getId)
+                .allMatch(this::selectionContainsId)
+                && removedItems.stream().map(dataProvider::getId)
+                        .noneMatch(this::selectionContainsId)) {
             return;
         }
 
@@ -446,8 +443,13 @@ protected void updateSelection(Set<T> addedItems, Set<T> removedItems,
 
         doUpdateSelection(set -> {
             // order of add / remove does not matter since no duplicates
-            set.removeAll(removedItems);
-            set.addAll(addedItems);
+            Set<Object> removedItemIds = removedItems.stream()
+                    .map(dataProvider::getId).collect(Collectors.toSet());
+            set.removeIf(
+                    item -> removedItemIds.contains(dataProvider.getId(item)));
+            addedItems.stream().filter(
+                    item -> !selectionContainsId(dataProvider.getId(item)))
+                    .forEach(set::add);
 
             // refresh method is NOOP for items that are not present client side
             DataCommunicator<T> dataCommunicator = getGrid()

server/src/test/java/com/vaadin/tests/components/grid/GridMultiSelectionModelProxyItemTest.java

@@ -0,0 +1,99 @@
+package com.vaadin.tests.components.grid;
+
+import com.liferay.portal.kernel.atom.AtomEntryContent;
+import com.vaadin.data.provider.CallbackDataProvider;
+import com.vaadin.ui.Grid;
+import com.vaadin.ui.components.grid.MultiSelectionModel;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.*;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+public class GridMultiSelectionModelProxyItemTest {
+
+    private List<String> data = IntStream.range(0, 100).boxed()
+            .map(i -> "String " + i).collect(Collectors.toList());
+    private Grid<AtomicReference<String>> proxyGrid = new Grid<>();
+    private MultiSelectionModel<AtomicReference<String>> model;
+    private AtomicReference<Set<AtomicReference<String>>> selectionEvent = new AtomicReference<>();
+
+    @Before
+    public void setup() {
+        proxyGrid.setDataProvider(new CallbackDataProvider<>(
+                q -> data.stream().map(AtomicReference::new).skip(q.getOffset())
+                        .limit(q.getLimit()),
+                q -> data.size(), AtomicReference::get));
+        model = (MultiSelectionModel<AtomicReference<String>>) proxyGrid
+                .setSelectionMode(Grid.SelectionMode.MULTI);
+        model.addSelectionListener(e -> {
+            selectionEvent.set(e.getAllSelectedItems());
+        });
+    }
+
+    @Test
+    public void testSelectAllWithProxyDataProvider() {
+        model.selectAll();
+        assertEquals("Item count mismatch on first select all", 100,
+                getSelectionEvent().size());
+        model.deselect(model.getFirstSelectedItem().orElseThrow(
+                () -> new IllegalStateException("Items should be selected")));
+        assertEquals("Item count mismatch on deselect", 99,
+                getSelectionEvent().size());
+        model.selectAll();
+        assertEquals("Item count mismatch on second select all", 100,
+                getSelectionEvent().size());
+    }
+
+    @Test
+    public void testUpdateSelectionWithDuplicateEntries() {
+        List<String> selection = data.stream().filter(s -> s.contains("1"))
+                .collect(Collectors.toList());
+        model.updateSelection(selection.stream().map(AtomicReference::new)
+                .collect(Collectors.toSet()), Collections.emptySet());
+        assertEquals("Failure in initial selection", selection.size(),
+                getSelectionEvent().size());
+
+        String toRemove = model.getFirstSelectedItem().map(AtomicReference::get)
+                .orElseThrow(() -> new IllegalStateException(
+                        "Items should be selected"));
+        model.updateSelection(
+                Stream.of(toRemove).map(AtomicReference::new)
+                        .collect(Collectors.toSet()),
+                Stream.of(toRemove).map(AtomicReference::new)
+                        .collect(Collectors.toSet()));
+        assertNull(
+                "Selection should not change when selecting and deselecting once",
+                selectionEvent.get());
+
+        Set<AtomicReference<String>> added = new LinkedHashSet<>();
+        Set<AtomicReference<String>> removed = new LinkedHashSet<>();
+        for (int i = 0; i < 20; ++i) {
+            added.add(new AtomicReference<>(toRemove));
+            removed.add(new AtomicReference<>(toRemove));
+        }
+        model.updateSelection(added, removed);
+        assertNull(
+                "Selection should not change when selecting and deselecting 20 times",
+                selectionEvent.get());
+
+        removed.add(new AtomicReference<>(toRemove));
+        model.updateSelection(added, removed);
+        assertEquals("Item should have been deselected", selection.size() - 1,
+                getSelectionEvent().size());
+    }
+
+    private Set<AtomicReference<String>> getSelectionEvent() {
+        Optional<Set<AtomicReference<String>>> eventOptional = Optional
+                .of(selectionEvent.get());
+        selectionEvent.set(null);
+        return eventOptional.orElseThrow(() -> new IllegalStateException(
+                "Expected selection event never happened"));
+    }
+}

server/src/test/java/com/vaadin/tests/components/grid/GridMultiSelectionModelTest.java

@@ -17,8 +17,11 @@
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import java.util.stream.Stream;
 
+import com.vaadin.data.provider.CallbackDataProvider;
 import org.easymock.Capture;
 import org.junit.Before;
 import org.junit.Test;
@@ -703,6 +706,5 @@ public void selectAllCheckboxVisible__lazyDataProvider() {
         assertFalse(model.isSelectAllCheckBoxVisible());
         assertEquals(SelectAllCheckBoxVisibility.DEFAULT,
                 model.getSelectAllCheckBoxVisibility());
-
     }
 }