Include components in joined cells when finding child components (#8772)

Fixes #8722

5 files changed, 58 insertions, 46 deletions

diff --git a/server/src/main/java/com/vaadin/ui/Grid.java b/server/src/main/java/com/vaadin/ui/Grid.java
index ed74b4c664..5d03825c6f 100644
--- a/server/src/main/java/com/vaadin/ui/Grid.java
+++ b/server/src/main/java/com/vaadin/ui/Grid.java
@@ -97,7 +97,6 @@ import com.vaadin.ui.components.grid.DetailsGenerator;
 import com.vaadin.ui.components.grid.Editor;
 import com.vaadin.ui.components.grid.EditorImpl;
 import com.vaadin.ui.components.grid.Footer;
-import com.vaadin.ui.components.grid.FooterCell;
 import com.vaadin.ui.components.grid.FooterRow;
 import com.vaadin.ui.components.grid.GridSelectionModel;
 import com.vaadin.ui.components.grid.Header;
@@ -2406,28 +2405,23 @@ public class Grid<T> extends AbstractListing<T> implements HasComponents,
         return column;
     }
 
+    /**
+     * {@inheritDoc}
+     * <p>
+     * Note that the order of the returned components it not specified.
+     */
     @Override
     public Iterator<Component> iterator() {
         Set<Component> componentSet = new LinkedHashSet<>(extensionComponents);
         Header header = getHeader();
         for (int i = 0; i < header.getRowCount(); ++i) {
             HeaderRow row = header.getRow(i);
-            getColumns().forEach(column -> {
-                HeaderCell cell = row.getCell(column);
-                if (cell.getCellType() == GridStaticCellType.WIDGET) {
-                    componentSet.add(cell.getComponent());
-                }
-            });
+            componentSet.addAll(row.getComponents());
         }
         Footer footer = getFooter();
         for (int i = 0; i < footer.getRowCount(); ++i) {
             FooterRow row = footer.getRow(i);
-            getColumns().forEach(column -> {
-                FooterCell cell = row.getCell(column);
-                if (cell.getCellType() == GridStaticCellType.WIDGET) {
-                    componentSet.add(cell.getComponent());
-                }
-            });
+            componentSet.addAll(row.getComponents());
         }
         return Collections.unmodifiableSet(componentSet).iterator();
     }
diff --git a/server/src/main/java/com/vaadin/ui/components/grid/FooterRow.java b/server/src/main/java/com/vaadin/ui/components/grid/FooterRow.java
index 42e4650471..9a3c1aadb2 100644
--- a/server/src/main/java/com/vaadin/ui/components/grid/FooterRow.java
+++ b/server/src/main/java/com/vaadin/ui/components/grid/FooterRow.java
@@ -16,8 +16,10 @@
 package com.vaadin.ui.components.grid;
 
 import java.io.Serializable;
+import java.util.Collection;
 import java.util.Set;
 
+import com.vaadin.ui.Component;
 import com.vaadin.ui.Grid;
 import com.vaadin.ui.Grid.Column;
 
@@ -130,4 +132,14 @@ public interface FooterRow extends Serializable {
      *            the style name to set or null to not use any style name
      */
     public void setStyleName(String styleName);
+
+    /**
+     * Gets a collection of all components inside this row.
+     * <p>
+     * The order of the components in the returned collection is not specified.
+     *
+     * @return a collection of components in the row
+     */
+    public Collection<? extends Component> getComponents();
+
 }
diff --git a/server/src/main/java/com/vaadin/ui/components/grid/HeaderRow.java b/server/src/main/java/com/vaadin/ui/components/grid/HeaderRow.java
index f11843e147..d6c740e720 100644
--- a/server/src/main/java/com/vaadin/ui/components/grid/HeaderRow.java
+++ b/server/src/main/java/com/vaadin/ui/components/grid/HeaderRow.java
@@ -16,8 +16,10 @@
 package com.vaadin.ui.components.grid;
 
 import java.io.Serializable;
+import java.util.Collection;
 import java.util.Set;
 
+import com.vaadin.ui.Component;
 import com.vaadin.ui.Grid;
 import com.vaadin.ui.Grid.Column;
 
@@ -130,4 +132,13 @@ public interface HeaderRow extends Serializable {
      *            the style name to set or null to not use any style name
      */
     public void setStyleName(String styleName);
+
+    /**
+     * Gets a collection of all components inside this row.
+     * <p>
+     * The order of the components in the returned collection is not specified.
+     *
+     * @return a collection of components in the row
+     */
+    public Collection<? extends Component> getComponents();
 }
diff --git a/server/src/main/java/com/vaadin/ui/components/grid/StaticSection.java b/server/src/main/java/com/vaadin/ui/components/grid/StaticSection.java
index c53cc9c976..840fa7b115 100644
--- a/server/src/main/java/com/vaadin/ui/components/grid/StaticSection.java
+++ b/server/src/main/java/com/vaadin/ui/components/grid/StaticSection.java
@@ -17,6 +17,7 @@ package com.vaadin.ui.components.grid;
 
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Iterator;
@@ -405,6 +406,21 @@ public abstract class StaticSection<ROW extends StaticSection.StaticRow<?>>
         void addMergedCell(CELL newCell, Set<String> columnGroup) {
             rowState.cellGroups.put(newCell.getCellState(), columnGroup);
         }
+
+        public Collection<? extends Component> getComponents() {
+            List<Component> components = new ArrayList<>();
+            cells.forEach((id, cell) -> {
+                if (cell.getCellType() == GridStaticCellType.WIDGET) {
+                    components.add(cell.getComponent());
+                }
+            });
+            rowState.cellGroups.forEach((cellState, columnIds) -> {
+                if (cellState.connector != null) {
+                    components.add((Component) cellState.connector);
+                }
+            });
+            return components;
+        }
     }
 
     /**
@@ -534,7 +550,7 @@ public abstract class StaticSection<ROW extends StaticSection.StaticRow<?>>
         public GridStaticCellType getCellType() {
             return cellState.type;
         }
-        
+
         /**
          * Returns the custom style name for this cell.
          *
@@ -548,8 +564,7 @@ public abstract class StaticSection<ROW extends StaticSection.StaticRow<?>>
          * Sets a custom style name for this cell.
          *
          * @param styleName
-         *            the style name to set or null to not use any style
-         *            name
+         *            the style name to set or null to not use any style name
          */
         public void setStyleName(String styleName) {
             cellState.styleName = styleName;
diff --git a/compatibility-server/src/test/java/com/vaadin/v7/tests/server/component/grid/GridChildrenTest.java b/server/src/test/java/com/vaadin/tests/server/component/grid/GridChildrenTest.java
index 2d8edddac1..4959169cf8 100644
--- a/compatibility-server/src/test/java/com/vaadin/v7/tests/server/component/grid/GridChildrenTest.java
+++ b/server/src/test/java/com/vaadin/tests/server/component/grid/GridChildrenTest.java
@@ -13,7 +13,7 @@
  * License for the specific language governing permissions and limitations under
  * the License.
  */
-package com.vaadin.v7.tests.server.component.grid;
+package com.vaadin.tests.server.component.grid;
 
 import java.util.Iterator;
 
@@ -21,22 +21,23 @@ import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
+import com.vaadin.tests.data.bean.Person;
 import com.vaadin.ui.Component;
+import com.vaadin.ui.Grid;
 import com.vaadin.ui.Label;
-import com.vaadin.v7.ui.Grid;
-import com.vaadin.v7.ui.Grid.FooterCell;
-import com.vaadin.v7.ui.Grid.HeaderCell;
+import com.vaadin.ui.components.grid.FooterCell;
+import com.vaadin.ui.components.grid.HeaderCell;
 
 public class GridChildrenTest {
 
-    private Grid grid;
+    private Grid<Person> grid;
 
     @Before
     public void createGrid() {
-        grid = new Grid();
-        grid.addColumn("foo");
-        grid.addColumn("bar");
-        grid.addColumn("baz");
+        grid = new Grid<>();
+        grid.addColumn(Person::getFirstName).setId("foo");
+        grid.addColumn(Person::getLastName).setId("bar");
+        grid.addColumn(Person::getEmail).setId("baz");
 
     }
 
@@ -63,17 +64,6 @@ public class GridChildrenTest {
     }
 
     @Test
-    public void removeHeaderWithComponentInMergedHeaderCell() {
-        HeaderCell merged = grid.getDefaultHeaderRow().join("foo", "bar",
-                "baz");
-        Label label = new Label();
-        merged.setComponent(label);
-        Assert.assertEquals(grid, label.getParent());
-        grid.removeHeaderRow(0);
-        Assert.assertNull(label.getParent());
-    }
-
-    @Test
     public void removeComponentInMergedFooterCell() {
         FooterCell merged = grid.addFooterRowAt(0).join("foo", "bar", "baz");
         Label label = new Label();
@@ -84,16 +74,6 @@ public class GridChildrenTest {
     }
 
     @Test
-    public void removeFooterWithComponentInMergedFooterCell() {
-        FooterCell merged = grid.addFooterRowAt(0).join("foo", "bar", "baz");
-        Label label = new Label();
-        merged.setComponent(label);
-        Assert.assertEquals(grid, label.getParent());
-        grid.removeFooterRow(0);
-        Assert.assertNull(label.getParent());
-    }
-
-    @Test
     public void componentsInMergedFooter() {
         FooterCell merged = grid.addFooterRowAt(0).join("foo", "bar", "baz");
         Label label = new Label();