Fix rendering of TreeGrid's frozen columns after hierarchy-column reset (#12028)

* Add tests

* Fix getVisibleFrozenColumnCount() if SelectionMode is multi

2 files changed, 115 insertions, 0 deletions

diff --git a/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridChangeHierarchyColumn.java b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridChangeHierarchyColumn.java
new file mode 100644
index 0000000000..0e6e650e78
--- /dev/null
+++ b/uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridChangeHierarchyColumn.java
@@ -0,0 +1,75 @@
+package com.vaadin.tests.components.treegrid;
+
+import com.vaadin.data.TreeData;
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.tests.components.AbstractTestUI;
+import com.vaadin.ui.Button;
+import com.vaadin.ui.Component;
+import com.vaadin.ui.CssLayout;
+import com.vaadin.ui.Grid;
+import com.vaadin.ui.Label;
+import com.vaadin.ui.TreeGrid;
+
+public class TreeGridChangeHierarchyColumn extends AbstractTestUI {
+
+    @Override
+    protected String getTestDescription() {
+        return "TreeGrid in MultiSelect mode should take hiden columns into account when"
+                + " rendering frozen columns after hierarchy-column reset.";
+    }
+
+    @Override
+    protected Integer getTicketNumber() {
+        return 12026;
+    }
+
+    @Override
+    protected void setup(VaadinRequest request) {
+        TreeGrid<String> treeGrid = new TreeGrid<>();
+        treeGrid.setId("TreeGrid");
+
+        for (int i = 0; i < 20; i++) {
+            String columnId = String.valueOf(i);
+            Grid.Column<String, Component> column = addColumn(treeGrid,
+                    columnId);
+            column.setCaption(columnId);
+            column.setId(columnId);
+        }
+
+        TreeData<String> data = treeGrid.getTreeData();
+        data.addItem(null, "child");
+        data.addItem("child", "grandChild");
+
+        treeGrid.setHierarchyColumn(treeGrid.getColumns().get(0));
+
+        Button hideHierCol = new Button("Hide Hierarchy Column");
+        hideHierCol.addClickListener(e -> {
+            treeGrid.getHierarchyColumn().setHidden(true);
+        });
+        hideHierCol.setId("hideHierColButton");
+
+        Button setHierCol = new Button("Set new Hierarchy Column");
+        setHierCol.addClickListener(e -> {
+            treeGrid.getColumns().stream().filter(column -> !column.isHidden())
+                    .findFirst().ifPresent(col -> {
+                        treeGrid.setHierarchyColumn(col.getId());
+                    });
+        });
+        setHierCol.setId("setHierColButton");
+
+        treeGrid.setSelectionMode(Grid.SelectionMode.MULTI);
+        treeGrid.setFrozenColumnCount(1);
+
+        addComponents(treeGrid, hideHierCol, setHierCol);
+    }
+
+    private Grid.Column<String, Component> addColumn(Grid<String> grid,
+            String columnId) {
+        return grid.addComponentColumn(val -> {
+            Label label = new Label(columnId);
+            label.setWidth(50, Unit.PIXELS);
+            return new CssLayout(label);
+        });
+    }
+
+}
diff --git a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridChangeHierarchyColumnTest.java b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridChangeHierarchyColumnTest.java
new file mode 100644
index 0000000000..77efba9d37
--- /dev/null
+++ b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridChangeHierarchyColumnTest.java
@@ -0,0 +1,40 @@
+package com.vaadin.tests.components.treegrid;
+
+import static org.junit.Assert.assertEquals;
+import java.util.List;
+
+import org.junit.Test;
+import org.openqa.selenium.WebElement;
+import com.vaadin.testbench.By;
+import com.vaadin.testbench.elements.ButtonElement;
+import com.vaadin.testbench.elements.TreeGridElement;
+import com.vaadin.tests.tb3.MultiBrowserTest;
+
+public class TreeGridChangeHierarchyColumnTest extends MultiBrowserTest {
+
+    @Test
+    public void renderingFrozenColumnsShouldFactorInHiddenColumns() {
+        openTestURL();
+        waitForElementPresent(By.id("TreeGrid"));
+        waitForElementPresent(By.id("hideHierColButton"));
+        waitForElementPresent(By.id("setHierColButton"));
+
+        TreeGridElement treeGrid = $(TreeGridElement.class).id("TreeGrid");
+        ButtonElement hideHierCol = $(ButtonElement.class)
+                .id("hideHierColButton");
+        ButtonElement setHierCol = $(ButtonElement.class)
+                .id("setHierColButton");
+
+        hideHierCol.click();
+        setHierCol.click();
+
+        // Wait for the new hierarchy column to be rendered
+        waitForElementPresent(By.className("v-treegrid-expander"));
+
+        List<WebElement> frozenCells = treeGrid
+                .findElements(By.className("frozen"));
+
+        assertEquals("Only the MultiSelect column should have frozen cells.", 2,
+                frozenCells.size());
+    }
+}