Return exception arg for Browser.wait_for_element

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,
+            self, locator, parent=None, visible=False, timeout=5, delay=0.2, exception=True,
             ensure_page_safe=False):
         """Wait for presence or visibility of elements specified by a locator.
 
@@ -283,11 +283,13 @@ 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 None.
             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.
+            to params. ``None`` if not found and ``exception=False``.
 
         Raises:
             :py:class:`selenium.common.exceptions.NoSuchElementException` if element not found.
@@ -309,8 +311,11 @@ def _element_lookup():
                               fail_condition=lambda elements: not bool(elements),
                               fail_func=self.plugin.ensure_page_safe if ensure_page_safe else None)
         except TimedOutError:
-            raise NoSuchElementException('Failed waiting for element with {} in {}'
-                                         .format(locator, parent))
+            if exception:
+                raise NoSuchElementException('Failed waiting for element with {} in {}'
+                                             .format(locator, parent))
+            else:
+                return None
         # wait_for returns NamedTuple, return first item from 'out', the WebElement
         return result.out[0]
 

testing/test_browser.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 from __future__ import unicode_literals
+
 import pytest
 
 from widgetastic.browser import BrowserParentWrapper, WebElement
@@ -71,11 +72,21 @@ def test_wait_for_element_visible(browser):
         pytest.fail('NoSuchElementException raised when webelement expected')
 
 
-def test_wait_for_element_visible_fail_except(browser):
-    # Click on the button
+@pytest.mark.parametrize('exception', [True, False], ids=['with_exception', 'without_exception'])
+def test_wait_for_element_exception_control(browser, exception):
+    # Click on the button, element will not appear
     browser.click('#invisible_appear_button')
-    with pytest.raises(NoSuchElementException):
-        browser.wait_for_element('#invisible_appear_p', visible=True, timeout=1.5)
+    wait_for_args = dict(
+        locator='#invisible_appear_p',
+        visible=True,
+        timeout=1.5,
+        exception=exception
+    )
+    if exception:
+        with pytest.raises(NoSuchElementException):
+            browser.wait_for_element(**wait_for_args)
+    else:
+        assert browser.wait_for_element(**wait_for_args) is None
 
 
 def test_element_only_invisible(browser):