Merge pull request #4185 from askerry/fix_race

Add wait_for_tag to avoid race condition in display isolation test

Author: takluyver

notebook/tests/selenium/test_display_isolation.py

@@ -3,6 +3,7 @@
 An object whose metadata contains an "isolated" tag must be isolated
 from the rest of the document.
 """
+from .utils import wait_for_tag
 
 
 def test_display_isolation(notebook):
@@ -26,7 +27,7 @@ def isolated_html(notebook):
     """
     red = 'rgb(255, 0, 0)'
     blue = 'rgb(0, 0, 255)'
-    test_str = "<div id='test'>Should be red from non-isolation</div>"
+    test_str = "<div id='test'>Should turn red from non-isolation</div>"
     notebook.add_and_execute_cell(content="display(HTML(%r))" % test_str)
     non_isolated = (
         "<style>div{color:%s;}</style>" % red +
@@ -41,6 +42,8 @@ def isolated_html(notebook):
         isolated)
     notebook.add_and_execute_cell(content=display_i)
 
+    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")
     assert non_isolated_div.value_of_css_property("color") == red
@@ -50,11 +53,13 @@ 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
     notebook.browser.switch_to.default_content()
+    # Clean up the html test cells
+    for i in range(1, len(notebook.cells)):
+        notebook.delete_cell(1)
 
 
 def isolated_svg(notebook):
@@ -73,16 +78,19 @@ 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))")
-    iframes = notebook.body.find_elements_by_tag_name("iframe")
+    iframes = wait_for_tag(notebook.browser, "iframe", wait_for_n=2)
 
     # The first rectangle will be red
-    notebook.browser.switch_to.frame(iframes[1])
+    notebook.browser.switch_to.frame(iframes[0])
     isolated_svg_1 = notebook.browser.find_element_by_id('r1')
     assert isolated_svg_1.value_of_css_property("fill") == yellow
     notebook.browser.switch_to.default_content()
 
     # The second rectangle will be black
-    notebook.browser.switch_to.frame(iframes[2])
+    notebook.browser.switch_to.frame(iframes[1])
     isolated_svg_2 = notebook.browser.find_element_by_id('r2')
     assert isolated_svg_2.value_of_css_property("fill") == black
 
+    # Clean up the svg test cells
+    for i in range(1, len(notebook.cells)):
+        notebook.delete_cell(1)

notebook/tests/selenium/utils.py

@@ -2,6 +2,7 @@
 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 @@
 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,31 @@ 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):
+        elements = driver.find_elements(locator_type, locator)
+        if visible:
+            elements = [e for e in elements if e.is_displayed()]
+        if len(elements) < wait_for_n:
+            return False
+        return elements
+
+    return wait.until(multiple_found)
+
+
 class CellTypeError(ValueError):
 
     def __init__(self, message=""):