2 files changed, 31 insertions, 2 deletions
diff --git a/uitest/src/main/java/com/vaadin/tests/components/grid/GridAddColumn.java b/uitest/src/main/java/com/vaadin/tests/components/grid/GridAddColumn.java
index e771b9cc78..8846091512 100644
--- a/uitest/src/main/java/com/vaadin/tests/components/grid/GridAddColumn.java
+++ b/uitest/src/main/java/com/vaadin/tests/components/grid/GridAddColumn.java
@@ -1,18 +1,34 @@
 package com.vaadin.tests.components.grid;
 
+import com.vaadin.annotations.Widgetset;
 import com.vaadin.data.ValueProvider;
 import com.vaadin.server.VaadinRequest;
 import com.vaadin.tests.components.AbstractTestUI;
+import com.vaadin.ui.Button;
 import com.vaadin.ui.Grid;
 import com.vaadin.ui.Grid.Column;
 import com.vaadin.ui.Label;
+import com.vaadin.ui.Notification;
 import com.vaadin.ui.renderers.NumberRenderer;
 
+@Widgetset("com.vaadin.DefaultWidgetSet")
 public class GridAddColumn extends AbstractTestUI {
 
     @Override
     protected void setup(VaadinRequest request) {
         Grid<String> grid = new Grid<>();
+        createColumns(grid);
+        grid.setItems("a", "aa", "aaa");
+        addComponent(grid);
+
+        addComponent(new Button("Replace columns", e -> {
+            grid.removeAllColumns();
+            createColumns(grid);
+            Notification.show("Columns replaced.");
+        }));
+    }
+
+    private void createColumns(Grid<String> grid) {
         Column<String, String> col0 = grid.addColumn(ValueProvider.identity())
                 .setCaption("First column");
         grid.getDefaultHeaderRow().getCell(col0)
@@ -21,7 +37,5 @@ public class GridAddColumn extends AbstractTestUI {
         grid.addColumn(String::length);
         grid.addColumn(string -> -string.length());
         grid.addColumn(string -> new Object());
-        grid.setItems("a", "aa", "aaa");
-        addComponent(grid);
     }
 }
diff --git a/uitest/src/test/java/com/vaadin/tests/components/grid/GridAddColumnTest.java b/uitest/src/test/java/com/vaadin/tests/components/grid/GridAddColumnTest.java
index ca5ad8ec50..e3595e7674 100644
--- a/uitest/src/test/java/com/vaadin/tests/components/grid/GridAddColumnTest.java
+++ b/uitest/src/test/java/com/vaadin/tests/components/grid/GridAddColumnTest.java
@@ -5,8 +5,10 @@ import org.junit.Before;
 import org.junit.Test;
 
 import com.vaadin.testbench.By;
+import com.vaadin.testbench.elements.ButtonElement;
 import com.vaadin.testbench.elements.GridElement;
 import com.vaadin.testbench.elements.GridElement.GridCellElement;
+import com.vaadin.testbench.elements.NotificationElement;
 import com.vaadin.tests.tb3.SingleBrowserTest;
 
 public class GridAddColumnTest extends SingleBrowserTest {
@@ -64,6 +66,19 @@ public class GridAddColumnTest extends SingleBrowserTest {
                 firstHeader.getText());
     }
 
+    @Test
+    public void replace_all_columns() {
+        $(ButtonElement.class).first().click();
+
+        // Verify button got clicked
+        Assert.assertTrue(isElementPresent(NotificationElement.class));
+        Assert.assertEquals("Columns replaced.",
+                $(NotificationElement.class).first().getText());
+
+        // Run default rendering test
+        columns_rendered_correctly();
+    }
+
     private void assertCellEquals(int rowIndex, int colIndex, String content) {
         Assert.assertEquals("Cell text should equal", content,
                 grid.getCell(rowIndex, colIndex).getText());