Checkbox get label improvement

htmlelements-java/src/main/java/ru/yandex/qatools/htmlelements/element/CheckBox.java

@@ -3,6 +3,10 @@
 import org.openqa.selenium.By;
 import org.openqa.selenium.NoSuchElementException;
 import org.openqa.selenium.WebElement;
+import ru.yandex.qatools.htmlelements.utils.HtmlElementUtils;
+
+import java.util.LinkedList;
+import java.util.List;
 
 /**
  * Represents web page checkbox control.
@@ -21,14 +25,52 @@ public CheckBox(WebElement wrappedElement) {
     }
 
     /**
-     * Finds label corresponding to this checkbox using "following-sibling::label" xpath.
+     * Finds label corresponding to checkbox element
      *
      * @return {@code WebElement} representing label or {@code null} if no label has been found.
      */
     public WebElement getLabel() {
+        WebElement label = null;
+        String id = null;
+
+        // Trying to find id attibute in element
         try {
-            return getWrappedElement().findElement(By.xpath("following-sibling::label"));
-        } catch (NoSuchElementException e) {
+            id = HtmlElementUtils.xpathLiteral(getWrappedElement().getAttribute("id"));
+        } catch (NullPointerException | NoSuchElementException ex) {
+            // Do nothing. Element doesn't have id attribute
+        }
+
+        if (id != null) {
+            // Label with matching "for" attribute
+            // Trying to find element from !ROOT! node
+            String xpath = String.format("(//label[@for = %s])[1]", id);
+            label = findLabel(By.xpath(xpath));
+        }
+
+        if (label == null) {
+            // Label wrapped around checkbox
+            label = findLabel(By.xpath("parent::label"));
+        }
+
+        if (label == null) {
+            // Label right next to checkbox
+            label = findLabel(By.xpath("following-sibling::label"));
+        }
+
+        return label;
+    }
+
+    /**
+     * Finds label in checkbox element
+     *
+     * @param by path to label
+     * @return {@code WebElement} representing label or {@code null} if label not found
+     */
+    private WebElement findLabel(By by) {
+        try {
+            return getWrappedElement().findElement(by);
+
+        } catch (NoSuchElementException ex) {
             return null;
         }
     }

htmlelements-java/src/main/java/ru/yandex/qatools/htmlelements/utils/HtmlElementUtils.java

@@ -13,6 +13,7 @@
 
 import java.lang.reflect.*;
 import java.net.URL;
+import java.util.LinkedList;
 import java.util.List;
 
 import static org.apache.commons.lang3.reflect.ConstructorUtils.invokeConstructor;
@@ -177,4 +178,40 @@ public static boolean existsInClasspath(final String fileName) {
     public static URL getResourceFromClasspath(final String fileName) {
         return Thread.currentThread().getContextClassLoader().getResource(fileName);
     }
+
+    public static String xpathLiteral(String s) {
+        if (s.indexOf("'") == -1) {
+            return String.format("'%s'", s);
+        }
+
+        if (s.indexOf("\"") == -1) {
+            return String.format("\"%s\"", s);
+        }
+
+        String string = s;
+        List<String> parts = new LinkedList<>();
+
+        while (true) {
+            int pos;
+            if ((pos = string.indexOf("'")) != -1) {
+                parts.add(string.substring(0, pos));
+                parts.add("\"'\"");
+
+                string = string.substring(pos + 1);
+            } else {
+                parts.add(String.format("'%s'", string));
+                break;
+            }
+        }
+
+        StringBuilder result = new StringBuilder();
+        for (String part: parts) {
+            result.append(part);
+            result.append(",");
+        }
+        result.deleteCharAt(result.length() - 1);
+
+        return result.toString();
+    }
+
 }