Added test of remove_tagged_cells in execute_notebook

nteract · Nov 9, 2023 · 332509e · 332509e
1 parent 7ad32a0
commit 332509e
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 2 deletions.
diff --git a/papermill/tests/notebooks/simple_with_tags.ipynb b/papermill/tests/notebooks/simple_with_tags.ipynb
@@ -22,7 +22,7 @@
    "id": "7979d49a-abb1-4815-9534-ad76e4505b56",
    "metadata": {
     "tags": [
-     "skipcell"
+     "assigncell"
     ]
    },
    "outputs": [],
@@ -35,7 +35,9 @@
    "execution_count": null,
    "id": "778581dd-f385-4039-be97-4050615fa271",
    "metadata": {
-    "tags": []
+    "tags": [
+     "printcell"
+    ]
    },
    "outputs": [],
    "source": [

diff --git a/papermill/tests/test_execute.py b/papermill/tests/test_execute.py
@@ -135,6 +135,49 @@ def test_prepare_only(self):
                 ['# Parameters', r'foo = "do\\ not\\ crash"', ''],
             )
 
+    def test_remove_tagged_cells(self):
+        notebook_name = 'simple_with_tags.ipynb'
+
+        # Default case, no cells are skipped
+        nb_test_executed_fname = os.path.join(self.test_dir, 'output_{}'.format(notebook_name))
+        execute_notebook(get_notebook_path(notebook_name), nb_test_executed_fname, {})
+        output_nb = load_notebook_node(nb_test_executed_fname)
+        assert len(output_nb.cells) == 4
+
+        # If a nonexistent tag is specified, no cells are skipped
+        nb_test_executed_fname = os.path.join(self.test_dir, 'output_{}'.format(notebook_name))
+        execute_notebook(
+            get_notebook_path(notebook_name),
+            nb_test_executed_fname,
+            {},
+            remove_tagged_cells="nonexistent",
+        )
+        output_nb = load_notebook_node(nb_test_executed_fname)
+        assert len(output_nb.cells) == 4
+
+        # If cells with the 'printcell' tag are skipped, the output notebook is missing one cell
+        nb_test_executed_fname = os.path.join(self.test_dir, 'output_{}'.format(notebook_name))
+        execute_notebook(
+            get_notebook_path(notebook_name),
+            nb_test_executed_fname,
+            {},
+            remove_tagged_cells="printcell",
+        )
+        output_nb = load_notebook_node(nb_test_executed_fname)
+        assert len(output_nb.cells) == 3
+
+        # If cells with the 'assigncell' tag are skipped, the execution raises an error
+        nb_test_executed_fname = os.path.join(self.test_dir, 'output_{}'.format(notebook_name))
+        with self.assertRaises(PapermillExecutionError):
+            execute_notebook(
+                get_notebook_path(notebook_name),
+                nb_test_executed_fname,
+                {},
+                remove_tagged_cells="assigncell",
+            )
+        output_nb = load_notebook_node(nb_test_executed_fname)
+        self.assertEqual(output_nb.cells[4].outputs[0]["evalue"], "name 'a' is not defined")
+
 
 class TestBrokenNotebook1(unittest.TestCase):
     def setUp(self):