Improve and add driver methods

Author: mdmintz

seleniumbase/core/browser_launcher.py

@@ -108,14 +108,18 @@ def make_driver_executable_if_not(driver_path):
 def extend_driver(driver):
     # Extend the driver with new methods
     driver.default_find_element = driver.find_element
+    driver.default_find_elements = driver.find_elements
     DM = sb_driver.DriverMethods(driver)
     driver.find_element = DM.find_element
+    driver.find_elements = DM.find_elements
     driver.locator = DM.locator
     page = types.SimpleNamespace()
     page.open = DM.open_url
     page.click = DM.click
     page.send_keys = DM.send_keys
+    page.press_keys = DM.press_keys
     page.type = DM.update_text
+    page.submit = DM.submit
     page.assert_element = DM.assert_element_visible
     page.assert_element_present = DM.assert_element_present
     page.assert_element_not_visible = DM.assert_element_not_visible
@@ -131,6 +135,9 @@ def extend_driver(driver):
     page.is_text_visible = DM.is_text_visible
     page.is_exact_text_visible = DM.is_exact_text_visible
     page.get_text = DM.get_text
+    page.find_element = DM.find_element
+    page.find_elements = DM.find_elements
+    page.locator = DM.locator
     driver.page = page
     js = types.SimpleNamespace()
     js.js_click = DM.js_click
@@ -142,8 +149,11 @@ def extend_driver(driver):
     driver.js = js
     driver.open = DM.open_url
     driver.click = DM.click
+    driver.click_link = DM.click_link
     driver.send_keys = DM.send_keys
+    driver.press_keys = DM.press_keys
     driver.type = DM.update_text
+    driver.submit = DM.submit
     driver.assert_element = DM.assert_element_visible
     driver.assert_element_present = DM.assert_element_present
     driver.assert_element_not_visible = DM.assert_element_not_visible

seleniumbase/core/sb_driver.py

@@ -8,14 +8,26 @@ class DriverMethods():
     def __init__(self, driver):
         self.driver = driver
 
-    def find_element(self, by, value=None):
+    def find_element(self, by=None, value=None):
         if not value:
             value = by
             by = "css selector"
+        elif not by:
+            by = "css selector"
         else:
             value, by = page_utils.swap_selector_and_by_if_reversed(value, by)
         return self.driver.default_find_element(by=by, value=value)
 
+    def find_elements(self, by=None, value=None):
+        if not value:
+            value = by
+            by = "css selector"
+        elif not by:
+            by = "css selector"
+        else:
+            value, by = page_utils.swap_selector_and_by_if_reversed(value, by)
+        return self.driver.default_find_elements(by=by, value=value)
+
     def locator(self, selector, by=None):
         if not by:
             by = "css selector"
@@ -35,12 +47,21 @@ def open_url(self, *args, **kwargs):
     def click(self, *args, **kwargs):
         page_actions.click(self.driver, *args, **kwargs)
 
+    def click_link(self, *args, **kwargs):
+        page_actions.click_link(self.driver, *args, **kwargs)
+
     def send_keys(self, *args, **kwargs):
         page_actions.send_keys(self.driver, *args, **kwargs)
 
+    def press_keys(self, *args, **kwargs):
+        page_actions.press_keys(self.driver, *args, **kwargs)
+
     def update_text(self, *args, **kwargs):
         page_actions.update_text(self.driver, *args, **kwargs)
 
+    def submit(self, *args, **kwargs):
+        page_actions.submit(self.driver, *args, **kwargs)
+
     def assert_element_visible(self, *args, **kwargs):
         page_actions.assert_element_visible(self.driver, *args, **kwargs)
 

seleniumbase/fixtures/page_actions.py

@@ -28,6 +28,7 @@
 from selenium.common.exceptions import NoSuchWindowException
 from selenium.common.exceptions import StaleElementReferenceException
 from selenium.webdriver.common.action_chains import ActionChains
+from selenium.webdriver.common.keys import Keys
 from seleniumbase.common.exceptions import LinkTextNotFoundException
 from seleniumbase.common.exceptions import TextNotVisibleException
 from seleniumbase.config import settings
@@ -1525,6 +1526,13 @@ def click(driver, selector, by="css selector", timeout=settings.SMALL_TIMEOUT):
     element.click()
 
 
+def click_link(driver, link_text, timeout=settings.SMALL_TIMEOUT):
+    element = wait_for_element_clickable(
+        driver, link_text, by="link text", timeout=timeout
+    )
+    element.click()
+
+
 def js_click(
     driver, selector, by="css selector", timeout=settings.SMALL_TIMEOUT
 ):
@@ -1564,6 +1572,22 @@ def send_keys(
         element.submit()
 
 
+def press_keys(
+    driver, selector, text, by="css selector", timeout=settings.LARGE_TIMEOUT
+):
+    selector, by = page_utils.recalculate_selector(selector, by)
+    element = wait_for_element_clickable(
+        driver, selector, by=by, timeout=timeout
+    )
+    if not text.endswith("\n"):
+        for key in text:
+            element.send_keys(key)
+    else:
+        for key in text[:-1]:
+            element.send_keys(key)
+        element.send_keys(Keys.RETURN)
+
+
 def update_text(
     driver, selector, text, by="css selector", timeout=settings.LARGE_TIMEOUT
 ):
@@ -1579,6 +1603,14 @@ def update_text(
         element.submit()
 
 
+def submit(driver, selector, by="css selector"):
+    selector, by = page_utils.recalculate_selector(selector, by)
+    element = wait_for_element_clickable(
+        driver, selector, by=by, timeout=settings.SMALL_TIMEOUT
+    )
+    element.submit()
+
+
 def assert_element_visible(
     driver, selector, by="css selector", timeout=settings.SMALL_TIMEOUT
 ):