Add custom expectation to wait for n elements

notebook/tests/selenium/test_display_isolation.py

@@ -42,7 +42,7 @@ def isolated_html(notebook):
         isolated)
     notebook.add_and_execute_cell(content=display_i)
 
-    wait_for_tag(notebook.browser, "iframe")
+    iframe = wait_for_tag(notebook.browser, "iframe", single=True)
 
     # The non-isolated div will be in the body
     non_isolated_div = notebook.body.find_element_by_id("non-isolated")
@@ -53,7 +53,6 @@ def isolated_html(notebook):
     assert test_div.value_of_css_property("color") == red
 
     # The isolated div will be in an iframe, only that element will be blue
-    iframe = notebook.body.find_element_by_tag_name("iframe")
     notebook.browser.switch_to.frame(iframe)
     isolated_div = notebook.browser.find_element_by_id("isolated")
     assert isolated_div.value_of_css_property("color") == blue
@@ -79,8 +78,7 @@ def isolated_svg(notebook):
         content="display_svg(SVG(s1), metadata=dict(isolated=True))")
     notebook.add_and_execute_cell(
         content="display_svg(SVG(s2), metadata=dict(isolated=True))")
-    wait_for_tag(notebook.browser, "iframe")
+    iframes = wait_for_tag(notebook.browser, "iframe", wait_for_n=2)
-    iframes = notebook.body.find_elements_by_tag_name("iframe")
 
     # The first rectangle will be red
     notebook.browser.switch_to.frame(iframes[0])

notebook/tests/selenium/utils.py

@@ -2,6 +2,7 @@ import os
 import time
 from selenium.webdriver import ActionChains
 from selenium.webdriver.common.by import By
+from selenium.common.exceptions import WebDriverException
 from selenium.webdriver.common.keys import Keys
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
@@ -12,13 +13,33 @@ from contextlib import contextmanager
 pjoin = os.path.join
 
 
-def wait_for_selector(driver, selector, timeout=10, visible=False, single=False):
+def wait_for_selector(driver, selector, timeout=10, visible=False, single=False, wait_for_n=1):
+    if wait_for_n > 1:
+            return _wait_for_multiple(
+                driver, By.CSS_SELECTOR, selector, timeout, wait_for_n, visible)
     return _wait_for(driver, By.CSS_SELECTOR, selector, timeout, visible, single)
 
-def wait_for_tag(driver, tag, timeout=10, visible=False, single=False):
+
+def wait_for_tag(driver, tag, timeout=10, visible=False, single=False, wait_for_n=1):
+    if wait_for_n > 1:
+            return _wait_for_multiple(
+                driver, By.TAG_NAME, tag, timeout, wait_for_n, visible)
     return _wait_for(driver, By.TAG_NAME, tag, timeout, visible, single)
 
+
 def _wait_for(driver, locator_type, locator, timeout=10, visible=False, single=False):
+    """Waits `timeout` seconds for the specified condition to be met. Condition is
+    met if any matching element is found. Returns located element(s) when found.
+
+    Args:
+        driver: Selenium web driver instance
+        locator_type: type of locator (e.g. By.CSS_SELECTOR or By.TAG_NAME)
+        locator: name of tag, class, etc. to wait for
+        timeout: how long to wait for presence/visibility of element
+        visible: if True, require that element is not only present, but visible
+        single: if True, return a single element, otherwise return a list of matching
+        elements
+    """
     wait = WebDriverWait(driver, timeout)
     if single:
         if visible:
@@ -33,6 +54,36 @@ def _wait_for(driver, locator_type, locator, timeout=10, visible=False, single=F
     return wait.until(conditional((locator_type, locator)))
 
 
+def _wait_for_multiple(driver, locator_type, locator, timeout, wait_for_n, visible=False):
+    """Waits until `wait_for_n` matching elements to be present (or visible).
+    Returns located elements when found.
+
+    Args:
+        driver: Selenium web driver instance
+        locator_type: type of locator (e.g. By.CSS_SELECTOR or By.TAG_NAME)
+        locator: name of tag, class, etc. to wait for
+        timeout: how long to wait for presence/visibility of element
+        wait_for_n: wait until this number of matching elements are present/visible
+        visible: if True, require that elements are not only present, but visible
+    """
+    wait = WebDriverWait(driver, timeout)
+
+    def multiple_found(driver):
+        try:
+            elements = driver.find_elements(locator_type, locator)
+        except WebDriverException as e:
+            raise e
+        if len(elements) < wait_for_n:
+            return False
+        if visible:
+            for element in elements:
+                if not element.is_displayed():
+                    return False
+        return elements
+
+    return wait.until(multiple_found)
+
+
 class CellTypeError(ValueError):
     
     def __init__(self, message=""):