From 966edac4ef6fc01618c070e29b05280e26592327 Mon Sep 17 00:00:00 2001
From: Alex Remedios <alex@treebeard.io>
Date: Fri, 3 Dec 2021 16:24:44 +0000
Subject: [PATCH] kernel flag (#58)

---
 README.md                   | 15 +++++++++++----
 src/nbmake/nb_run.py        |  7 +++++++
 src/nbmake/pytest_items.py  |  1 +
 src/nbmake/pytest_plugin.py |  8 +++++++-
 tests/test_pytest_plugin.py | 18 ++++++++++++++++++
 5 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 788002b..c33f604 100644
--- a/README.md
+++ b/README.md
@@ -25,15 +25,15 @@ pytest --nbmake **/*ipynb
 
 ## Allow errors and Configure Cell Timeouts
 
-You can configure the notebook run timeout with a the following pytest flag:
+You can configure the cell timeout with the following pytest flag:
 
 ```sh
-pytest --nbmake --nbmake-timeout=3000 # allows each notebook 3000 seconds to finish
+pytest --nbmake --nbmake-timeout=3000 # allows each cell 3000 seconds to finish
 ```
 
 Each notebook can also be separately overidden to allow errors and fail if running exceeds a timeout.
 
-This configuration must be placed in the notebook's top-level metadata (not cell-level metadata).
+This configuration must be placed in the notebook's **top-level metadata** (not cell-level metadata).
 
 Your notebook should look like this:
 
@@ -50,10 +50,17 @@ Your notebook should look like this:
 }
 ```
 
+## Override Notebook Kernels when Testing
+
+Regardless of the kernel configured in the notebook JSON, you can force nbmake to use a specific kernel when testing:
+
+```
+pytest --nbmake --nbmake-kernel=mycustomkernel
+```
 
 ## Add Missing Jupyter Kernel to Your CI Environment
 
-If you are using a kernel name other than the default ‘python3’. You will see an error message when executing your notebooks in a fresh CI environment: `Error - No such kernel: 'mycustomkernel'`
+If you are not using the flag above and are using a kernel name other than the default ‘python3’, you will see an error message when executing your notebooks in a fresh CI environment: `Error - No such kernel: 'mycustomkernel'`
 
 Use ipykernel to install the custom kernel:
 
diff --git a/src/nbmake/nb_run.py b/src/nbmake/nb_run.py
index c0df0b1..2d39273 100644
--- a/src/nbmake/nb_run.py
+++ b/src/nbmake/nb_run.py
@@ -23,10 +23,12 @@ def __init__(
         filename: Path,
         default_timeout: int,
         verbose: bool = False,
+        kernel: Optional[str] = None,
     ) -> None:
         self.filename = filename
         self.verbose = verbose
         self.default_timeout = default_timeout
+        self.kernel = kernel
 
     def execute(
         self,
@@ -46,12 +48,17 @@ def execute(
 
         error: Optional[NotebookError] = None
 
+        extra_kwargs = {}
+        if self.kernel:
+            extra_kwargs["kernel_name"] = self.kernel
+
         try:
             c = NotebookClient(
                 nb,
                 timeout=timeout,
                 allow_errors=allow_errors,
                 record_timing=True,
+                **extra_kwargs,
             )
             c.execute(cwd=self.filename.parent)
         except CellExecutionError:
diff --git a/src/nbmake/pytest_items.py b/src/nbmake/pytest_items.py
index 342781b..01d0b4c 100644
--- a/src/nbmake/pytest_items.py
+++ b/src/nbmake/pytest_items.py
@@ -46,6 +46,7 @@ def runtest(self):
             source,
             option.nbmake_timeout,
             verbose=bool(option.verbose),
+            kernel=option.nbmake_kernel,
         )
 
         res: NotebookResult = run.execute()
diff --git a/src/nbmake/pytest_plugin.py b/src/nbmake/pytest_plugin.py
index d5130d7..831bb8f 100644
--- a/src/nbmake/pytest_plugin.py
+++ b/src/nbmake/pytest_plugin.py
@@ -29,10 +29,16 @@ def pytest_addoption(parser: Any):
     group.addoption(
         "--nbmake-timeout",
         action="store",
-        help="Sets the default timeout for a notebook (seconds)",
+        help="Sets the default cell timeout (seconds)",
         default=300,
         type=int,
     )
+    group.addoption(
+        "--nbmake-kernel",
+        action="store",
+        help="Overrides the kernel used for all notebooks",
+        type=str,
+    )
 
 
 def pytest_collect_file(path: str, parent: Any) -> Optional[Any]:
diff --git a/tests/test_pytest_plugin.py b/tests/test_pytest_plugin.py
index 10e7652..f71265f 100644
--- a/tests/test_pytest_plugin.py
+++ b/tests/test_pytest_plugin.py
@@ -163,3 +163,21 @@ def test_when_explicit_metadata_then_ignore_timeout(testdir: Testdir):
 
     hook_recorder = testdir.inline_run("--nbmake", "--nbmake-timeout=1")
     assert hook_recorder.ret == ExitCode.OK
+
+
+def test_when_kernel_passed_then_override(testdir: Testdir):
+    write_nb(
+        passing_nb,
+        Path(f"x.ipynb"),
+        metadata={
+            "kernelspec": {
+                "display_name": "Python 3",
+                "language": "python",
+                "name": "blah",
+            }
+        },
+    )
+
+    hook_recorder = testdir.inline_run("--nbmake", "--nbmake-kernel=python3")
+
+    assert hook_recorder.ret == ExitCode.OK