py: fix RelativeBy#near to take 2 parameters (#13082)

* add failing test to check RelativeLocator#near accept single int

* fix s.w.support.RelativeBy#near to take 2 parameters

* more strict typing on s.w.support.relative_locator

* add some tests for s.w.support.relative_locator

* remove test case calling RelativeLocator#near wrong way

* fix linting issues

---------

Co-authored-by: Diego Molina <diemol@users.noreply.github.com>

py/selenium/webdriver/common/by.py

@@ -16,6 +16,8 @@
 # under the License.
 """The By implementation."""
 
+from typing import Literal
+
 
 class By:
     """Set of supported locator strategies."""
@@ -28,3 +30,6 @@ class By:
     TAG_NAME = "tag name"
     CLASS_NAME = "class name"
     CSS_SELECTOR = "css selector"
+
+
+ByType = Literal["id", "xpath", "link text", "partial link text", "name", "tag name", "class name", "css selector"]

py/selenium/webdriver/support/relative_locator.py

@@ -16,11 +16,14 @@
 # under the License.
 from typing import Dict
 from typing import List
+from typing import NoReturn
 from typing import Optional
 from typing import Union
+from typing import overload
 
 from selenium.common.exceptions import WebDriverException
 from selenium.webdriver.common.by import By
+from selenium.webdriver.common.by import ByType
 from selenium.webdriver.remote.webelement import WebElement
 
 
@@ -37,10 +40,10 @@ def with_tag_name(tag_name: str) -> "RelativeBy":
     """
     if not tag_name:
         raise WebDriverException("tag_name can not be null")
-    return RelativeBy({"css selector": tag_name})
+    return RelativeBy({By.CSS_SELECTOR: tag_name})
 
 
-def locate_with(by: By, using: str) -> "RelativeBy":
+def locate_with(by: ByType, using: str) -> "RelativeBy":
     """Start searching for relative objects your search criteria with By.
 
     :Args:
@@ -70,7 +73,9 @@ class RelativeBy:
         assert "mid" in ids
     """
 
-    def __init__(self, root: Optional[Dict[Union[By, str], str]] = None, filters: Optional[List] = None):
+    LocatorType = Dict[ByType, str]
+
+    def __init__(self, root: Optional[Dict[ByType, str]] = None, filters: Optional[List] = None):
         """Creates a new RelativeBy object. It is preferred if you use the
         `locate_with` method as this signature could change.
 
@@ -82,7 +87,13 @@ def __init__(self, root: Optional[Dict[Union[By, str], str]] = None, filters: Op
         self.root = root
         self.filters = filters or []
 
-    def above(self, element_or_locator: Union[WebElement, Dict] = None) -> "RelativeBy":
+    @overload
+    def above(self, element_or_locator: Union[WebElement, LocatorType]) -> "RelativeBy": ...
+
+    @overload
+    def above(self, element_or_locator: None = None) -> "NoReturn": ...
+
+    def above(self, element_or_locator: Union[WebElement, LocatorType, None] = None) -> "RelativeBy":
         """Add a filter to look for elements above.
 
         :Args:
@@ -94,7 +105,13 @@ def above(self, element_or_locator: Union[WebElement, Dict] = None) -> "Relative
         self.filters.append({"kind": "above", "args": [element_or_locator]})
         return self
 
-    def below(self, element_or_locator: Union[WebElement, Dict] = None) -> "RelativeBy":
+    @overload
+    def below(self, element_or_locator: Union[WebElement, LocatorType]) -> "RelativeBy": ...
+
+    @overload
+    def below(self, element_or_locator: None = None) -> "NoReturn": ...
+
+    def below(self, element_or_locator: Union[WebElement, Dict, None] = None) -> "RelativeBy":
         """Add a filter to look for elements below.
 
         :Args:
@@ -106,7 +123,13 @@ def below(self, element_or_locator: Union[WebElement, Dict] = None) -> "Relative
         self.filters.append({"kind": "below", "args": [element_or_locator]})
         return self
 
-    def to_left_of(self, element_or_locator: Union[WebElement, Dict] = None) -> "RelativeBy":
+    @overload
+    def to_left_of(self, element_or_locator: Union[WebElement, LocatorType]) -> "RelativeBy": ...
+
+    @overload
+    def to_left_of(self, element_or_locator: None = None) -> "NoReturn": ...
+
+    def to_left_of(self, element_or_locator: Union[WebElement, Dict, None] = None) -> "RelativeBy":
         """Add a filter to look for elements to the left of.
 
         :Args:
@@ -118,7 +141,13 @@ def to_left_of(self, element_or_locator: Union[WebElement, Dict] = None) -> "Rel
         self.filters.append({"kind": "left", "args": [element_or_locator]})
         return self
 
-    def to_right_of(self, element_or_locator: Union[WebElement, Dict] = None) -> "RelativeBy":
+    @overload
+    def to_right_of(self, element_or_locator: Union[WebElement, LocatorType]) -> "RelativeBy": ...
+
+    @overload
+    def to_right_of(self, element_or_locator: None = None) -> "NoReturn": ...
+
+    def to_right_of(self, element_or_locator: Union[WebElement, Dict, None] = None) -> "RelativeBy":
         """Add a filter to look for elements right of.
 
         :Args:
@@ -130,16 +159,25 @@ def to_right_of(self, element_or_locator: Union[WebElement, Dict] = None) -> "Re
         self.filters.append({"kind": "right", "args": [element_or_locator]})
         return self
 
-    def near(self, element_or_locator_distance: Union[WebElement, Dict, int] = None) -> "RelativeBy":
+    @overload
+    def near(self, element_or_locator: Union[WebElement, LocatorType], distance: int = 50) -> "RelativeBy": ...
+
+    @overload
+    def near(self, element_or_locator: None = None, distance: int = 50) -> "NoReturn": ...
+
+    def near(self, element_or_locator: Union[WebElement, LocatorType, None] = None, distance: int = 50) -> "RelativeBy":
         """Add a filter to look for elements near.
 
         :Args:
-            - element_or_locator_distance: Element to look near by the element or within a distance
+            - element_or_locator: Element to look near by the element or within a distance
+            - distance: distance in pixel
         """
-        if not element_or_locator_distance:
-            raise WebDriverException("Element or locator or distance must be given when calling near method")
+        if not element_or_locator:
+            raise WebDriverException("Element or locator must be given when calling near method")
+        if distance <= 0:
+            raise WebDriverException("Distance must be positive")
 
-        self.filters.append({"kind": "near", "args": [element_or_locator_distance]})
+        self.filters.append({"kind": "near", "args": [element_or_locator, distance]})
         return self
 
     def to_dict(self) -> Dict:

py/test/selenium/webdriver/support/relative_by_tests.py

@@ -31,6 +31,14 @@ def test_should_be_able_to_find_first_one(driver, pages):
     assert el.get_attribute("id") == "mid"
 
 
+def test_should_be_able_to_find_first_one_by_locator(driver, pages):
+    pages.load("relative_locators.html")
+
+    el = driver.find_element(with_tag_name("p").above({By.ID: "below"}))
+
+    assert el.get_attribute("id") == "mid"
+
+
 def test_should_be_able_to_find_elements_above_another(driver, pages):
     pages.load("relative_locators.html")
     lowest = driver.find_element(By.ID, "below")
@@ -42,6 +50,16 @@ def test_should_be_able_to_find_elements_above_another(driver, pages):
     assert "mid" in ids
 
 
+def test_should_be_able_to_find_elements_above_another_by_locator(driver, pages):
+    pages.load("relative_locators.html")
+
+    elements = driver.find_elements(with_tag_name("p").above({By.ID: "below"}))
+
+    ids = [el.get_attribute("id") for el in elements]
+    assert "above" in ids
+    assert "mid" in ids
+
+
 def test_should_be_able_to_combine_filters(driver, pages):
     pages.load("relative_locators.html")
 
@@ -55,6 +73,15 @@ def test_should_be_able_to_combine_filters(driver, pages):
     assert "third" in ids
 
 
+def test_should_be_able_to_combine_filters_by_locator(driver, pages):
+    pages.load("relative_locators.html")
+
+    elements = driver.find_elements(with_tag_name("td").above({By.ID: "center"}).to_right_of({By.ID: "second"}))
+
+    ids = [el.get_attribute("id") for el in elements]
+    assert "third" in ids
+
+
 def test_should_be_able_to_use_css_selectors(driver, pages):
     pages.load("relative_locators.html")
 
@@ -68,6 +95,17 @@ def test_should_be_able_to_use_css_selectors(driver, pages):
     assert "third" in ids
 
 
+def test_should_be_able_to_use_css_selectors_by_locator(driver, pages):
+    pages.load("relative_locators.html")
+
+    elements = driver.find_elements(
+        locate_with(By.CSS_SELECTOR, "td").above({By.ID: "center"}).to_right_of({By.ID: "second"})
+    )
+
+    ids = [el.get_attribute("id") for el in elements]
+    assert "third" in ids
+
+
 def test_should_be_able_to_use_xpath(driver, pages):
     pages.load("relative_locators.html")
 
@@ -81,6 +119,15 @@ def test_should_be_able_to_use_xpath(driver, pages):
     assert "fourth" in ids
 
 
+def test_should_be_able_to_use_xpath_by_locator(driver, pages):
+    pages.load("relative_locators.html")
+
+    elements = driver.find_elements(locate_with(By.XPATH, "//td[1]").below({By.ID: "second"}).above({By.ID: "seventh"}))
+
+    ids = [el.get_attribute("id") for el in elements]
+    assert "fourth" in ids
+
+
 def test_no_such_element_is_raised_rather_than_index_error(driver, pages):
     pages.load("relative_locators.html")
     with pytest.raises(NoSuchElementException) as exc:
@@ -89,6 +136,13 @@ def test_no_such_element_is_raised_rather_than_index_error(driver, pages):
     assert "Cannot locate relative element with: {'id': 'nonexistentid'}" in exc.value.msg
 
 
+def test_no_such_element_is_raised_rather_than_index_error_by_locator(driver, pages):
+    pages.load("relative_locators.html")
+    with pytest.raises(NoSuchElementException) as exc:
+        driver.find_element(locate_with(By.ID, "nonexistentid").above({By.ID: "second"}))
+    assert "Cannot locate relative element with: {'id': 'nonexistentid'}" in exc.value.msg
+
+
 def test_near_locator_should_find_near_elements(driver, pages):
     pages.load("relative_locators.html")
     rect1 = driver.find_element(By.ID, "rect1")
@@ -98,6 +152,14 @@ def test_near_locator_should_find_near_elements(driver, pages):
     assert el.get_attribute("id") == "rect2"
 
 
+def test_near_locator_should_find_near_elements_by_locator(driver, pages):
+    pages.load("relative_locators.html")
+
+    el = driver.find_element(locate_with(By.ID, "rect2").near({By.ID: "rect1"}))
+
+    assert el.get_attribute("id") == "rect2"
+
+
 def test_near_locator_should_not_find_far_elements(driver, pages):
     pages.load("relative_locators.html")
     rect3 = driver.find_element(By.ID, "rect3")
@@ -106,3 +168,29 @@ def test_near_locator_should_not_find_far_elements(driver, pages):
         driver.find_element(locate_with(By.ID, "rect4").near(rect3))
 
     assert "Cannot locate relative element with: {'id': 'rect4'}" in exc.value.msg
+
+
+def test_near_locator_should_not_find_far_elements_by_locator(driver, pages):
+    pages.load("relative_locators.html")
+
+    with pytest.raises(NoSuchElementException) as exc:
+        driver.find_element(locate_with(By.ID, "rect4").near({By.ID: "rect3"}))
+
+    assert "Cannot locate relative element with: {'id': 'rect4'}" in exc.value.msg
+
+
+def test_near_locator_should_find_far_elements(driver, pages):
+    pages.load("relative_locators.html")
+    rect3 = driver.find_element(By.ID, "rect3")
+
+    el = driver.find_element(locate_with(By.ID, "rect4").near(rect3, 100))
+
+    assert el.get_attribute("id") == "rect4"
+
+
+def test_near_locator_should_find_far_elements_by_locator(driver, pages):
+    pages.load("relative_locators.html")
+
+    el = driver.find_element(locate_with(By.ID, "rect4").near({By.ID: "rect3"}, 100))
+
+    assert el.get_attribute("id") == "rect4"