Implemented decorator pattern with notebook_loader

testbook/execute.py

@@ -6,15 +6,48 @@
 from testbook.client import TestbookNotebookClient
 
 
-@pytest.fixture
-def notebook_loader():
-    @contextmanager
-    def notebook_helper(nb_path, **kwargs):
-        with open(nb_path) as f:
+class notebook_loader:
+    def __init__(self, nb_path, prerun=None):
+        self.nb_path = nb_path
+        self.prerun = prerun
+
+        with open(self.nb_path) as f:
             nb = nbformat.read(f, as_version=4)
 
         client = TestbookNotebookClient(nb)
-        with client.setup_kernel(**kwargs):
-            yield client
 
-    return notebook_helper
+        if self.prerun is not None:
+            with client.setup_kernel():
+                client.execute_cell(self.prerun)
+
+        self.client = client
+
+    def _start_kernel(self):
+        if self.client.km is None:
+            self.client.start_kernel_manager()
+
+        if not self.client.km.has_kernel:
+            self.client.start_new_kernel_client()
+
+    def __enter__(self):
+        self._start_kernel()
+        return self.client
+
+    def __exit__(self, *args):
+        self.client._cleanup_kernel()
+
+    def __call__(self, func):
+        def wrapper():
+            self._start_kernel()
+
+            def inner(*args, **kwargs):
+                return func(self.client, *args, **kwargs)
+
+            self.client._cleanup_kernel()
+
+            inner.__name__ = func.__name__
+            inner.__doc__ = func.__doc__
+
+            return inner
+
+        return wrapper

testbook/tests/test_client.py

@@ -3,7 +3,7 @@
 from testbook import notebook_loader
 
 
-def test_get_cell_index(notebook_loader):
+def test_get_cell_index():
     with notebook_loader('testbook/tests/resources/foo.ipynb') as notebook:
         with pytest.raises(TypeError):
             notebook._get_cell_index(1)

testbook/tests/test_execute.py

@@ -1,7 +1,7 @@
 from testbook import notebook_loader
 
 
-def test_execute_cell(notebook_loader):
+def test_execute_cell():
     with notebook_loader('testbook/tests/resources/foo.ipynb') as notebook:
         notebook.execute_cell(1)
         assert notebook.cell_output_text(1) == 'hello world\n[1, 2, 3]\n'
@@ -10,10 +10,24 @@ def test_execute_cell(notebook_loader):
         assert notebook.cell_output_text(3) == 'foo\n'
 
 
-def test_execute_cell_tags(notebook_loader):
+def test_execute_cell_tags():
     with notebook_loader('testbook/tests/resources/foo.ipynb') as notebook:
         notebook.execute_cell('test1')
         assert notebook.cell_output_text('test1') == 'hello world\n[1, 2, 3]\n'
 
         notebook.execute_cell(['prepare_foo', 'execute_foo'])
         assert notebook.cell_output_text('execute_foo') == 'foo\n'
+
+
+@notebook_loader("testbook/tests/resources/foo.ipynb")
+def test_notebook(notebook):
+    notebook.execute_cell('test1')
+    assert notebook.cell_output_text('test1') == 'hello world\n[1, 2, 3]\n'
+
+    notebook.execute_cell(['prepare_foo', 'execute_foo'])
+    assert notebook.cell_output_text('execute_foo') == 'foo\n'
+
+
+@notebook_loader("testbook/tests/resources/foo.ipynb", prerun='test1')
+def test_notebook_with_prerun(notebook):
+    assert notebook.cell_output_text(1) == 'hello world\n[1, 2, 3]\n'

testbook/tests/test_inject.py

@@ -25,7 +25,7 @@ def helper_2(arg1, arg2):
         ),
     ],
 )
-def test_inject(helper_func, args, expected_text, notebook_loader):
+def test_inject(helper_func, args, expected_text):
 
     with notebook_loader('testbook/tests/resources/foo.ipynb') as notebook:
         notebook.inject(helper_func, args).assert_output_text(expected_text)
@@ -52,12 +52,12 @@ def foo(arg):
         ),
     ],
 )
-def test_inject_code_block(code_block, expected_text, notebook_loader):
+def test_inject_code_block(code_block, expected_text):
     with notebook_loader('testbook/tests/resources/foo.ipynb') as notebook:
         notebook.inject(code_block).assert_output_text(expected_text)
 
 
-def test_inject_raises_exception(notebook_loader):
+def test_inject_raises_exception():
     with notebook_loader('testbook/tests/resources/foo.ipynb') as notebook:
 
         values = [3, {'key': 'value'}, ['a', 'b', 'c'], (1, 2, 3), {1, 2, 3}]