Optimize methods for switching into iframes

Author: mdmintz

seleniumbase/fixtures/base_case.py

@@ -1061,7 +1061,7 @@ def is_element_in_an_iframe(self, selector, by=By.CSS_SELECTOR):
         return False
 
     def switch_to_frame_of_element(self, selector, by=By.CSS_SELECTOR):
-        """ Set driver control to the iframe of the element (assuming the
+        """ Set driver control to the iframe containing element (assuming the
             element is in a single-nested iframe) and returns the iframe name.
             If element is not in an iframe, returns None, and nothing happens.
             May not work if multiple iframes are nested within each other. """
@@ -1076,13 +1076,26 @@ def switch_to_frame_of_element(self, selector, by=By.CSS_SELECTOR):
                 iframe_identifier = iframe['name']
             elif iframe.has_attr('id') and len(iframe['id']) > 0:
                 iframe_identifier = iframe['id']
+            elif iframe.has_attr('class') and len(iframe['class']) > 0:
+                iframe_class = " ".join(iframe["class"])
+                iframe_identifier = '[class="%s"]' % iframe_class
             else:
                 continue
-            self.switch_to_frame(iframe_identifier)
-            if self.is_element_present(selector, by=by):
-                return iframe_identifier
+            try:
+                self.switch_to_frame(iframe_identifier, timeout=1)
+                if self.is_element_present(selector, by=by):
+                    return iframe_identifier
+            except Exception:
+                pass
             self.switch_to_default_content()
-        return None
+        try:
+            self.switch_to_frame(selector, timeout=1)
+            return selector
+        except Exception:
+            if self.is_element_present(selector, by=by):
+                return ""
+            raise Exception("Could not switch to iframe containing "
+                            "element {%s}!" % selector)
 
     def hover_on_element(self, selector, by=By.CSS_SELECTOR):
         selector, by = self.__recalculate_selector(selector, by)
@@ -1310,7 +1323,13 @@ def maximize_window(self):
         self.__demo_mode_pause_if_active()
 
     def switch_to_frame(self, frame, timeout=None):
-        """ Sets driver control to the specified browser frame. """
+        """
+        Wait for an iframe to appear, and switch to it. This should be
+        usable as a drop-in replacement for driver.switch_to.frame().
+        @Params
+        frame - the frame element, name, id, index, or selector
+        timeout - the time to wait for the alert in seconds
+        """
         if not timeout:
             timeout = settings.SMALL_TIMEOUT
         if self.timeout_multiplier and timeout == settings.SMALL_TIMEOUT:

seleniumbase/fixtures/page_actions.py

@@ -33,6 +33,7 @@
 from selenium.webdriver.remote.errorhandler import NoSuchWindowException
 from seleniumbase.config import settings
 from seleniumbase.core import log_helper
+from seleniumbase.fixtures import page_utils
 
 
 def is_element_present(driver, selector, by=By.CSS_SELECTOR):
@@ -604,11 +605,11 @@ def wait_for_and_switch_to_alert(driver, timeout=settings.LARGE_TIMEOUT):
 
 def switch_to_frame(driver, frame, timeout=settings.SMALL_TIMEOUT):
     """
-    Wait for an iframe to appear, and switch to it. This should be usable
-    as a drop-in replacement for driver.switch_to.frame().
+    Wait for an iframe to appear, and switch to it. This should be
+    usable as a drop-in replacement for driver.switch_to.frame().
     @Params
     driver - the webdriver object (required)
-    frame - the frame element, name, or index
+    frame - the frame element, name, id, index, or selector
     timeout - the time to wait for the alert in seconds
     """
     start_ms = time.time() * 1000.0
@@ -618,6 +619,19 @@ def switch_to_frame(driver, frame, timeout=settings.SMALL_TIMEOUT):
             driver.switch_to.frame(frame)
             return True
         except NoSuchFrameException:
+            if type(frame) is str:
+                by = None
+                if page_utils.is_xpath_selector(frame):
+                    by = By.XPATH
+                else:
+                    by = By.CSS_SELECTOR
+                if is_element_visible(driver, frame, by=by):
+                    try:
+                        element = driver.find_element(by=by, value=frame)
+                        driver.switch_to.frame(element)
+                        return True
+                    except Exception:
+                        pass
             now_ms = time.time() * 1000.0
             if now_ms >= stop_ms:
                 break