Fix Grid not removing columns from sort order

Fixes #10928

2 files changed, 31 insertions, 4 deletions

diff --git a/server/src/main/java/com/vaadin/ui/Grid.java b/server/src/main/java/com/vaadin/ui/Grid.java
index 3bf65c6c9b..956b81c3eb 100644
--- a/server/src/main/java/com/vaadin/ui/Grid.java
+++ b/server/src/main/java/com/vaadin/ui/Grid.java
@@ -2894,6 +2894,14 @@ public class Grid<T> extends AbstractListing<T> implements HasComponents,
             getFooter().removeColumn(columnId);
             getState(true).columnOrder.remove(columnId);
 
+            // Remove column from sorted columns.
+            List<GridSortOrder<T>> filteredSortOrder = sortOrder.stream()
+                    .filter(order -> !order.getSorted().equals(column))
+                    .collect(Collectors.toList());
+            if (filteredSortOrder.size() < sortOrder.size()) {
+                setSortOrder(filteredSortOrder);
+            }
+
             if (displayIndex < getFrozenColumnCount()) {
                 setFrozenColumnCount(getFrozenColumnCount() - 1);
             }
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);
     }