unit test support for mcp notebook

Signed-off-by: Bryan Bednarski <[email protected]>

Author: bbednarski9

examples/notebooks/mcp_setup_and_integration.ipynb

@@ -303,7 +303,9 @@
     "    print(\"  - Enter a branch name (e.g., 'develop', 'main') to clone and install\")\n",
     "    print(\"=\" * 60)\n",
     "\n",
-    "    branch_name = input(\"\\nEnter your choice: \").strip()\n",
+    "    branch_name = os.getenv(\"NAT_NOTEBOOK_INSTALL_MODE\", \"\").strip()\n",
+    "    if not branch_name:\n",
+    "        branch_name = input(\"\\nEnter your choice: \").strip()\n",
     "\n",
     "    # Default to 'main' if empty\n",
     "    if not branch_name:\n",
@@ -499,22 +501,32 @@
    "id": "c8e1c22d",
    "metadata": {},
    "source": [
-    "Wait about 10 seconds before running the next cell...\n",
-    "\n",
     "The `--bg` command will run your server in the background of the notebook, so you might not see confirmation that the server is running. To see the startup logs, you can remove this flag and run the cell again. However, subsequent steps require the server to be run as a background process.\n",
     "\n",
     "Next, we will list the set of available tools from `simple_calculator` MCP server."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c2d15d0d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Wait about 10 seconds before running the next cell...\n",
+    "import time\n",
+    "time.sleep(10)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "afdef7a1",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# For streamable-http transport (default)\n",
-    "!nat mcp client tool list --url http://localhost:9901/mcp"
+    "%%bash\n",
+    "nat mcp client tool list --url http://localhost:9901/mcp"
    ]
   },
   {
@@ -572,11 +584,15 @@
    ]
   },
   {
-   "cell_type": "markdown",
-   "id": "9b013a00",
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "111c81a7",
    "metadata": {},
+   "outputs": [],
    "source": [
-    "Wait about 10 seconds before running the next cell..."
+    "# Wait about 10 seconds before running the next cell...\n",
+    "import time\n",
+    "time.sleep(10)"
    ]
   },
   {
@@ -629,11 +645,15 @@
    ]
   },
   {
-   "cell_type": "markdown",
-   "id": "664a9ed6",
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7ce2eebb",
    "metadata": {},
+   "outputs": [],
    "source": [
-    "Wait about 10 seconds before running the next cell..."
+    "# Wait about 10 seconds before running the next cell...\n",
+    "import time\n",
+    "time.sleep(10)"
    ]
   },
   {
@@ -954,11 +974,15 @@
    ]
   },
   {
-   "cell_type": "markdown",
-   "id": "feae920e",
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b7efca58",
    "metadata": {},
+   "outputs": [],
    "source": [
-    "Wait about 10 seconds before running the next cell..."
+    "# Wait about 10 seconds before running the next cell...\n",
+    "import time\n",
+    "time.sleep(10)"
    ]
   },
   {
@@ -1071,11 +1095,15 @@
    ]
   },
   {
-   "cell_type": "markdown",
-   "id": "b5d33205",
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7dc57c53",
    "metadata": {},
+   "outputs": [],
    "source": [
-    "Wait about 10 seconds before running the next cell..."
+    "# Wait about 10 seconds before running the next cell...\n",
+    "import time\n",
+    "time.sleep(10)"
    ]
   },
   {
@@ -1119,11 +1147,15 @@
    ]
   },
   {
-   "cell_type": "markdown",
-   "id": "4835e115",
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cc1de7ca",
    "metadata": {},
+   "outputs": [],
    "source": [
-    "Wait about 10 seconds before running the next cell..."
+    "# Wait about 10 seconds before running the next cell...\n",
+    "import time\n",
+    "time.sleep(10)"
    ]
   },
   {

examples/notebooks/tests/test_notebooks_e2e.py

@@ -13,9 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import shutil
 import subprocess
 from pathlib import Path
+import os
 
 import pytest
 
@@ -90,9 +90,6 @@ def _delete_other_files(notebooks_dir: Path):
     for file in _OTHER_FILES:
         file_path = notebooks_dir / file
         if file_path.exists():
-            # if file_path.is_dir():
-            #     shutil.rmtree(file_path)
-            # else:
             file_path.unlink()
 
 
@@ -118,9 +115,12 @@ def _run_notebook(notebook_path: Path, expected_packages: list[str], timeout_sec
         str(notebook_path.absolute()),
     ]
 
+    env = os.environ.copy()
+    env["NAT_NOTEBOOK_INSTALL_MODE"] = "local"
+
     # Ideally if the notebook times out we want jupyter to catch it and exit gracefully with the most informative error
     # possible. However in the potential situation where jupyter itself hangs, we add a 10s buffer to the timeout.
-    result = subprocess.run(cmd, check=False, capture_output=True, text=True, timeout=timeout_seconds + 10)
+    result = subprocess.run(cmd, check=False, capture_output=True, text=True, timeout=timeout_seconds + 10, env=env)
     assert result.returncode == 0, f"Notebook execution failed:\n{result.stderr}"
 
     for package in expected_packages: