Make select all work with TreeGrid (#9412)

Fixes #9403

3 files changed, 133 insertions, 3 deletions

diff --git a/server/src/main/java/com/vaadin/ui/components/grid/MultiSelectionModelImpl.java b/server/src/main/java/com/vaadin/ui/components/grid/MultiSelectionModelImpl.java
index 01635302f0..289db2a04f 100644
--- a/server/src/main/java/com/vaadin/ui/components/grid/MultiSelectionModelImpl.java
+++ b/server/src/main/java/com/vaadin/ui/components/grid/MultiSelectionModelImpl.java
@@ -29,6 +29,8 @@ 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.event.selection.MultiSelectionEvent;
 import com.vaadin.event.selection.MultiSelectionListener;
@@ -312,15 +314,70 @@ public class MultiSelectionModelImpl<T> extends AbstractSelectionModel<T>
             getState().allSelected = true;
         }
 
-        DataProvider<T, ?> dataSource = getGrid().getDataProvider();
+        Stream<T> allItemsStream;
+        DataProvider<T, ?> dataProvider = getGrid().getDataProvider();
         // this will fetch everything from backend
-        Stream<T> stream = dataSource.fetch(new Query<>());
+        if (dataProvider instanceof HierarchicalDataProvider) {
+            allItemsStream = fetchAllHierarchical(
+                    (HierarchicalDataProvider<T, ?>) dataProvider);
+        } else {
+            allItemsStream = fetchAll(dataProvider);
+        }
         LinkedHashSet<T> allItems = new LinkedHashSet<>();
-        stream.forEach(allItems::add);
+        allItemsStream.forEach(allItems::add);
         updateSelection(allItems, Collections.emptySet(), userOriginated);
     }
 
     /**
+     * Fetch all items from the given hierarchical data provider.
+     *
+     * @since 8.1
+     * @param dataProvider
+     *            the data provider to fetch from
+     * @return all items in the data provider
+     */
+    private Stream<T> fetchAllHierarchical(
+            HierarchicalDataProvider<T, ?> dataProvider) {
+        return fetchAllDescendants(null, dataProvider);
+    }
+
+    /**
+     * Fetch all the descendants of the given parent item from the given data
+     * provider.
+     *
+     * @since 8.1
+     * @param parent
+     *            the parent item to fetch descendants for
+     * @param dataProvider
+     *            the data provider to fetch from
+     * @return the stream of all descendant items
+     */
+    private Stream<T> fetchAllDescendants(T parent,
+            HierarchicalDataProvider<T, ?> dataProvider) {
+        List<T> children = dataProvider
+                .fetchChildren(new HierarchicalQuery<>(null, parent))
+                .collect(Collectors.toList());
+        if (children.isEmpty()) {
+            return Stream.empty();
+        }
+        return children.stream()
+                .flatMap(child -> Stream.concat(Stream.of(child),
+                        fetchAllDescendants(child, dataProvider)));
+    }
+
+    /**
+     * Fetch all items from the given data provider.
+     *
+     * @since 8.1
+     * @param dataProvider
+     *            the data provider to fetch from
+     * @return all items in this data provider
+     */
+    private Stream<T> fetchAll(DataProvider<T, ?> dataProvider) {
+        return dataProvider.fetch(new Query<>());
+    }
+
+    /**
      * Triggered when the user unchecks the select all checkbox.
      *
      * @param userOriginated
diff --git a/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeatures.java b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeatures.java
index 911c13f449..0c29f71294 100644
--- a/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeatures.java
+++ b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeatures.java
@@ -17,6 +17,7 @@ import com.vaadin.server.SerializablePredicate;
 import com.vaadin.shared.Range;
 import com.vaadin.tests.components.AbstractComponentTest;
 import com.vaadin.tests.data.bean.HierarchicalTestBean;
+import com.vaadin.ui.Grid.SelectionMode;
 import com.vaadin.ui.ItemCollapseAllowedProvider;
 import com.vaadin.ui.TreeGrid;
 
@@ -68,6 +69,7 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
         createExpandMenu();
         createCollapseMenu();
         createListenerMenu();
+        createSelectionModeMenu();
     }
 
     private void initializeDataProviders() {
@@ -137,6 +139,7 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
                 (treeGrid, value, data) -> treeGrid.setHierarchyColumn(value));
     }
 
+    @SuppressWarnings("unchecked")
     private void createCollapseAllowedSelect() {
         LinkedHashMap<String, ItemCollapseAllowedProvider<HierarchicalTestBean>> options = new LinkedHashMap<>();
         options.put("all allowed", t -> true);
@@ -190,4 +193,14 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
                                 + event.isUserOriginated() + "): "
                                 + event.getExpandedItem())));
     }
+
+    private void createSelectionModeMenu() {
+        LinkedHashMap<String, SelectionMode> options = new LinkedHashMap<>();
+        options.put("none", SelectionMode.NONE);
+        options.put("single", SelectionMode.SINGLE);
+        options.put("multi", SelectionMode.MULTI);
+
+        createSelectAction("Selection mode", "State", options, "single",
+                (treeGrid, value, data) -> treeGrid.setSelectionMode(value));
+    }
 }
diff --git a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridSelectTest.java b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridSelectTest.java
new file mode 100644
index 0000000000..e167dc6ac5
--- /dev/null
+++ b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridSelectTest.java
@@ -0,0 +1,60 @@
+package com.vaadin.tests.components.treegrid;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import com.vaadin.testbench.By;
+import com.vaadin.testbench.elements.TreeGridElement;
+import com.vaadin.tests.tb3.SingleBrowserTest;
+
+public class TreeGridSelectTest extends SingleBrowserTest {
+
+    @Override
+    public Class<?> getUIClass() {
+        return TreeGridBasicFeatures.class;
+    }
+
+    @Test
+    public void select_and_deselect_all() {
+        openTestURL();
+
+        selectMenuPath("Component", "Features", "Set data provider",
+                "TreeDataProvider");
+        selectMenuPath("Component", "State", "Selection mode", "multi");
+
+        TreeGridElement grid = $(TreeGridElement.class).first();
+
+        assertAllRowsDeselected(grid);
+        clickSelectAll(grid);
+        assertAllRowsSelected(grid);
+        grid.expandWithClick(1, 1);
+        grid.expandWithClick(2, 1);
+        assertAllRowsSelected(grid);
+        clickSelectAll(grid);
+        assertAllRowsDeselected(grid);
+        clickSelectAll(grid);
+        grid.collapseWithClick(2, 1);
+        grid.expandWithClick(2, 1);
+        assertAllRowsSelected(grid);
+        grid.collapseWithClick(2, 1);
+        clickSelectAll(grid);
+        grid.expandWithClick(2, 1);
+        assertAllRowsDeselected(grid);
+    }
+
+    private void assertAllRowsSelected(TreeGridElement grid) {
+        for (int i = 0; i < grid.getRowCount(); i++) {
+            Assert.assertTrue(grid.getRow(i).isSelected());
+        }
+    }
+
+    private void assertAllRowsDeselected(TreeGridElement grid) {
+        for (int i = 0; i < grid.getRowCount(); i++) {
+            Assert.assertFalse(grid.getRow(i).isSelected());
+        }
+    }
+
+    private void clickSelectAll(TreeGridElement grid) {
+        grid.getHeaderCell(0, 0).findElement(By.tagName("input")).click();
+    }
+}