2 files changed, 13 insertions, 1 deletions
diff --git a/server/src/main/java/com/vaadin/ui/TreeGrid.java b/server/src/main/java/com/vaadin/ui/TreeGrid.java
index 703a83192a..8ee154a56e 100644
--- a/server/src/main/java/com/vaadin/ui/TreeGrid.java
+++ b/server/src/main/java/com/vaadin/ui/TreeGrid.java
@@ -248,7 +248,7 @@ public class TreeGrid<T> extends Grid<T>
      *         has been explicitly set
      */
     public Column<T, ?> getHierarchyColumn() {
-        return getColumn(getState(false).hierarchyColumnId);
+        return getColumnByInternalId(getState(false).hierarchyColumnId);
     }
 
     /**
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 fdaacae48b..1876d8de35 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
@@ -5,6 +5,7 @@ import org.junit.Test;
 
 import com.vaadin.data.TreeData;
 import com.vaadin.data.provider.TreeDataProvider;
+import com.vaadin.ui.Grid.Column;
 import com.vaadin.ui.TreeGrid;
 import com.vaadin.ui.renderers.TextRenderer;
 
@@ -14,6 +15,7 @@ public class TreeGridTest {
     private boolean expandEventFired = false;
     private boolean collapseEventFired = false;
 
+    @Test
     public void testChangeRendererOfHierarchyColumn() {
         treeGrid.addColumn(Object::toString).setId("foo");
         treeGrid.setHierarchyColumn("foo");
@@ -44,4 +46,14 @@ public class TreeGridTest {
         Assert.assertFalse("Item not collapsed", treeGrid.isExpanded("Foo"));
         Assert.assertTrue("Collapse event not fired", collapseEventFired);
     }
+
+    @Test
+    public void testSetAndGetHierarchyColumn() {
+        Column<String, String> column = treeGrid.addColumn(Object::toString)
+                .setId("foo");
+        treeGrid.setHierarchyColumn("foo");
+        Assert.assertEquals("Hierarchy column was not correctly returned",
+                column, treeGrid.getHierarchyColumn());
+    }
+
 }