tests for custom element test

Author: navin772

py/test/selenium/webdriver/remote/custom_element_tests.py

@@ -0,0 +1,61 @@
+# Licensed to the Software Freedom Conservancy (SFC) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The SFC licenses this file
+# to you 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.
+import pytest
+
+from selenium import webdriver
+from selenium.webdriver.common.by import By
+from selenium.webdriver.remote.webelement import WebElement
+
+
+# Custom element class
+class MyCustomElement(WebElement):
+    def custom_method(self):
+        return "Custom element method"
+
+
+@pytest.fixture
+def driver():
+    options = webdriver.ChromeOptions()
+    driver = webdriver.Chrome(options=options)
+    yield driver
+    driver.quit()
+
+
+def test_find_element_with_custom_class(driver, pages):
+    """Test to ensure custom element class is used for a single element."""
+    driver._web_element_cls = MyCustomElement
+    pages.load("simpleTest.html")
+    element = driver.find_element(By.TAG_NAME, "body")
+    assert isinstance(element, MyCustomElement)
+    assert element.custom_method() == "Custom element method"
+
+
+def test_find_elements_with_custom_class(driver, pages):
+    """Test to ensure custom element class is used for multiple elements."""
+    driver._web_element_cls = MyCustomElement
+    pages.load("simpleTest.html")
+    elements = driver.find_elements(By.TAG_NAME, "div")
+    assert all(isinstance(el, MyCustomElement) for el in elements)
+    assert all(el.custom_method() == "Custom element method" for el in elements)
+
+
+def test_default_element_class(driver, pages):
+    """Test to ensure default WebElement class is used."""
+    pages.load("simpleTest.html")
+    element = driver.find_element(By.TAG_NAME, "body")
+    assert isinstance(element, WebElement)
+    assert not hasattr(element, "custom_method")

py/test/selenium/webdriver/remote/remote_custom_locator_tests.py

@@ -31,14 +31,12 @@ def convert(self, by, value):
 
 @pytest.fixture
 def driver() -> WebDriver:
-    # Setup for driver
     options = webdriver.ChromeOptions()
     driver = webdriver.Chrome(options=options)
     driver.locator_converter = CustomLocatorConverter()
-
     driver.get("data:text/html,<div custom-attr='example'>Test</div>")
+
     yield driver
-    # Teardown after test
     driver.quit()
 
 
@@ -49,6 +47,8 @@ def test_find_element_with_custom_locator(driver: WebDriver) -> None:
 
 
 def test_find_elements_with_custom_locator(driver: WebDriver) -> None:
+    driver.get("data:text/html,<div custom-attr='example'>Test1</div><div custom-attr='example'>Test2</div>")
     elements = driver.find_elements("custom", "example")
-    assert len(elements) == 1
-    assert elements[0].text == "Test"
+    assert len(elements) == 2
+    assert elements[0].text == "Test1"
+    assert elements[1].text == "Test2"