wait_for_page_load.py

class wait_for_page_load(object):

    def __init__(self, browser, url):
        self.browser = browser
        self.url = url

    def __enter__(self):
        self.old_page = self.browser.find_element_by_tag_name('html')

    def page_has_loaded(self):
        new_page = self.browser.find_element_by_tag_name('html')
        return new_page.id != self.old_page.id

    def wait_for(self, condition_function):
        start_time = time.time()
        while time.time() < start_time + 10:
            if condition_function():
                self.browser.implicitly_wait(10)
                return True
            else:
                time.sleep(0.3)
        log.debug(
            'Timeout waiting for {0}, wasted {1} sec.'.format(
                self.url,
                time.time() - start_time,
            )
        )

    def __exit__(self, *_):
        self.wait_for(self.page_has_loaded)