[py] Fix custom locator tests

Author: cgoldberg

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

@@ -29,9 +29,11 @@ def custom_method(self):
 
 @pytest.fixture()
 def custom_element_driver(driver):
-    driver._web_element_cls = MyCustomElement
-    yield driver
-    driver._web_element_cls = WebElement
+    try:
+        driver._web_element_cls = MyCustomElement
+        yield driver
+    finally:
+        driver._web_element_cls = WebElement
 
 
 def test_find_element_with_custom_class(custom_element_driver, pages):

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

@@ -17,6 +17,8 @@
 
 import pytest
 
+from selenium import webdriver
+from selenium.webdriver.firefox.options import Options
 from selenium.webdriver.remote.locator_converter import LocatorConverter
 
 
@@ -28,18 +30,28 @@ def convert(self, by, value):
         return super().convert(by, value)
 
 
-@pytest.mark.skip(reason="Needs to be updated")
-def test_find_element_with_custom_locator(driver):
-    driver.get("data:text/html,<div custom-attr='example'>Test</div>")
-    element = driver.find_element("custom", "example")
+@pytest.fixture()
+def custom_locator_driver(headless):
+    options = Options()
+    if headless:
+        options.add_argument("-headless")
+    try:
+        driver = webdriver.Remote(options=options, locator_converter=CustomLocatorConverter())
+        yield driver
+    finally:
+        driver.quit()
+
+
+def test_find_element_with_custom_locator(custom_locator_driver):
+    custom_locator_driver.get("data:text/html,<div custom-attr='example'>Test</div>")
+    element = custom_locator_driver.find_element("custom", "example")
     assert element is not None
     assert element.text == "Test"
 
 
-@pytest.mark.skip(reason="Needs to be updated")
-def test_find_elements_with_custom_locator(driver):
-    driver.get("data:text/html,<div custom-attr='example'>Test1</div><div custom-attr='example'>Test2</div>")
-    elements = driver.find_elements("custom", "example")
+def test_find_elements_with_custom_locator(custom_locator_driver):
+    custom_locator_driver.get("data:text/html,<div custom-attr='example'>Test1</div><div custom-attr='example'>Test2</div>")
+    elements = custom_locator_driver.find_elements("custom", "example")
     assert len(elements) == 2
     assert elements[0].text == "Test1"
     assert elements[1].text == "Test2"