Take into account lab as a potential endpoint for the driver

notebook/tests/selenium/quick_selenium.py

@@ -2,15 +2,14 @@ from selenium.webdriver import Firefox
 from notebook.notebookapp import list_running_servers
 
 
-def quick_driver():
+def quick_driver(lab=False):
     """Quickly create a selenium driver pointing at an active noteboook server.
 
-    Usage example
+    Usage example:
     
-    from inside the selenium test directory:
-    
-        import quick_selenium, test_markdown, utils
-        nb = utils.Notebook(test_markdown.notebook(quick_selenium.quick_driver()))
+        from notebook.tests.selenium.quick_selenium import quick_driver
+        driver = quick_driver
+        
     """
     try:
         server = list(list_running_servers())[0]
@@ -19,4 +18,10 @@ def quick_driver():
     driver = Firefox()
     auth_url = '{url}?token={token}'.format(**server)
     driver.get(auth_url)
-    return driver
+
+    # If this redirects us to a lab page and we don't want that;
+    # then we need to redirect ourselves to the classic notebook view
+    if driver.current_url.endswith('/lab') and not lab:
+        driver.get(driver.current_url.rstrip('lab')+'tree')
+    return driver
+