try:
    from urllib import unquote # Python 2
    from urlparse import urljoin
except ImportError:
    from urllib.parse import urljoin, unquote # Python 3

from selenium.webdriver.common.by import By
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.ui import WebDriverWait
from selenium.common.exceptions import TimeoutException


class BaseTestFormgrade(object):
    """Do NOT add any test cases to this base class. This class is for helper
    functions ONLY. If you add test cases, it will break the tests in
    test_formgrader.py and/or test_gradebook.py!

    """

    def formgrade_url(self, url=""):
        return urljoin(self.manager.base_formgrade_url, url).rstrip("/")

    def notebook_url(self, url=""):
        return urljoin(self.manager.base_notebook_url, url).rstrip("/")

    def _check_url(self, url):
        if not url.startswith("http"):
            url = self.formgrade_url(url)
        assert unquote(self.browser.current_url.rstrip("/")) == url

    def _check_breadcrumbs(self, *breadcrumbs):
        # check that breadcrumbs are correct
        elements = self.browser.find_elements_by_css_selector("ul.breadcrumb li")
        assert tuple([e.text for e in elements]) == breadcrumbs

        # check that the active breadcrumb is correct
        element = self.browser.find_element_by_css_selector("ul.breadcrumb li.active")
        assert element.text == breadcrumbs[-1]

    def _click_link(self, link_text, partial=False):
        if partial:
            element = self.browser.find_element_by_partial_link_text(link_text)
        else:
            element = self.browser.find_element_by_link_text(link_text)
        element.click()

    def _wait_for_element(self, element_id, time=30):
        return WebDriverWait(self.browser, time).until(
            EC.presence_of_element_located((By.ID, element_id))
        )

    def _wait_for_gradebook_page(self, url):
        self._wait_for_element("gradebook")
        self._check_url(url)

    def _get(self, url, retries=5):
        for _ in range(retries):
            try:
                self.browser.get(url)
            except TimeoutException:
                print("Failed to load '{}', trying again...".format(url))
                continue
            else:
                return

        raise TimeoutException

    def _load_gradebook_page(self, url):
        self._get(self.formgrade_url(url))
        self._wait_for_gradebook_page(url)

    def _wait_for_notebook_page(self, url):
        self._wait_for_element("notebook-container")
        self._check_url(url)

    def _wait_for_formgrader(self, url):
        self._wait_for_element("notebook-container")
        page_loaded = lambda browser: browser.execute_script(
            """
            if (!(typeof MathJax !== "undefined" && MathJax !== undefined && MathJax.loaded)) {
                return false;
            }

            if (!(typeof formgrader !== "undefined" && formgrader !== undefined)) {
                return false;
            }

            if (!(formgrader.grades !== undefined && formgrader.grades.loaded)) {
                return false;
            }

            if (!(formgrader.comments !== undefined && formgrader.comments.loaded)) {
                return false;
            }

            if (!(typeof autosize !== "undefined" && autosize !== undefined)) {
                return false;
            }

            return true;
            """)
        WebDriverWait(self.browser, 30).until(page_loaded)
        self._check_url(url)

@@ -2052,45 +2052,8 @@
 5):%0A
-        for _ in range(retries):%0A    
     
@@ -2069,28 +2069,24 @@
             
-  
-  
 self.browser
@@ -2091,28 +2091,24 @@
 er.get(url)%0A
-    
         exce
@@ -2128,16 +2128,85 @@
 eption:%0A
+            if retries == 0:%0A                raise%0A            else:%0A
         
@@ -2291,89 +2291,43 @@
     
-continue%0A            else:%0A                r
+self._g
 et
+(
 ur
-n%0A%0A        raise TimeoutException
+l, retries=retries - 1)
 %0A%0A  
@@ -2609,32 +2609,43 @@
 grader(self, url
+, retries=5
 ):%0A        self.
@@ -3457,16 +3457,33 @@
    %22%22%22)%0A
+        try:%0A    
         
@@ -3533,16 +3533,330 @@
 loaded)%0A
+        except TimeoutException:%0A            if retries == 0:%0A                raise%0A            else:%0A                print(%22Failed to load formgrader (url '%7B%7D'), trying again...%22.format(url))%0A                self._get(self.browser.current_url)%0A                self._wait_for_formgrader(url, retries=retries - 1)%0A%0A
         

Add retries to _wait_for_formgrader