From 7fb461ddf77af2bc590ed387eca23bd5d6f700e6 Mon Sep 17 00:00:00 2001
From: Thomas Kluyver <thomas@kluyver.me.uk>
Date: Mon, 11 Nov 2019 12:12:34 +0000
Subject: [PATCH] Selenium test fixture to open a notebook with prefilled
 content

This makes several tests shorter and (hopefully) more efficient
---
 notebook/tests/selenium/conftest.py           | 20 +++++++
 notebook/tests/selenium/test_buffering.py     | 18 ++----
 .../selenium/test_clipboard_multiselect.py    |  7 +--
 notebook/tests/selenium/test_deletecell.py    | 13 ++---
 .../tests/selenium/test_dualmode_clipboard.py | 19 +++---
 .../tests/selenium/test_dualmode_execute.py   | 17 ++----
 .../selenium/test_dualmode_insertcell.py      | 12 +---
 .../tests/selenium/test_find_and_replace.py   | 24 +++-----
 notebook/tests/selenium/test_markdown.py      | 31 +++++-----
 notebook/tests/selenium/test_merge_cells.py   | 22 +++----
 notebook/tests/selenium/test_multiselect.py   | 12 ++--
 .../tests/selenium/test_multiselect_toggle.py | 13 ++---
 .../tests/selenium/test_prompt_numbers.py     | 58 ++++++++++---------
 notebook/tests/selenium/test_save.py          |  2 +-
 notebook/tests/selenium/test_undelete.py      | 54 ++++++++---------
 notebook/tests/selenium/utils.py              | 27 ++++-----
 16 files changed, 157 insertions(+), 192 deletions(-)

diff --git a/notebook/tests/selenium/conftest.py b/notebook/tests/selenium/conftest.py
index 0228094f0..64cdfa23b 100644
--- a/notebook/tests/selenium/conftest.py
+++ b/notebook/tests/selenium/conftest.py
@@ -1,9 +1,12 @@
 import json
+import nbformat
+from nbformat.v4 import new_notebook, new_code_cell
 import os
 import pytest
 import requests
 from subprocess import Popen
 import sys
+from tempfile import mkstemp
 from testpath.tempdir import TemporaryDirectory
 import time
 from urllib.parse import urljoin
@@ -121,3 +124,20 @@ def notebook(authenticated_browser):
     tree_wh = authenticated_browser.current_window_handle
     yield Notebook.new_notebook(authenticated_browser)
     authenticated_browser.switch_to.window(tree_wh)
+
+@pytest.fixture
+def prefill_notebook(selenium_driver, notebook_server):
+    def inner(cells):
+        cells = [new_code_cell(c) if isinstance(c, str) else c
+                 for c in cells]
+        nb = new_notebook(cells=cells)
+        fd, path = mkstemp(dir=notebook_server['nbdir'], suffix='.ipynb')
+        with open(fd, 'w', encoding='utf-8') as f:
+            nbformat.write(nb, f)
+        fname = os.path.basename(path)
+        selenium_driver.get(
+            "{url}notebooks/{}?token={token}".format(fname, **notebook_server)
+        )
+        return Notebook(selenium_driver)
+
+    return inner
diff --git a/notebook/tests/selenium/test_buffering.py b/notebook/tests/selenium/test_buffering.py
index 4c2e748b7..4e4f45692 100644
--- a/notebook/tests/selenium/test_buffering.py
+++ b/notebook/tests/selenium/test_buffering.py
@@ -13,23 +13,22 @@ def wait_for_kernel_ready(notebook):
     wait_for_selector(notebook.browser, ".kernel_idle_icon")
 
 
-def kernels_buffer_without_conn(notebook):
+def test_kernels_buffer_without_conn(prefill_notebook):
     """Test that execution request made while disconnected is buffered."""
-    # Assert that cell executed while kernel is disconnected still
-    # executes when reconnected
+    notebook = prefill_notebook(["print(1 + 2)"])
+
     wait_for_kernel_ready(notebook)
     notebook.browser.execute_script("IPython.notebook.kernel.stop_channels();")
-    notebook.edit_cell(index=0, content="print(1 + 2)")
     notebook.execute_cell(0)
     notebook.browser.execute_script("IPython.notebook.kernel.reconnect();")
     wait_for_kernel_ready(notebook)
+
     assert wait_for_cell_text_output(notebook, 0) == "3"
-    notebook.delete_cell(0)
 
 
-def buffered_cells_execute_in_order(notebook):
+def test_buffered_cells_execute_in_order(prefill_notebook):
     """Test that buffered requests execute in order."""
-    notebook.append('k=1', 'k+=1', 'k*=3', 'print(k)')
+    notebook = prefill_notebook(['', 'k=1', 'k+=1', 'k*=3', 'print(k)'])
 
     # Repeated execution of cell queued up in the kernel executes
     # each execution request in order.
@@ -49,8 +48,3 @@ def buffered_cells_execute_in_order(notebook):
 
     # Check that current value of k is 7
     assert wait_for_cell_text_output(notebook, 4) == "7"
-
-
-def test_buffering(notebook):
-    kernels_buffer_without_conn(notebook)
-    buffered_cells_execute_in_order(notebook)
diff --git a/notebook/tests/selenium/test_clipboard_multiselect.py b/notebook/tests/selenium/test_clipboard_multiselect.py
index 360f016da..e75380c8e 100644
--- a/notebook/tests/selenium/test_clipboard_multiselect.py
+++ b/notebook/tests/selenium/test_clipboard_multiselect.py
@@ -2,11 +2,8 @@
 from selenium.webdriver.common.keys import Keys
 from .utils import wait_for_selector, wait_for_xpath
 
-def test_clipboard_multiselect(notebook):
-
-    # Cell contents setup
-    values = ['1', '2', '3', '4', '5a', '6b', '7c', '8d']
-    notebook.extend(values)
+def test_clipboard_multiselect(prefill_notebook):
+    notebook = prefill_notebook(['', '1', '2', '3', '4', '5a', '6b', '7c', '8d'])
 
     assert notebook.get_cells_contents() == ['', '1', '2', '3', '4', '5a', '6b', '7c', '8d']
     
diff --git a/notebook/tests/selenium/test_deletecell.py b/notebook/tests/selenium/test_deletecell.py
index 3b262fe94..8253b9e63 100644
--- a/notebook/tests/selenium/test_deletecell.py
+++ b/notebook/tests/selenium/test_deletecell.py
@@ -7,15 +7,12 @@ def remove_all_cells(notebook):
     for i in range(len(notebook.cells)):
         notebook.delete_cell(0)
 
-def test_delete_cells(notebook):
-    a = 'print("a")'
-    b = 'print("b")'
-    c = 'print("c")'
+INITIAL_CELLS = ['print("a")', 'print("b")', 'print("c")']
+
+def test_delete_cells(prefill_notebook):
+    a, b, c = INITIAL_CELLS
+    notebook = prefill_notebook(INITIAL_CELLS)
 
-    notebook.edit_cell(index=0, content=a)
-    notebook.append(b, c)
-    notebook.to_command_mode()
-    
     # Validate initial state
     assert notebook.get_cells_contents() == [a, b, c]
     for cell in range(0, 3):
diff --git a/notebook/tests/selenium/test_dualmode_clipboard.py b/notebook/tests/selenium/test_dualmode_clipboard.py
index bec4a4657..66d718d91 100644
--- a/notebook/tests/selenium/test_dualmode_clipboard.py
+++ b/notebook/tests/selenium/test_dualmode_clipboard.py
@@ -1,18 +1,13 @@
 """Test"""
 from .utils import shift, validate_dualmode_state
 
-def test_dualmode_clipboard(notebook):
-    a = 'print("a")'
-    notebook.append(a)
-    notebook.execute_cell(1)
+INITIAL_CELLS = ['', 'print("a")', 'print("b")', 'print("c")']
 
-    b = 'print("b")'
-    notebook.append(b)
-    notebook.execute_cell(2)
-
-    c = 'print("c")'
-    notebook.append(c)
-    notebook.execute_cell(3)
+def test_dualmode_clipboard(prefill_notebook):
+    notebook = prefill_notebook(INITIAL_CELLS)
+    _, a, b, c = INITIAL_CELLS
+    for i in range(1, 4):
+        notebook.execute_cell(i)
 
     #Copy/past/cut
     num_cells = len(notebook.cells)
@@ -56,4 +51,4 @@ def test_dualmode_clipboard(notebook):
     shift(notebook.browser, 'v') #Paste
     validate_dualmode_state(notebook, 'command', 0)
     assert notebook.get_cell_contents(0) == c #Cell 0 has the copied contents
-    assert len(notebook.cells) == num_cells+3 #A cell was added 
\ No newline at end of file
+    assert len(notebook.cells) == num_cells+3 #A cell was added
diff --git a/notebook/tests/selenium/test_dualmode_execute.py b/notebook/tests/selenium/test_dualmode_execute.py
index 4b646bb52..08bc3346f 100644
--- a/notebook/tests/selenium/test_dualmode_execute.py
+++ b/notebook/tests/selenium/test_dualmode_execute.py
@@ -4,19 +4,12 @@ from selenium.webdriver.common.keys import Keys
 
 from .utils import shift, cmdtrl, alt, validate_dualmode_state
 
+INITIAL_CELLS = ['', 'print("a")', 'print("b")', 'print("c")']
 
-def test_dualmode_execute(notebook):
-    a = 'print("a")'
-    notebook.append(a)
-    notebook.execute_cell(1)
-
-    b = 'print("b")'
-    notebook.append(b)
-    notebook.execute_cell(2)
-
-    c = 'print("c")'
-    notebook.append(c)
-    notebook.execute_cell(3)
+def test_dualmode_execute(prefill_notebook):
+    notebook = prefill_notebook(INITIAL_CELLS)
+    for i in range(1, 4):
+        notebook.execute_cell(i)
 
     #shift-enter
     #last cell in notebook
diff --git a/notebook/tests/selenium/test_dualmode_insertcell.py b/notebook/tests/selenium/test_dualmode_insertcell.py
index 4ab100ced..77f38ea2c 100644
--- a/notebook/tests/selenium/test_dualmode_insertcell.py
+++ b/notebook/tests/selenium/test_dualmode_insertcell.py
@@ -1,16 +1,10 @@
 from selenium.webdriver.common.keys import Keys
 from .utils import shift
 
-def test_insert_cell(notebook):
-    a = "print('a')"
-    b = "print('b')"
-    c = "print('c')"
+INITIAL_CELLS = ['print("a")', 'print("b")', 'print("c")']
 
-    notebook.edit_cell(index=0, content=a)
-    notebook.append(b, c)
-    notebook.to_command_mode()
-
-    assert notebook.get_cells_contents() == [a, b, c] 
+def test_insert_cell(prefill_notebook):
+    notebook = prefill_notebook(INITIAL_CELLS)
 
     notebook.to_command_mode()
     notebook.focus_cell(2)
diff --git a/notebook/tests/selenium/test_find_and_replace.py b/notebook/tests/selenium/test_find_and_replace.py
index 11143138c..07e2f1b00 100755
--- a/notebook/tests/selenium/test_find_and_replace.py
+++ b/notebook/tests/selenium/test_find_and_replace.py
@@ -1,24 +1,16 @@
-import pytest
+INITIAL_CELLS = ["hello", "hellohello", "abc", "ello"]
 
-
-def test_find_and_replace(notebook):
+def test_find_and_replace(prefill_notebook):
     """ test find and replace on all the cells """
-    cell_0, cell_1, cell_2, cell_3 = "hello", "hellohello", "abc", "ello"
-
-    find_str = "ello"                            # string to replace
-    replace_str = "foo"                          # string to replace to
+    notebook = prefill_notebook(INITIAL_CELLS)
 
-    # set the contents of the cells
-    notebook.add_cell(index=0, content=cell_0);
-    notebook.add_cell(index=1, content=cell_1);
-    notebook.add_cell(index=2, content=cell_2);
-    notebook.add_cell(index=3, content=cell_3);
+    find_str = "ello"
+    replace_str = "foo"
 
     # replace the strings
     notebook.find_and_replace(index=0, find_txt=find_str, replace_txt=replace_str)
 
     # check content of the cells
-    assert notebook.get_cell_contents(0) == cell_0.replace(find_str, replace_str)
-    assert notebook.get_cell_contents(1) == cell_1.replace(find_str, replace_str)
-    assert notebook.get_cell_contents(2) == cell_2.replace(find_str, replace_str)
-    assert notebook.get_cell_contents(3) == cell_3.replace(find_str, replace_str)
+    assert notebook.get_cells_contents() == [
+        s.replace(find_str, replace_str) for s in INITIAL_CELLS
+    ]
diff --git a/notebook/tests/selenium/test_markdown.py b/notebook/tests/selenium/test_markdown.py
index db4a44c9b..7d068ce4e 100644
--- a/notebook/tests/selenium/test_markdown.py
+++ b/notebook/tests/selenium/test_markdown.py
@@ -1,3 +1,4 @@
+from nbformat.v4 import new_markdown_cell
 
 def get_rendered_contents(nb):
     cl = ["text_cell", "render"]
@@ -9,21 +10,15 @@ def get_rendered_contents(nb):
             if x is not None]
 
             
-def test_markdown_cell(notebook):
-    nb = notebook
-    cell_text = ["# Foo",
-                 '**Bar**',
-                 '*Baz*',
-                 '```\nx = 1\n```',
-                 '```aaaa\nx = 1\n```',
-                 ]
-    expected_contents = ['<h1 id="Foo">Foo<a class="anchor-link" href="#Foo">¶</a></h1>',
-                         '<p><strong>Bar</strong></p>',
-                         '<p><em>Baz</em></p>', 
-                         '<pre><code>x = 1</code></pre>',
-                         '<pre><code class="cm-s-ipython language-aaaa">x = 1</code></pre>'
-                         ]
-    nb.append(*cell_text, cell_type="markdown")
-    nb.run_all()
-    rendered_contents = get_rendered_contents(nb)
-    assert rendered_contents == expected_contents
+def test_markdown_cell(prefill_notebook):
+    nb = prefill_notebook([new_markdown_cell(md) for md in [
+        '# Foo', '**Bar**', '*Baz*', '```\nx = 1\n```', '```aaaa\nx = 1\n```',
+    ]])
+
+    assert get_rendered_contents(nb) == [
+        '<h1 id="Foo">Foo<a class="anchor-link" href="#Foo">¶</a></h1>',
+        '<p><strong>Bar</strong></p>',
+        '<p><em>Baz</em></p>',
+        '<pre><code>x = 1</code></pre>',
+        '<pre><code class="cm-s-ipython language-aaaa">x = 1</code></pre>',
+    ]
diff --git a/notebook/tests/selenium/test_merge_cells.py b/notebook/tests/selenium/test_merge_cells.py
index 7311949fb..23c2b87e8 100644
--- a/notebook/tests/selenium/test_merge_cells.py
+++ b/notebook/tests/selenium/test_merge_cells.py
@@ -1,15 +1,17 @@
 """Tests the merge cell api."""
 
-def test_merge_cells(notebook):
-    # Add cells to notebook
-    a = "foo = 5"
-    b = "bar = 10"
-    c = "baz = 15"
-    d = "print(foo)"
-    e = "print(bar)"
-    f = "print(baz)"
-    notebook.edit_cell(index=0, content=a)
-    notebook.append(b, c, d, e, f)
+INITIAL_CELLS = [
+    "foo = 5",
+    "bar = 10",
+    "baz = 15",
+    "print(foo)",
+    "print(bar)",
+    "print(baz)",
+]
+
+def test_merge_cells(prefill_notebook):
+    notebook = prefill_notebook(INITIAL_CELLS)
+    a, b, c, d, e, f = INITIAL_CELLS
 
     # Before merging, there are 6 separate cells
     assert notebook.get_cells_contents() == [a, b, c, d, e, f]
diff --git a/notebook/tests/selenium/test_multiselect.py b/notebook/tests/selenium/test_multiselect.py
index 5ce6d4904..6e09b25c0 100644
--- a/notebook/tests/selenium/test_multiselect.py
+++ b/notebook/tests/selenium/test_multiselect.py
@@ -1,4 +1,8 @@
-def test_multiselect(notebook):
+INITIAL_CELLS = ['print("a")', 'print("b")', 'print("c")']
+
+def test_multiselect(prefill_notebook):
+    notebook = prefill_notebook(INITIAL_CELLS)
+
     def extend_selection_by(delta):
         notebook.browser.execute_script(
             "Jupyter.notebook.extend_selection_by(arguments[0]);", delta)
@@ -7,12 +11,6 @@ def test_multiselect(notebook):
         return notebook.browser.execute_script(
             "return Jupyter.notebook.get_selected_cells().length;")
 
-    a = 'print("a")'
-    b = 'print("b")'
-    c = 'print("c")'
-    notebook.edit_cell(index=0, content=a)
-    notebook.append(b, c)
-
     notebook.focus_cell(0)
     assert n_selected_cells() == 1
 
diff --git a/notebook/tests/selenium/test_multiselect_toggle.py b/notebook/tests/selenium/test_multiselect_toggle.py
index 245da169e..4ba6f667c 100644
--- a/notebook/tests/selenium/test_multiselect_toggle.py
+++ b/notebook/tests/selenium/test_multiselect_toggle.py
@@ -1,4 +1,8 @@
-def test_multiselect_toggle(notebook):
+INITIAL_CELLS = ['print("a")', 'print("b")', 'print("c")']
+
+def test_multiselect_toggle(prefill_notebook):
+    notebook = prefill_notebook(INITIAL_CELLS)
+
     def extend_selection_by(delta):
         notebook.browser.execute_script(
             "Jupyter.notebook.extend_selection_by(arguments[0]);", delta)
@@ -11,13 +15,6 @@ def test_multiselect_toggle(notebook):
         notebook.focus_cell(0)
         extend_selection_by(2)
 
-    a = 'print("a")'
-    b = 'print("b")'
-    c = 'print("c")'
-
-    notebook.edit_cell(index=0, content=a)
-    notebook.append(b, c)
-
     # Test that cells, which start off not collapsed, are collapsed after
     # calling the multiselected cell toggle.
     select_cells()
diff --git a/notebook/tests/selenium/test_prompt_numbers.py b/notebook/tests/selenium/test_prompt_numbers.py
index 57ff5bb18..42e27775e 100755
--- a/notebook/tests/selenium/test_prompt_numbers.py
+++ b/notebook/tests/selenium/test_prompt_numbers.py
@@ -1,28 +1,30 @@
-import os
-import pytest
-import time
-
-# selenium test version for 'prompt_numbers.js' 
-
-def get_prompt(nb, index):
-    cell = nb.cells[0]
-    return cell.find_element_by_class_name('input').find_element_by_class_name('input_prompt').get_attribute('innerHTML').strip() 
-
-def set_prompt(nb, index, value):
-    nb.set_cell_input_prompt(index, value)
-
-def test_prompt_numbers(notebook):
-    cell_index = 0
-    a = 'print("a")'
-    notebook.edit_cell(index=cell_index, content=a)
-    assert get_prompt(notebook, cell_index) == "<bdi>In</bdi>&nbsp;[&nbsp;]:"
-    set_prompt(notebook, cell_index, 2);
-    assert get_prompt(notebook, cell_index) == "<bdi>In</bdi>&nbsp;[2]:"
-    set_prompt(notebook, cell_index, 0);
-    assert get_prompt(notebook, cell_index) == "<bdi>In</bdi>&nbsp;[0]:"
-    set_prompt(notebook, cell_index, "'*'");
-    assert get_prompt(notebook, cell_index) == "<bdi>In</bdi>&nbsp;[*]:"
-    set_prompt(notebook, cell_index, "undefined");
-    assert get_prompt(notebook, cell_index) == "<bdi>In</bdi>&nbsp;[&nbsp;]:"
-    set_prompt(notebook, cell_index, "null");
-    assert get_prompt(notebook, cell_index) == "<bdi>In</bdi>&nbsp;[&nbsp;]:"
+
+def test_prompt_numbers(prefill_notebook):
+    notebook = prefill_notebook(['print("a")'])
+
+    def get_prompt():
+        return (
+            notebook.cells[0].find_element_by_class_name('input')
+            .find_element_by_class_name('input_prompt')
+            .get_attribute('innerHTML').strip()
+        )
+
+    def set_prompt(value):
+        notebook.set_cell_input_prompt(0, value)
+
+    assert get_prompt() == "<bdi>In</bdi>&nbsp;[&nbsp;]:"
+
+    set_prompt(2)
+    assert get_prompt() == "<bdi>In</bdi>&nbsp;[2]:"
+
+    set_prompt(0)
+    assert get_prompt() == "<bdi>In</bdi>&nbsp;[0]:"
+
+    set_prompt("'*'")
+    assert get_prompt() == "<bdi>In</bdi>&nbsp;[*]:"
+
+    set_prompt("undefined")
+    assert get_prompt() == "<bdi>In</bdi>&nbsp;[&nbsp;]:"
+
+    set_prompt("null")
+    assert get_prompt() == "<bdi>In</bdi>&nbsp;[&nbsp;]:"
diff --git a/notebook/tests/selenium/test_save.py b/notebook/tests/selenium/test_save.py
index 2ce609a5a..3ed7721f3 100644
--- a/notebook/tests/selenium/test_save.py
+++ b/notebook/tests/selenium/test_save.py
@@ -2,7 +2,7 @@
 """
 
 from urllib.parse import quote
-from .utils import wait_for_selector, new_window
+from .utils import wait_for_selector
 
 promise_js = """
 var done = arguments[arguments.length - 1];
diff --git a/notebook/tests/selenium/test_undelete.py b/notebook/tests/selenium/test_undelete.py
index 65b618c5a..d9823cda5 100644
--- a/notebook/tests/selenium/test_undelete.py
+++ b/notebook/tests/selenium/test_undelete.py
@@ -1,100 +1,92 @@
 from selenium.webdriver.common.keys import Keys
 from .utils import shift
 
-def get_cells_contents(nb):
-    JS = 'return Jupyter.notebook.get_cells().map(function(c) {return c.get_text();})'
-    return nb.browser.execute_script(JS)
-
 def undelete(nb):
     nb.browser.execute_script('Jupyter.notebook.undelete_cell();')
 
-def test_undelete_cells(notebook):
-    a = 'print("a")'
-    b = 'print("b")'
-    c = 'print("c")'
-    d = 'print("d")'
+INITIAL_CELLS = ['print("a")', 'print("b")', 'print("c")', 'print("d")']
 
-    notebook.edit_cell(index=0, content=a)
-    notebook.append(b, c, d)
-    notebook.to_command_mode()
+def test_undelete_cells(prefill_notebook):
+    notebook = prefill_notebook(INITIAL_CELLS)
+    a, b, c, d = INITIAL_CELLS
 
     # Verify initial state
-    assert get_cells_contents(notebook) == [a, b, c, d]
+    assert notebook.get_cells_contents() == [a, b, c, d]
 
     # Delete cells [1, 2]
     notebook.focus_cell(1)
     shift(notebook.browser, Keys.DOWN)
     notebook.current_cell.send_keys('dd')
-    assert get_cells_contents(notebook) == [a, d]
+    assert notebook.get_cells_contents() == [a, d]
 
     # Delete new cell 1 (which contains d)
     notebook.focus_cell(1)
     notebook.current_cell.send_keys('dd')
-    assert get_cells_contents(notebook) == [a]
+    assert notebook.get_cells_contents() == [a]
 
     # Undelete d
     undelete(notebook)
-    assert get_cells_contents(notebook) == [a, d]
+    assert notebook.get_cells_contents() == [a, d]
 
     # Undelete b, c
     undelete(notebook)
-    assert get_cells_contents(notebook) == [a, b, c, d]
+    assert notebook.get_cells_contents() == [a, b, c, d]
 
     # Nothing more to undelete
     undelete(notebook)
-    assert get_cells_contents(notebook) == [a, b, c, d]
+    assert notebook.get_cells_contents() == [a, b, c, d]
 
     # Delete first two cells and restore
     notebook.focus_cell(0)
     shift(notebook.browser, Keys.DOWN)
     notebook.current_cell.send_keys('dd')
-    assert get_cells_contents(notebook) == [c, d]
+    assert notebook.get_cells_contents() == [c, d]
     undelete(notebook)
-    assert get_cells_contents(notebook) == [a, b, c, d]
+    assert notebook.get_cells_contents() == [a, b, c, d]
 
     # Delete last two cells and restore
     notebook.focus_cell(-1)
     shift(notebook.browser, Keys.UP)
     notebook.current_cell.send_keys('dd')
-    assert get_cells_contents(notebook) == [a, b]
+    assert notebook.get_cells_contents() == [a, b]
     undelete(notebook)
-    assert get_cells_contents(notebook) == [a, b, c, d]
+    assert notebook.get_cells_contents() == [a, b, c, d]
 
     # Merge cells [1, 2], restore the deleted one
     bc = b + "\n\n" + c
     notebook.focus_cell(1)
     shift(notebook.browser, 'j')
     shift(notebook.browser, 'm')
-    assert get_cells_contents(notebook) == [a, bc, d]
+    assert notebook.get_cells_contents() == [a, bc, d]
     undelete(notebook)
-    assert get_cells_contents(notebook) == [a, bc, c, d]
+    assert notebook.get_cells_contents() == [a, bc, c, d]
 
     # Merge cells [2, 3], restore the deleted one
     cd = c + "\n\n" + d
     notebook.focus_cell(-1)
     shift(notebook.browser, 'k')
     shift(notebook.browser, 'm')
-    assert get_cells_contents(notebook) == [a, bc, cd]
+    assert notebook.get_cells_contents() == [a, bc, cd]
     undelete(notebook)
-    assert get_cells_contents(notebook) == [a, bc, cd, d]
+    assert notebook.get_cells_contents() == [a, bc, cd, d]
 
     # Reset contents to [a, b, c, d] --------------------------------------
     notebook.edit_cell(index=1, content=b)
     notebook.edit_cell(index=2, content=c)
-    assert get_cells_contents(notebook) == [a, b, c, d]
+    assert notebook.get_cells_contents() == [a, b, c, d]
 
     # Merge cell below, restore the deleted one
     ab = a + "\n\n" + b
     notebook.focus_cell(0)
     notebook.browser.execute_script("Jupyter.notebook.merge_cell_below();")
-    assert get_cells_contents(notebook) == [ab, c, d]
+    assert notebook.get_cells_contents() == [ab, c, d]
     undelete(notebook)
-    assert get_cells_contents(notebook) == [ab, b, c, d]
+    assert notebook.get_cells_contents() == [ab, b, c, d]
 
     # Merge cell above, restore the deleted one
     cd = c + "\n\n" + d
     notebook.focus_cell(-1)
     notebook.browser.execute_script("Jupyter.notebook.merge_cell_above();")
-    assert get_cells_contents(notebook) == [ab, b, cd]
+    assert notebook.get_cells_contents() == [ab, b, cd]
     undelete(notebook)
-    assert get_cells_contents(notebook) == [ab, b, c, cd]
+    assert notebook.get_cells_contents() == [ab, b, c, cd]
diff --git a/notebook/tests/selenium/utils.py b/notebook/tests/selenium/utils.py
index 4d9b7bc28..304b9c1b9 100644
--- a/notebook/tests/selenium/utils.py
+++ b/notebook/tests/selenium/utils.py
@@ -1,8 +1,6 @@
 import os
-import time
 from selenium.webdriver import ActionChains
 from selenium.webdriver.common.by import By
-from selenium.common.exceptions import WebDriverException
 from selenium.webdriver.common.keys import Keys
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
@@ -99,6 +97,7 @@ class Notebook:
     
     def __init__(self, browser):
         self.browser = browser
+        self._wait_for_start()
         self.disable_autosave_and_onbeforeunload()
         
     def __len__(self):
@@ -121,6 +120,13 @@ class Notebook:
     def __iter__(self):
         return (cell for cell in self.cells)
 
+    def _wait_for_start(self):
+        """Wait until the notebook interface is loaded and the kernel started"""
+        wait_for_selector(self.browser, '.cell')
+        WebDriverWait(self.browser, 10).until(
+            lambda drvr: self.is_kernel_running()
+        )
+
     @property
     def body(self):
         return self.browser.find_element_by_tag_name("body")
@@ -315,12 +321,9 @@ class Notebook:
 
     @classmethod
     def new_notebook(cls, browser, kernel_name='kernel-python3'):
-        with new_window(browser, selector=".cell"):
+        with new_window(browser):
             select_kernel(browser, kernel_name=kernel_name)
-        wait = WebDriverWait(browser, 10)
-        nb = cls(browser)
-        wait.until(lambda driver: nb.is_kernel_running())
-        return nb
+        return cls(browser)
 
 
 def select_kernel(browser, kernel_name='kernel-python3'):
@@ -335,23 +338,19 @@ def select_kernel(browser, kernel_name='kernel-python3'):
 
 
 @contextmanager
-def new_window(browser, selector=None):
+def new_window(browser):
     """Contextmanager for switching to & waiting for a window created. 
     
     This context manager gives you the ability to create a new window inside 
     the created context and it will switch you to that new window.
     
-    If you know a CSS selector that can be expected to appear on the window, 
-    then this utility can wait on that selector appearing on the page before
-    releasing the context.
-    
     Usage example:
     
         from notebook.tests.selenium.utils import new_window, Notebook
         
         ⋮ # something that creates a browser object
         
-        with new_window(browser, selector=".cell"):
+        with new_window(browser):
             select_kernel(browser, kernel_name=kernel_name)
         nb = Notebook(browser)
 
@@ -363,8 +362,6 @@ def new_window(browser, selector=None):
     if not new_window_handles:
         raise Exception("No new windows opened during context")
     browser.switch_to.window(new_window_handles[0])
-    if selector is not None:
-        wait_for_selector(browser, selector)
 
 def shift(browser, k):
     """Send key combination Shift+(k)"""