Make select all work with TreeGrid (#9412)

Fixes #9403

2 files changed, 73 insertions, 0 deletions

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();
+    }
+}