Allow changing the renderer of hierarchy column in TreeGrid (#9514)

Addresses #9465

8 files changed, 116 insertions, 85 deletions

diff --git a/client/src/main/java/com/vaadin/client/connectors/grid/ColumnConnector.java b/client/src/main/java/com/vaadin/client/connectors/grid/ColumnConnector.java
index 89914ea664..c47488f78e 100644
--- a/client/src/main/java/com/vaadin/client/connectors/grid/ColumnConnector.java
+++ b/client/src/main/java/com/vaadin/client/connectors/grid/ColumnConnector.java
@@ -36,7 +36,8 @@ import elemental.json.JsonValue;
 @Connect(com.vaadin.ui.Grid.Column.class)
 public class ColumnConnector extends AbstractExtensionConnector {
 
-    static abstract class CustomColumn extends Column<Object, JsonObject> {
+    public static abstract class CustomColumn
+            extends Column<Object, JsonObject> {
 
         private final String connectorId;
 
@@ -99,6 +100,7 @@ public class ColumnConnector extends AbstractExtensionConnector {
     @OnStateChange("renderer")
     void updateRenderer() {
         column.setRenderer(getRendererConnector().getRenderer());
+        getParent().onColumnRendererChanged(column);
     }
 
     @OnStateChange("hidingToggleCaption")
diff --git a/client/src/main/java/com/vaadin/client/connectors/grid/GridConnector.java b/client/src/main/java/com/vaadin/client/connectors/grid/GridConnector.java
index 661289ffb1..0e234cc968 100644
--- a/client/src/main/java/com/vaadin/client/connectors/grid/GridConnector.java
+++ b/client/src/main/java/com/vaadin/client/connectors/grid/GridConnector.java
@@ -457,6 +457,20 @@ public class GridConnector extends AbstractListingConnector
         idToColumn.remove(id);
     }
 
+    /**
+     * Method called by {@code CustomColumn} when its renderer changes. This
+     * method is used to maintain hierarchical renderer wrap in
+     * {@code TreeGrid}.
+     * 
+     * @param column
+     *            the column which now has a new renderer
+     * 
+     * @since 8.1
+     */
+    public void onColumnRendererChanged(CustomColumn column) {
+        // NO-OP
+    }
+
     @Override
     public void onUnregister() {
         super.onUnregister();
diff --git a/client/src/main/java/com/vaadin/client/ui/treegrid/TreeGridConnector.java b/client/src/main/java/com/vaadin/client/ui/treegrid/TreeGridConnector.java
index 71c286aedf..e92747dde9 100644
--- a/client/src/main/java/com/vaadin/client/ui/treegrid/TreeGridConnector.java
+++ b/client/src/main/java/com/vaadin/client/ui/treegrid/TreeGridConnector.java
@@ -18,6 +18,7 @@ package com.vaadin.client.ui.treegrid;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 import java.util.logging.Logger;
 
@@ -27,6 +28,7 @@ import com.google.gwt.dom.client.Element;
 import com.google.gwt.event.dom.client.KeyCodes;
 import com.google.gwt.user.client.Event;
 import com.vaadin.client.annotations.OnStateChange;
+import com.vaadin.client.connectors.grid.ColumnConnector.CustomColumn;
 import com.vaadin.client.connectors.grid.GridConnector;
 import com.vaadin.client.data.AbstractRemoteDataSource;
 import com.vaadin.client.data.DataChangeHandler;
@@ -229,7 +231,7 @@ public class TreeGridConnector extends GridConnector {
             GridEventHandler<?> eventHandler)
     /*-{
         var browserEventHandlers = grid.@com.vaadin.client.widgets.Grid::browserEventHandlers;
-
+    
         // FocusEventHandler is initially 5th in the list of browser event handlers
         browserEventHandlers.@java.util.List::set(*)(5, eventHandler);
     }-*/;
@@ -352,7 +354,8 @@ public class TreeGridConnector extends GridConnector {
                 // Hierarchy metadata
                 JsonObject rowData = cell.getRow();
                 if (rowData == null) {
-                    // Row data is lost from the cache, i.e. the row is at least outside the visual area,
+                    // Row data is lost from the cache, i.e. the row is at least
+                    // outside the visual area,
                     // let's scroll the row into the view
                     getWidget().scrollToRow(cell.getRowIndex());
                 } else if (rowData.hasKey(
@@ -363,21 +366,21 @@ public class TreeGridConnector extends GridConnector {
                             HierarchicalDataCommunicatorConstants.ROW_LEAF);
                     boolean collapsed = isCollapsed(rowData);
                     switch (domEvent.getKeyCode()) {
-                        case KeyCodes.KEY_RIGHT:
-                            if (collapsed && !leaf) {
-                                setCollapsed(cell.getRowIndex(), false);
-                            }
-                            break;
-                        case KeyCodes.KEY_LEFT:
-                            if (collapsed || leaf) {
-                                // navigate up
-                                int columnIndex = cell.getColumnIndex();
-                                getRpcProxy(FocusParentRpc.class).focusParent(
-                                        cell.getRowIndex(), columnIndex);
-                            } else if (isCollapseAllowed(rowDescription)) {
-                                setCollapsed(cell.getRowIndex(), true);
-                            }
-                            break;
+                    case KeyCodes.KEY_RIGHT:
+                        if (collapsed && !leaf) {
+                            setCollapsed(cell.getRowIndex(), false);
+                        }
+                        break;
+                    case KeyCodes.KEY_LEFT:
+                        if (collapsed || leaf) {
+                            // navigate up
+                            int columnIndex = cell.getColumnIndex();
+                            getRpcProxy(FocusParentRpc.class).focusParent(
+                                    cell.getRowIndex(), columnIndex);
+                        } else if (isCollapseAllowed(rowDescription)) {
+                            setCollapsed(cell.getRowIndex(), true);
+                        }
+                        break;
                     }
 
                 }
@@ -419,17 +422,29 @@ public class TreeGridConnector extends GridConnector {
         return rowData
                 .getObject(
                         HierarchicalDataCommunicatorConstants.ROW_HIERARCHY_DESCRIPTION)
-                .getBoolean(HierarchicalDataCommunicatorConstants.ROW_COLLAPSED);
+                .getBoolean(
+                        HierarchicalDataCommunicatorConstants.ROW_COLLAPSED);
     }
 
     /**
-     * Checks if the item can be collapsed
+     * Checks if the item can be collapsed.
      *
-     * @param row the item row
-     * @return {@code true} if the item is allowed to be collapsed, {@code false} otherwise.
+     * @param row
+     *            the item row
+     * @return {@code true} if the item is allowed to be collapsed,
+     *         {@code false} otherwise.
      */
     public static boolean isCollapseAllowed(JsonObject row) {
         return row.getBoolean(
                 HierarchicalDataCommunicatorConstants.ROW_COLLAPSE_ALLOWED);
     }
+
+    @Override
+    public void onColumnRendererChanged(CustomColumn column) {
+        super.onColumnRendererChanged(column);
+
+        if (Objects.equals(getColumnId(column), hierarchyColumnId)) {
+            updateHierarchyColumn();
+        }
+    }
 }
diff --git a/server/src/main/java/com/vaadin/ui/Grid.java b/server/src/main/java/com/vaadin/ui/Grid.java
index 3ae03679cf..7e283a38ea 100644
--- a/server/src/main/java/com/vaadin/ui/Grid.java
+++ b/server/src/main/java/com/vaadin/ui/Grid.java
@@ -1906,6 +1906,16 @@ public class Grid<T> extends AbstractListing<T> implements HasComponents,
         }
 
         /**
+         * Gets the Renderer for this Column.
+         * 
+         * @return the renderer
+         * @since
+         */
+        public Renderer<? super V> getRenderer() {
+            return (Renderer<? super V>) getState().renderer;
+        }
+
+        /**
          * Gets the grid that this column belongs to.
          *
          * @return the grid that this column belongs to, or <code>null</code> if
@@ -4006,7 +4016,7 @@ public class Grid<T> extends AbstractListing<T> implements HasComponents,
                 column = addColumn(id);
             } else {
                 DeclarativeValueProvider<T> provider = new DeclarativeValueProvider<>();
-                column = new Column<>(provider, new HtmlRenderer());
+                column = createColumn(provider, new HtmlRenderer());
                 addColumn(getGeneratedIdentifier(), column);
                 if (id != null) {
                     column.setId(id);
diff --git a/server/src/main/java/com/vaadin/ui/TreeGrid.java b/server/src/main/java/com/vaadin/ui/TreeGrid.java
index 310a7c6cf5..fdda7226fe 100644
--- a/server/src/main/java/com/vaadin/ui/TreeGrid.java
+++ b/server/src/main/java/com/vaadin/ui/TreeGrid.java
@@ -486,26 +486,6 @@ public class TreeGrid<T> extends Grid<T>
                 childItem -> writeRow(container, childItem, item, context));
     }
 
-    @Override
-    protected <V> Column<T, V> createColumn(ValueProvider<T, V> valueProvider,
-            AbstractRenderer<? super T, ? super V> renderer) {
-        return new Column<T, V>(valueProvider, renderer) {
-
-            @Override
-            public com.vaadin.ui.Grid.Column<T, V> setRenderer(
-                    Renderer<? super V> renderer) {
-                // Disallow changing renderer for the hierarchy column
-                if (getInternalIdForColumn(this).equals(
-                        TreeGrid.this.getState(false).hierarchyColumnId)) {
-                    throw new IllegalStateException(
-                            "Changing the renderer of the hierarchy column is not allowed.");
-                }
-
-                return super.setRenderer(renderer);
-            }
-        };
-    }
-
     /**
      * Emit an expand event.
      *
diff --git a/server/src/test/java/com/vaadin/tests/components/treegrid/TreeGridTest.java b/server/src/test/java/com/vaadin/tests/components/treegrid/TreeGridTest.java
index 1544a4f69d..1f900688c4 100644
--- a/server/src/test/java/com/vaadin/tests/components/treegrid/TreeGridTest.java
+++ b/server/src/test/java/com/vaadin/tests/components/treegrid/TreeGridTest.java
@@ -14,11 +14,10 @@ public class TreeGridTest {
     private boolean expandEventFired = false;
     private boolean collapseEventFired = false;
 
-    @Test(expected = IllegalStateException.class)
     public void testChangeRendererOfHierarchyColumn() {
         treeGrid.addColumn(Object::toString).setId("foo");
         treeGrid.setHierarchyColumn("foo");
-        // This should not be allowed.
+        // This should be allowed.
         treeGrid.getColumn("foo").setRenderer(new TextRenderer());
     }
 
@@ -28,8 +27,7 @@ public class TreeGridTest {
         treeData.addItem(null, "Foo");
         treeData.addItem("Foo", "Bar");
         treeData.addItem("Foo", "Baz");
-        treeGrid.setDataProvider(
-                new TreeDataProvider<>(treeData));
+        treeGrid.setDataProvider(new TreeDataProvider<>(treeData));
 
         treeGrid.addExpandListener(e -> expandEventFired = true);
         treeGrid.addCollapseListener(e -> collapseEventFired = true);
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 0c29f71294..2ac0c0673d 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
@@ -4,6 +4,7 @@ import java.util.Arrays;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Optional;
+import java.util.function.Consumer;
 import java.util.stream.Stream;
 
 import com.vaadin.annotations.Theme;
@@ -17,9 +18,12 @@ 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.Column;
 import com.vaadin.ui.Grid.SelectionMode;
 import com.vaadin.ui.ItemCollapseAllowedProvider;
 import com.vaadin.ui.TreeGrid;
+import com.vaadin.ui.renderers.HtmlRenderer;
+import com.vaadin.ui.renderers.TextRenderer;
 
 @Theme("valo")
 @Widgetset("com.vaadin.DefaultWidgetSet")
@@ -29,6 +33,7 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
     private TreeDataProvider<HierarchicalTestBean> inMemoryDataProvider;
     private LazyHierarchicalDataProvider lazyDataProvider;
     private HierarchicalDataProvider<HierarchicalTestBean, ?> loggingDataProvider;
+    private Column<HierarchicalTestBean, String> hierarchyColumn;
 
     @Override
     public TreeGrid getComponent() {
@@ -45,8 +50,9 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
         initializeDataProviders();
         grid = new TreeGrid<>();
         grid.setSizeFull();
-        grid.addColumn(HierarchicalTestBean::toString).setCaption("String")
-                .setId("string");
+        hierarchyColumn = grid.addColumn(HierarchicalTestBean::toString);
+        hierarchyColumn.setCaption("String").setId("string").setStyleGenerator(
+                t -> hierarchyColumn.getRenderer().getClass().getSimpleName());
         grid.addColumn(HierarchicalTestBean::getDepth).setCaption("Depth")
                 .setId("depth").setDescriptionGenerator(
                         t -> "Hierarchy depth: " + t.getDepth());
@@ -66,6 +72,7 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
         createDataProviderSelect();
         createHierarchyColumnSelect();
         createCollapseAllowedSelect();
+        createRendererSelect();
         createExpandMenu();
         createCollapseMenu();
         createListenerMenu();
@@ -94,8 +101,7 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
 
         inMemoryDataProvider = new TreeDataProvider<>(data);
         lazyDataProvider = new LazyHierarchicalDataProvider(3, 2);
-        loggingDataProvider = new TreeDataProvider<HierarchicalTestBean>(
-                data) {
+        loggingDataProvider = new TreeDataProvider<HierarchicalTestBean>(data) {
 
             @Override
             public Stream<HierarchicalTestBean> fetchChildren(
@@ -152,6 +158,16 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
                         .setItemCollapseAllowedProvider(value));
     }
 
+    private void createRendererSelect() {
+        LinkedHashMap<String, Consumer<Column<?, String>>> options = new LinkedHashMap<>();
+        options.put("text", c -> c.setRenderer(new TextRenderer()));
+        options.put("html", c -> c.setRenderer(new HtmlRenderer()));
+
+        createSelectAction("Hierarchy column renderer", CATEGORY_FEATURES,
+                options, "text",
+                (treeGrid, consumer, data) -> consumer.accept(hierarchyColumn));
+    }
+
     @SuppressWarnings("unchecked")
     private void createExpandMenu() {
         createCategory("Server-side expand", CATEGORY_FEATURES);
diff --git a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeaturesTest.java b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeaturesTest.java
index c73418675a..43c08415df 100644
--- a/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeaturesTest.java
+++ b/uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeaturesTest.java
@@ -136,58 +136,48 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
         new Actions(getDriver()).sendKeys(Keys.RIGHT).perform();
         assertEquals(6, grid.getRowCount());
         assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "1 | 2" });
-        assertTrue(
-                grid.getRow(0).hasClassName("v-treegrid-row-focused"));
-        assertFalse(
-                grid.getRow(1).hasClassName("v-treegrid-row-focused"));
+        assertTrue(grid.getRow(0).hasClassName("v-treegrid-row-focused"));
+        assertFalse(grid.getRow(1).hasClassName("v-treegrid-row-focused"));
 
         // Should navigate 2 times down to "1 | 1"
         new Actions(getDriver()).sendKeys(Keys.DOWN, Keys.DOWN).perform();
         assertEquals(6, grid.getRowCount());
         assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "1 | 2" });
-        assertFalse(
-                grid.getRow(0).hasClassName("v-treegrid-row-focused"));
-        assertFalse(
-                grid.getRow(1).hasClassName("v-treegrid-row-focused"));
-        assertTrue(
-                grid.getRow(2).hasClassName("v-treegrid-row-focused"));
+        assertFalse(grid.getRow(0).hasClassName("v-treegrid-row-focused"));
+        assertFalse(grid.getRow(1).hasClassName("v-treegrid-row-focused"));
+        assertTrue(grid.getRow(2).hasClassName("v-treegrid-row-focused"));
 
         // Should expand "1 | 1" without moving focus
         new Actions(getDriver()).sendKeys(Keys.RIGHT).perform();
         assertEquals(9, grid.getRowCount());
         assertCellTexts(2, 0,
                 new String[] { "1 | 1", "2 | 0", "2 | 1", "2 | 2", "1 | 2" });
-        assertTrue(
-                grid.getRow(2).hasClassName("v-treegrid-row-focused"));
+        assertTrue(grid.getRow(2).hasClassName("v-treegrid-row-focused"));
 
         // Should collapse "1 | 1"
         new Actions(getDriver()).sendKeys(Keys.LEFT).perform();
         assertEquals(6, grid.getRowCount());
         assertCellTexts(2, 0, new String[] { "1 | 1", "1 | 2", "0 | 1" });
-        assertTrue(
-                grid.getRow(2).hasClassName("v-treegrid-row-focused"));
+        assertTrue(grid.getRow(2).hasClassName("v-treegrid-row-focused"));
 
         // Should navigate to "0 | 0"
         new Actions(getDriver()).sendKeys(Keys.LEFT).perform();
         assertEquals(6, grid.getRowCount());
         assertCellTexts(0, 0,
                 new String[] { "0 | 0", "1 | 0", "1 | 1", "1 | 2", "0 | 1" });
-        assertTrue(
-                grid.getRow(0).hasClassName("v-treegrid-row-focused"));
+        assertTrue(grid.getRow(0).hasClassName("v-treegrid-row-focused"));
 
         // Should collapse "0 | 0"
         new Actions(getDriver()).sendKeys(Keys.LEFT).perform();
         assertEquals(3, grid.getRowCount());
         assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
-        assertTrue(
-                grid.getRow(0).hasClassName("v-treegrid-row-focused"));
+        assertTrue(grid.getRow(0).hasClassName("v-treegrid-row-focused"));
 
         // Nothing should happen
         new Actions(getDriver()).sendKeys(Keys.LEFT).perform();
         assertEquals(3, grid.getRowCount());
         assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
-        assertTrue(
-                grid.getRow(0).hasClassName("v-treegrid-row-focused"));
+        assertTrue(grid.getRow(0).hasClassName("v-treegrid-row-focused"));
 
         assertNoErrorNotifications();
     }
@@ -205,12 +195,9 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
         new Actions(getDriver()).sendKeys(Keys.DOWN, Keys.DOWN).perform();
         assertEquals(6, grid.getRowCount());
         assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "1 | 2" });
-        assertFalse(
-                grid.getRow(0).hasClassName("v-treegrid-row-focused"));
-        assertFalse(
-                grid.getRow(1).hasClassName("v-treegrid-row-focused"));
-        assertTrue(
-                grid.getRow(2).hasClassName("v-treegrid-row-focused"));
+        assertFalse(grid.getRow(0).hasClassName("v-treegrid-row-focused"));
+        assertFalse(grid.getRow(1).hasClassName("v-treegrid-row-focused"));
+        assertTrue(grid.getRow(2).hasClassName("v-treegrid-row-focused"));
 
         // Should select "1 | 1" without moving focus
         new Actions(getDriver()).sendKeys(Keys.SPACE).perform();
@@ -218,19 +205,15 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
 
         // Should move focus but not selection
         new Actions(getDriver()).sendKeys(Keys.UP).perform();
-        assertTrue(
-                grid.getRow(1).hasClassName("v-treegrid-row-focused"));
-        assertFalse(
-                grid.getRow(2).hasClassName("v-treegrid-row-focused"));
+        assertTrue(grid.getRow(1).hasClassName("v-treegrid-row-focused"));
+        assertFalse(grid.getRow(2).hasClassName("v-treegrid-row-focused"));
         assertFalse(grid.getRow(1).hasClassName("v-treegrid-row-selected"));
         assertTrue(grid.getRow(2).hasClassName("v-treegrid-row-selected"));
 
         // Should select "1 | 0" without moving focus
         new Actions(getDriver()).sendKeys(Keys.SPACE).perform();
-        assertTrue(
-                grid.getRow(1).hasClassName("v-treegrid-row-focused"));
-        assertFalse(
-                grid.getRow(2).hasClassName("v-treegrid-row-focused"));
+        assertTrue(grid.getRow(1).hasClassName("v-treegrid-row-focused"));
+        assertFalse(grid.getRow(2).hasClassName("v-treegrid-row-focused"));
         assertTrue(grid.getRow(1).hasClassName("v-treegrid-row-selected"));
         assertFalse(grid.getRow(2).hasClassName("v-treegrid-row-selected"));
 
@@ -339,6 +322,19 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
         assertEquals(9, grid.getRowCount());
     }
 
+    @Test
+    public void change_renderer_of_hierarchy_column() {
+        assertTrue("Cell style names should contain renderer name", grid
+                .getCell(0, 0).getAttribute("class").contains("TextRenderer"));
+        selectMenuPath("Component", "Features", "Hierarchy column renderer",
+                "html");
+        assertTrue("Cell style names should contain renderer name", grid
+                .getCell(0, 0).getAttribute("class").contains("HtmlRenderer"));
+
+        grid.expandWithClick(0);
+        assertEquals("Not expanded", "1 | 0", grid.getCell(1, 0).getText());
+    }
+
     private void assertCellTexts(int startRowIndex, int cellIndex,
             String[] cellTexts) {
         int index = startRowIndex;