Add more methods directly into the raw "driver"

Author: mdmintz

seleniumbase/core/browser_launcher.py

@@ -116,6 +116,9 @@ def extend_driver(driver):
     page = types.SimpleNamespace()
     page.open = DM.open_url
     page.click = DM.click
+    page.click_link = DM.click_link
+    page.click_if_visible = DM.click_if_visible
+    page.click_active_element = DM.click_active_element
     page.send_keys = DM.send_keys
     page.press_keys = DM.press_keys
     page.type = DM.update_text
@@ -134,10 +137,13 @@ def extend_driver(driver):
     page.is_element_visible = DM.is_element_visible
     page.is_text_visible = DM.is_text_visible
     page.is_exact_text_visible = DM.is_exact_text_visible
+    page.is_attribute_present = DM.is_attribute_present
     page.get_text = DM.get_text
     page.find_element = DM.find_element
     page.find_elements = DM.find_elements
     page.locator = DM.locator
+    page.get_page_source = DM.get_page_source
+    page.get_title = DM.get_title
     driver.page = page
     js = types.SimpleNamespace()
     js.js_click = DM.js_click
@@ -150,6 +156,8 @@ def extend_driver(driver):
     driver.open = DM.open_url
     driver.click = DM.click
     driver.click_link = DM.click_link
+    driver.click_if_visible = DM.click_if_visible
+    driver.click_active_element = DM.click_active_element
     driver.send_keys = DM.send_keys
     driver.press_keys = DM.press_keys
     driver.type = DM.update_text
@@ -168,6 +176,7 @@ def extend_driver(driver):
     driver.is_element_visible = DM.is_element_visible
     driver.is_text_visible = DM.is_text_visible
     driver.is_exact_text_visible = DM.is_exact_text_visible
+    driver.is_attribute_present = DM.is_attribute_present
     driver.get_text = DM.get_text
     driver.js_click = DM.js_click
     driver.get_active_element_css = DM.get_active_element_css
@@ -176,6 +185,10 @@ def extend_driver(driver):
     driver.get_user_agent = DM.get_user_agent
     driver.highlight = DM.highlight
     driver.sleep = time.sleep
+    driver.get_page_source = DM.get_page_source
+    driver.get_title = DM.get_title
+    if hasattr(driver, "proxy"):
+        driver.set_wire_proxy = DM.set_wire_proxy
     return driver
 
 

seleniumbase/core/sb_driver.py

@@ -41,6 +41,12 @@ def locator(self, selector, by=None):
             pass
         raise Exception('No such Element: {%s} (by="%s")!' % (selector, by))
 
+    def get_page_source(self):
+        return self.driver.page_source
+
+    def get_title(self):
+        return self.driver.title
+
     def open_url(self, *args, **kwargs):
         page_actions.open_url(self.driver, *args, **kwargs)
 
@@ -50,6 +56,12 @@ def click(self, *args, **kwargs):
     def click_link(self, *args, **kwargs):
         page_actions.click_link(self.driver, *args, **kwargs)
 
+    def click_if_visible(self, *args, **kwargs):
+        page_actions.click_if_visible(self.driver, *args, **kwargs)
+
+    def click_active_element(self, *args, **kwargs):
+        page_actions.click_active_element(self.driver, *args, **kwargs)
+
     def send_keys(self, *args, **kwargs):
         page_actions.send_keys(self.driver, *args, **kwargs)
 
@@ -114,6 +126,9 @@ def get_text(self, *args, **kwargs):
     def js_click(self, *args, **kwargs):
         return page_actions.js_click(self.driver, *args, **kwargs)
 
+    def is_attribute_present(self, *args, **kwargs):
+        return page_actions.has_attribute(self.driver, *args, **kwargs)
+
     def get_active_element_css(self, *args, **kwargs):
         return js_utils.get_active_element_css(self.driver, *args, **kwargs)
 
@@ -128,3 +143,26 @@ def get_user_agent(self, *args, **kwargs):
 
     def highlight(self, *args, **kwargs):
         js_utils.highlight(self.driver, *args, **kwargs)
+
+    def set_wire_proxy(self, string):
+        """Set a proxy server for selenium-wire mode ("--wire")
+        Examples:  (ONLY avilable if using selenium-wire mode!)
+        driver.set_wire_proxy("SERVER:PORT")
+        driver.set_wire_proxy("socks5://SERVER:PORT")
+        driver.set_wire_proxy("USERNAME:PASSWORD@SERVER:PORT")
+        """
+        the_http = "http"
+        the_https = "https"
+        if string.startswith("socks4://"):
+            the_http = "socks4"
+            the_https = "socks4"
+        elif string.startswith("socks5://"):
+            the_http = "socks5"
+            the_https = "socks5"
+        string = string.split("//")[-1]
+        if hasattr(self.driver, "proxy"):
+            self.driver.proxy = {
+                "http": "%s://%s" % (the_http, string),
+                "https": "%s://%s" % (the_https, string),
+                "no_proxy": "localhost,127.0.0.1",
+            }

seleniumbase/fixtures/page_actions.py

@@ -1533,6 +1533,27 @@ def click_link(driver, link_text, timeout=settings.SMALL_TIMEOUT):
     element.click()
 
 
+def click_if_visible(
+    driver, selector, by="css selector", timeout=0
+):
+    selector, by = page_utils.recalculate_selector(selector, by)
+    if is_element_visible(driver, selector, by=by):
+        click(driver, selector, by=by, timeout=1)
+    elif timeout > 0:
+        try:
+            wait_for_element_visible(
+                driver, selector, by=by, timeout=timeout
+            )
+        except Exception:
+            pass
+        if is_element_visible(driver, selector, by=by):
+            click(driver, selector, by=by, timeout=1)
+
+
+def click_active_element(driver):
+    driver.execute_script("document.activeElement.click();")
+
+
 def js_click(
     driver, selector, by="css selector", timeout=settings.SMALL_TIMEOUT
 ):
@@ -1611,6 +1632,15 @@ def submit(driver, selector, by="css selector"):
     element.submit()
 
 
+def has_attribute(
+    driver, selector, attribute, value=None, by="css selector"
+):
+    selector, by = page_utils.recalculate_selector(selector, by)
+    return is_attribute_present(
+        driver, selector, attribute, value=value, by=by
+    )
+
+
 def assert_element_visible(
     driver, selector, by="css selector", timeout=settings.SMALL_TIMEOUT
 ):