Merge pull request #131 from mshriver/update-wait-for-element

Update wait_for_element, change exception behavior

src/widgetastic/browser.py

@@ -273,7 +273,7 @@ def elements(
         return result
 
     def wait_for_element(
-            self, locator, parent=None, visible=False, timeout=5, delay=0.2, exception=True,
+            self, locator, parent=None, visible=False, timeout=5, delay=0.2,
             ensure_page_safe=False):
         """Wait for presence or visibility of elements specified by a locator.
 
@@ -283,31 +283,36 @@ def wait_for_element(
                      also checks visibility.
             timeout: How long to wait for.
             delay: How often to check.
-            exception: If True (default), in case of element not being found an exception will be
-                       raised. If False, it returns False.
             ensure_page_safe: Whether to call the ``ensure_page_safe`` hook on repeat.
 
         Returns:
             :py:class:`selenium.webdriver.remote.webelement.WebElement` if element found according
-            to params. ``None`` if not found and ``exception=False``.
+            to params.
 
         Raises:
-            :py:class:`selenium.common.exceptions.NoSuchElementException` if element not found and
-            ``exception=True``.
+            :py:class:`selenium.common.exceptions.NoSuchElementException` if element not found.
         """
+        def _element_lookup():
+            try:
+                return self.elements(locator,
+                                     parent=parent,
+                                     check_visibility=visible,
+                                     check_safe=ensure_page_safe)
+            # allow other exceptions through to caller on first wait
+            except NoSuchElementException:
+                return False
+        # turn the timeout into NoSuchElement
         try:
-            result = wait_for(
-                lambda: self.elements(locator, parent=parent, check_visibility=visible,
-                                      check_safe=ensure_page_safe),
-                num_sec=timeout, delay=delay, fail_condition=lambda elements: not bool(elements),
-                fail_func=self.plugin.ensure_page_safe if ensure_page_safe else None)
+            result = wait_for(_element_lookup,
+                              num_sec=timeout,
+                              delay=delay,
+                              fail_condition=lambda elements: not bool(elements),
+                              fail_func=self.plugin.ensure_page_safe if ensure_page_safe else None)
         except TimedOutError:
-            if exception:
-                raise NoSuchElementException('Could not wait for element {!r}'.format(locator))
-            else:
-                return None
-        else:
-            return result.out[0]
+            raise NoSuchElementException('Failed waiting for element with {} in {}'
+                                         .format(locator, parent))
+        # wait_for returns NamedTuple, return first item from 'out', the WebElement
+        return result.out[0]
 
     def element(self, locator, *args, **kwargs):
         """Returns one :py:class:`selenium.webdriver.remote.webelement.WebElement`

testing/test_browser.py

@@ -2,7 +2,7 @@
 from __future__ import unicode_literals
 import pytest
 
-from widgetastic.browser import BrowserParentWrapper
+from widgetastic.browser import BrowserParentWrapper, WebElement
 from widgetastic.exceptions import NoSuchElementException, LocatorNotImplemented
 from widgetastic.widget import View, Text
 
@@ -65,7 +65,10 @@ def test_elements_check_visibility(browser):
 def test_wait_for_element_visible(browser):
     # Click on the button
     browser.click('#invisible_appear_button')
-    assert browser.wait_for_element('#invisible_appear_p', visible=True) is not None
+    try:
+        assert isinstance(browser.wait_for_element('#invisible_appear_p', visible=True), WebElement)
+    except NoSuchElementException:
+        pytest.fail('NoSuchElementException raised when webelement expected')
 
 
 def test_wait_for_element_visible_fail_except(browser):
@@ -75,13 +78,6 @@ def test_wait_for_element_visible_fail_except(browser):
         browser.wait_for_element('#invisible_appear_p', visible=True, timeout=1.5)
 
 
-def test_wait_for_element_visible_fail_none(browser):
-    # Click on the button
-    browser.click('#invisible_appear_button')
-    assert browser.wait_for_element(
-        '#invisible_appear_p', visible=True, timeout=1.5, exception=False) is None
-
-
 def test_element_only_invisible(browser):
     browser.element('#hello', check_visibility=False)