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

Fixes #8722

5 files changed, 134 insertions, 16 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/server/src/test/java/com/vaadin/tests/server/component/grid/GridChildrenTest.java b/server/src/test/java/com/vaadin/tests/server/component/grid/GridChildrenTest.java
new file mode 100644
index 0000000000..4959169cf8
--- /dev/null
+++ b/server/src/test/java/com/vaadin/tests/server/component/grid/GridChildrenTest.java
@@ -0,0 +1,86 @@
+/*
+ * Copyright 2000-2016 Vaadin Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package com.vaadin.tests.server.component.grid;
+
+import java.util.Iterator;
+
+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.ui.components.grid.FooterCell;
+import com.vaadin.ui.components.grid.HeaderCell;
+
+public class GridChildrenTest {
+
+    private Grid<Person> grid;
+
+    @Before
+    public void createGrid() {
+        grid = new Grid<>();
+        grid.addColumn(Person::getFirstName).setId("foo");
+        grid.addColumn(Person::getLastName).setId("bar");
+        grid.addColumn(Person::getEmail).setId("baz");
+
+    }
+
+    @Test
+    public void iteratorFindsComponentsInMergedHeader() {
+        HeaderCell merged = grid.getDefaultHeaderRow().join("foo", "bar",
+                "baz");
+        Label label = new Label();
+        merged.setComponent(label);
+        Iterator<Component> i = grid.iterator();
+        Assert.assertEquals(label, i.next());
+        Assert.assertFalse(i.hasNext());
+    }
+
+    @Test
+    public void removeComponentInMergedHeaderCell() {
+        HeaderCell merged = grid.getDefaultHeaderRow().join("foo", "bar",
+                "baz");
+        Label label = new Label();
+        merged.setComponent(label);
+        Assert.assertEquals(grid, label.getParent());
+        merged.setText("foo");
+        Assert.assertNull(label.getParent());
+    }
+
+    @Test
+    public void removeComponentInMergedFooterCell() {
+        FooterCell merged = grid.addFooterRowAt(0).join("foo", "bar", "baz");
+        Label label = new Label();
+        merged.setComponent(label);
+        Assert.assertEquals(grid, label.getParent());
+        merged.setText("foo");
+        Assert.assertNull(label.getParent());
+    }
+
+    @Test
+    public void componentsInMergedFooter() {
+        FooterCell merged = grid.addFooterRowAt(0).join("foo", "bar", "baz");
+        Label label = new Label();
+        merged.setComponent(label);
+        Iterator<Component> i = grid.iterator();
+        Assert.assertEquals(label, i.next());
+        Assert.assertFalse(i.hasNext());
+        Assert.assertEquals(grid, label.getParent());
+    }
+}