Add getHeaderCellByCaption method to tb-api GridElement (#8248)

* Add getHeaderCellByCaption method to tb-api GridElement

testbench-api/src/main/java/com/vaadin/testbench/elements/GridElement.java

@@ -29,8 +29,7 @@ import com.vaadin.testbench.elementsbase.ServerClass;
 
 /**
  * TestBench Element API for Grid
- * 
- * @since
+ *
  * @author Vaadin Ltd
  */
 @ServerClass("com.vaadin.ui.Grid")
@@ -65,8 +64,8 @@ public class GridElement extends AbstractComponentElement {
         }
 
         public GridCellElement getCell(int columnIndex) {
-            TestBenchElement e = (TestBenchElement) findElement(
-                    By.xpath("./td[" + (columnIndex + 1) + "]"));
+            TestBenchElement e = (TestBenchElement) findElement(By
+                    .xpath("./td[" + (columnIndex + 1) + "]"));
             return e.wrap(GridCellElement.class);
         }
     }
@@ -207,13 +206,66 @@ public class GridElement extends AbstractComponentElement {
      * @return Header cell element with given indices.
      */
     public GridCellElement getHeaderCell(int rowIndex, int colIndex) {
-        return getSubPart("#header[" + rowIndex + "][" + colIndex + "]")
-                .wrap(GridCellElement.class);
+        return getSubPart("#header[" + rowIndex + "][" + colIndex + "]").wrap(
+                GridCellElement.class);
+    }
+
+    /**
+     * Finds the header cell element with the given caption. If
+     * there are multiple headers with the same name, the first one is returned.
+     *
+     * @param caption
+     *            The header caption
+     * @return The first header cell element with a given caption.
+     * @throws NoSuchElementException
+     *             if there is no header row or no header cell with the given
+     *             text.
+     */
+    public GridCellElement getHeaderCellByCaption(String caption) {
+        List<WebElement> headerRows = findElement(By.vaadin("#header"))
+            .findElements(By.xpath("./tr/th"));
+        for (WebElement header : headerRows) {
+            if (caption.equals(header.getText())) {
+                return TestBenchElement
+                    .wrapElement(header, getCommandExecutor())
+                    .wrap(GridCellElement.class);
+            }
+        }
+        String errorMessage = String
+            .format("There is no header cell with %s caption. ", caption);
+        throw new NoSuchElementException(errorMessage);
+    }
+
+    /**
+     * Gets the header cell element with the given caption in the given header
+     * row. If there are multiple headers with the same name, the first one is
+     * returned.
+     *
+     * @param rowIndex
+     *            The index of the header row
+     * @param caption
+     *            The header caption
+     * @return The first header cell element with a given caption.
+     * @throws NoSuchElementException
+     *             if there is no header row or no header cell with the given
+     *             text.
+     */
+    public GridCellElement getHeaderCellByCaption(int rowIndex, String caption) {
+        List<GridCellElement> headerCells = getHeaderCells(rowIndex);
+        for (GridCellElement cell : headerCells) {
+            if (caption.equals(cell.getText())) {
+                return cell;
+            }
+        }
+        String errorMessage = String.format(
+                "The row with index %d does not have header with %s caption. ",
+                rowIndex, caption);
+        throw new NoSuchElementException(errorMessage);
     }
 
     /**
      * Gets footer cell element with given row and column index.
-     * 
+     *
      * @param rowIndex
      *            Row index
      * @param colIndex

uitest/src/main/java/com/vaadin/tests/elements/grid/GridUI.java

@@ -7,6 +7,7 @@ import com.vaadin.server.VaadinRequest;
 import com.vaadin.tests.components.AbstractTestUI;
 import com.vaadin.ui.Grid;
 import com.vaadin.ui.Label;
+import com.vaadin.ui.components.grid.HeaderRow;
 
 public class GridUI extends AbstractTestUI {
 
@@ -19,7 +20,10 @@ public class GridUI extends AbstractTestUI {
 
         final Grid<Item> grid = new Grid<Item>();
         grid.setItems(getMockData(rowCount));
-        grid.addColumn(Item::getFoo).setCaption("foo");
+        Grid.Column<Item, String> column = grid.addColumn(Item::getFoo)
+            .setCaption("foo");
+        HeaderRow row =grid.addHeaderRowAt(1);
+        row.getCell(column).setText("extra row");
         grid.addColumn(Item::getBar).setCaption("bar");
 
         grid.setDetailsGenerator(item -> {

uitest/src/test/java/com/vaadin/tests/elements/grid/GridUITest.java

@@ -2,6 +2,7 @@ package com.vaadin.tests.elements.grid;
 
 import org.junit.Assert;
 import org.junit.Test;
+import org.openqa.selenium.NoSuchElementException;
 
 import com.vaadin.testbench.elements.GridElement;
 import com.vaadin.testbench.elements.GridElement.GridRowElement;
@@ -41,6 +42,48 @@ public class GridUITest extends MultiBrowserTest {
         Assert.assertEquals(100, checkRows());
     }
 
+    @Test
+    public void testGetHeadersByCaptionFirstRowFirstColumn() {
+        openTestURL("rowCount=10&restartApplication");
+        GridElement grid = $(GridElement.class).first();
+        grid.getHeaderCellByCaption("foo");
+    }
+
+    @Test
+    public void testGetHeadersByCaptionFirstRowNotFirstColumn() {
+        openTestURL("rowCount=10&restartApplication");
+        GridElement grid = $(GridElement.class).first();
+        grid.getHeaderCellByCaption("bar");
+    }
+
+    @Test(expected = NoSuchElementException.class)
+    public void testGetHeadersByCaptionNoHeader() {
+        openTestURL("rowCount=10&restartApplication");
+        GridElement grid = $(GridElement.class).first();
+        grid.getHeaderCellByCaption("not existing caption");
+    }
+
+    @Test(expected = NoSuchElementException.class)
+    public void testGetHeadersByCaptionByIndexNoHeader() {
+        openTestURL("rowCount=10&restartApplication");
+        GridElement grid = $(GridElement.class).first();
+        grid.getHeaderCellByCaption(0, "not existing caption");
+    }
+
+    @Test
+    public void testGetHeadersByCaptionNotFirstRow() {
+        openTestURL("rowCount=10&restartApplication");
+        GridElement grid = $(GridElement.class).first();
+        grid.getHeaderCellByCaption("extra row");
+    }
+
+    @Test
+    public void testGetHeadersByCaptionByIndexNotFirstRow() {
+        openTestURL("rowCount=10&restartApplication");
+        GridElement grid = $(GridElement.class).first();
+        grid.getHeaderCellByCaption(1, "extra row");
+    }
+    
     private int checkRows() {
         int rowCount = 0;
         for (final GridRowElement row : getRows()) {