1 files changed, 23 insertions, 4 deletions
diff --git a/server/src/test/java/com/vaadin/tests/server/component/grid/GridTest.java b/server/src/test/java/com/vaadin/tests/server/component/grid/GridTest.java
index 04d7d338a7..6f1a227a1c 100644
--- a/server/src/test/java/com/vaadin/tests/server/component/grid/GridTest.java
+++ b/server/src/test/java/com/vaadin/tests/server/component/grid/GridTest.java
@@ -45,6 +45,7 @@ import com.vaadin.data.provider.QuerySortOrder;
 import com.vaadin.data.provider.bov.Person;
 import com.vaadin.event.selection.SelectionEvent;
 import com.vaadin.server.SerializableComparator;
+import com.vaadin.shared.communication.SharedState;
 import com.vaadin.shared.data.sort.SortDirection;
 import com.vaadin.shared.ui.grid.GridState;
 import com.vaadin.shared.ui.grid.HeightMode;
@@ -66,13 +67,18 @@ public class GridTest {
     private Column<String, Integer> lengthColumn;
     private Column<String, Object> objectColumn;
     private Column<String, String> randomColumn;
+    private GridState state;
 
     @Rule
     public ExpectedException thrown = ExpectedException.none();
 
     @Before
     public void setUp() {
-        grid = new Grid<>();
+        grid = new Grid<String>() {
+            {
+                state = getState(false);
+            }
+        };
 
         fooColumn = grid.addColumn(ValueProvider.identity()).setId("foo");
         lengthColumn = grid.addColumn(String::length, new NumberRenderer())
@@ -295,9 +301,22 @@ public class GridTest {
         assertEquals(0, grid.getSortOrder().size());
 
         // Make sure state is updated.
-        Method stateMethod = grid.getClass().getDeclaredMethod("getState");
-        stateMethod.setAccessible(true);
-        GridState state = (GridState) stateMethod.invoke(grid);
+        assertEquals(0, state.sortColumns.length);
+        assertEquals(0, state.sortDirs.length);
+    }
+
+    @Test
+    public void sortOrderDoesnotContainRemovedColumns() {
+        Column<String, ?> sortColumn = grid.getColumns().get(1);
+        grid.sort(sortColumn);
+
+        // Get id of column and check it's sorted.
+        String id = state.columnOrder.get(1);
+        assertEquals(id, state.sortColumns[0]);
+
+        // Remove column and make sure it's cleared correctly
+        grid.removeColumn(sortColumn);
+        assertFalse("Column not removed", state.columnOrder.contains(id));
         assertEquals(0, state.sortColumns.length);
         assertEquals(0, state.sortDirs.length);
     }