mocks

Author: alex-treebeard

src/nbmake/nb_run.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Optional
+from typing import Any, Dict, Optional
 
 import nbformat
 from nbclient.client import (
@@ -60,6 +60,26 @@ def execute(
                 record_timing=True,
                 **extra_kwargs,
             )
+
+            async def apply_mocks(
+                cell: NotebookNode, cell_index: int, execute_reply: Dict[str, Any]
+            ):
+                if c.kc is None:
+                    raise Exception("there is no kernelclient")
+                mocks: Dict[str, Any] = (
+                    cell.get("metadata", {}).get("nbmake", {}).get("mock", {})
+                )
+                for v in mocks:
+                    if isinstance(mocks[v], str):
+                        out = await c.kc.execute_interactive(f"{v} = '{mocks[v]}'")
+                    else:
+                        out = await c.kc.execute_interactive(f"{v} = {mocks[v]}")
+
+                    if out["content"]["status"] != "ok":
+                        raise Exception(f"Failed to apply mock {v}\n\n{str(out)}")
+
+            c.on_cell_executed = apply_mocks
+
             c.execute(cwd=self.filename.parent)
         except CellExecutionError:
             error = self._get_error(nb)

tests/resources/mock.ipynb

@@ -0,0 +1,42 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "nbmake": {
+     "mock": {
+      "x": 2,
+      "y": "fish",
+      "z": {
+       "x": 42
+      }
+     }
+    }
+   },
+   "outputs": [],
+   "source": [
+    "x = 5\n",
+    "y = 'y'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert x == 2\n",
+    "assert y == \"fish\"\n",
+    "assert z == { \"x\": 42 }"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tests/test_nb_run.py

@@ -113,3 +113,9 @@ def test_when_raises_exc_tag_then_succeeds(self, testdir2: Testdir):
         run = NotebookRun(nb, 300)
         res: NotebookResult = run.execute()
         assert res.error == None
+
+    def test_when_mock_then_succeeds(self, testdir2: Testdir):
+        nb = Path(__file__).parent / "resources" / "mock.ipynb"
+        run = NotebookRun(nb, 300)
+        res: NotebookResult = run.execute()
+        assert res.error == None