Implemented decorator pattern with notebook_loader

.gitignore

@@ -130,3 +130,4 @@ dmypy.json
 
 # Editor configs
 .vscode/
+.idea/

testbook/exceptions.py

@@ -1,3 +1,4 @@
 class CellTagNotFoundError(Exception):
     """Raised when cell tag is not declared in notebook"""
+
     pass

testbook/execute.py

@@ -6,15 +6,42 @@
 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(*args, **kwargs):
+            with self.client.setup_kernel():
+                func(self.client, *args, **kwargs)
+
+        wrapper.__name__ = func.__name__
+        wrapper.__doc__ = func.__doc__
+
+        return wrapper

testbook/testbooknode.py

@@ -3,7 +3,7 @@
 
 class TestbookNode(NotebookNode):
     """
-    Extends `NodebookNode` to perform assertions
+    Extends `NotebookNode` to perform assertions
     """
 
     def __init__(self, *args, **kw):

testbook/tests/test_client.py

@@ -5,7 +5,7 @@
 
 
 @pytest.mark.parametrize("get_call_index_args, expected_result", [(2, 2), ('hello', 1)])
-def test_get_cell_index(get_call_index_args, expected_result, notebook_loader):
+def test_get_cell_index(get_call_index_args, expected_result):
     with notebook_loader('testbook/tests/resources/inject.ipynb') as notebook:
         assert notebook._get_cell_index(get_call_index_args) == expected_result
 
@@ -14,7 +14,7 @@ def test_get_cell_index(get_call_index_args, expected_result, notebook_loader):
     "get_call_index_args, expected_error",
     [([1, 2, 3], TypeError), ('non-existent-tag', CellTagNotFoundError)],
 )
-def test_get_cell_index_raises_error(get_call_index_args, expected_error, notebook_loader):
+def test_get_cell_index_raises_error(get_call_index_args, expected_error):
     with notebook_loader('testbook/tests/resources/inject.ipynb') as notebook:
         with pytest.raises(expected_error):
             notebook._get_cell_index(get_call_index_args)

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:
         assert notebook.inject(helper_func, args).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:
         assert notebook.inject(code_block).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}]
@@ -95,6 +95,6 @@ def test_inject_raises_exception(notebook_loader):
         ),
     ],
 )
-def test_inject_with_prerun(prerun, code_block, expected_text, notebook_loader):
+def test_inject_with_prerun(prerun, code_block, expected_text):
     with notebook_loader('testbook/tests/resources/inject.ipynb') as notebook:
         assert notebook.inject(code_block, prerun=prerun).output_text == expected_text