Add driver methods

Author: mdmintz

seleniumbase/core/browser_launcher.py

@@ -187,8 +187,10 @@ def extend_driver(driver):
     driver.highlight = DM.highlight
     driver.highlight_click = DM.highlight_click
     driver.sleep = time.sleep
+    driver.get_attribute = DM.get_attribute
     driver.get_page_source = DM.get_page_source
     driver.get_title = DM.get_title
+    driver.switch_to_frame = DM.switch_to_frame
     if hasattr(driver, "proxy"):
         driver.set_wire_proxy = DM.set_wire_proxy
     return driver

seleniumbase/core/sb_driver.py

@@ -1,4 +1,5 @@
 """Add new methods to extend the driver"""
+from selenium.webdriver.remote.webelement import WebElement
 from seleniumbase.fixtures import js_utils
 from seleniumbase.fixtures import page_actions
 from seleniumbase.fixtures import page_utils
@@ -41,6 +42,10 @@ def locator(self, selector, by=None):
             pass
         raise Exception('No such Element: {%s} (by="%s")!' % (selector, by))
 
+    def get_attribute(self, selector, attribute, by="css selector"):
+        element = self.locator(selector, by=by)
+        return element.get_attribute(attribute)
+
     def get_page_source(self):
         return self.driver.page_source
 
@@ -158,6 +163,13 @@ def highlight_click(self, *args, **kwargs):
             kwargs.pop("loops")
         page_actions.click(self.driver, *args, **kwargs)
 
+    def switch_to_frame(self, frame):
+        if isinstance(frame, WebElement):
+            self.driver.switch_to.frame(frame)
+        else:
+            iframe = self.locator(frame)
+            self.driver.switch_to.frame(iframe)
+
     def set_wire_proxy(self, string):
         """Set a proxy server for selenium-wire mode ("--wire")
         Examples:  (ONLY avilable if using selenium-wire mode!)