example tests: remove checkpoints after test, minor adjustments

Author: vpratz

tests/test_examples/test_examples.py

@@ -3,12 +3,12 @@
 from tests.utils import run_notebook
 
 
+@pytest.mark.skip(reason="requires setting up Stan")
 @pytest.mark.slow
 def test_bayesian_experimental_design(examples_path):
     run_notebook(examples_path / "Bayesian_Experimental_Design.ipynb")
 
 
-@pytest.mark.skip(reason="Requires setting up Stan.")
 @pytest.mark.slow
 def test_from_abc_to_bayesflow(examples_path):
     run_notebook(examples_path / "From_ABC_to_BayesFlow.ipynb")

tests/utils/jupyter.py

@@ -2,15 +2,26 @@
 from nbconvert.preprocessors import ExecutePreprocessor
 
 from pathlib import Path
+import shutil
 
 
 def run_notebook(path):
+    path = Path(path)
+    checkpoint_path = path.parent / "checkpoints"
+    # only clean up if the directory did not exist before the test
+    cleanup_checkpoints = not checkpoint_path.exists()
     with open(str(path)) as f:
         nb = nbformat.read(f, nbformat.NO_CONVERT)
 
-    kernel = ExecutePreprocessor(
-        timeout=600, kernel_name="python3", resources={"metadata": {"path": Path(path).parent}}
-    )
+    kernel = ExecutePreprocessor(timeout=600, kernel_name="python3", resources={"metadata": {"path": path.parent}})
+
+    try:
+        result = kernel.preprocess(nb)
+    except Exception as e:
+        raise e
+    finally:
+        if cleanup_checkpoints and checkpoint_path.exists():
+            # clean up if the directory was created by the test
+            shutil.rmtree(checkpoint_path)
 
-    result = kernel.preprocess(nb)
     return result