Prepare testing utils for testing specific notebooks

Author: classiqdor

tests/utils_for_testbook.py

@@ -0,0 +1,75 @@
+import json
+import os
+from typing import Any, Callable
+import pytest
+
+from testbook import testbook
+from tests.utils_for_tests import resolve_notebook_path, should_skip_notebook
+
+from classiq.interface.generator.quantum_program import QuantumProgram
+
+
+def wrap_testbook(notebook_name: str, timeout_seconds: float = 10) -> Callable:
+    def inner_decorator(func: Callable) -> Any:
+        notebook_path = resolve_notebook_path(notebook_name)
+
+        for decorator in [
+            testbook(notebook_path, execute=True, timeout=timeout_seconds),
+            _build_cd_decorator(notebook_path),
+            _build_skip_decorator(notebook_path),
+        ]:
+            func = decorator(func)
+        return func
+
+    return inner_decorator
+
+
+# The purpose of the `cd_decorator` is to execute the test in the same folder as the `ipynb` file
+#   so that relative files (images, csv, etc.) will be available
+def _build_cd_decorator(file_path: str) -> Callable:
+    def cd_decorator(func: Callable) -> Any:
+        def inner(*args: Any, **kwargs: Any) -> Any:
+            previous_dir = os.getcwd()
+            os.chdir(os.path.dirname(file_path))
+
+            func(*args, **kwargs)
+
+            os.chdir(previous_dir)
+
+        return inner
+
+    return cd_decorator
+
+
+def _build_skip_decorator(notebook_path: str) -> Callable:
+    notebook_name = os.path.basename(notebook_path)
+    return pytest.mark.skipif(
+        should_skip_notebook(notebook_name),
+        reason="Didn't change",
+    )
+
+
+def validate_quantum_model(model: str) -> None:
+    # currently that's some dummy test - only checking that it's a valid dict
+    assert isinstance(json.loads(model), dict)
+
+
+def validate_quantum_program_size(
+    quantum_program: str,
+    expected_width: int | None = None,
+    expected_depth: int | None = None,
+) -> None:
+    qp = QuantumProgram.model_validate_json(quantum_program)
+
+    actual_width = qp.data.width
+    if expected_width is not None:
+        assert (
+            actual_width <= expected_width
+        ), f"The width of the circuit changed! (for the worse!). From {expected_width} to {actual_width}"
+
+    assert qp.transpiled_circuit is not None  # for mypy
+    actual_depth = qp.transpiled_circuit.depth
+    if expected_depth is not None:
+        assert (
+            actual_depth <= expected_depth
+        ), f"The depth of the circuit changed! (for the worse!). From {expected_depth} to {actual_depth}"

tests/utils_for_tests.py

@@ -1,11 +1,11 @@
 import os
-from collections.abc import Iterable
+from functools import lru_cache
 from pathlib import Path
 
 ROOT_DIRECTORY = Path(__file__).parents[1]
 
 
-def iterate_notebooks() -> Iterable[str]:
+def iterate_notebooks() -> list[str]:
     if os.environ.get("SHOULD_TEST_ALL_FILES", "") == "true":
         notebooks_to_test = _get_all_notebooks()
     else:
@@ -14,8 +14,39 @@ def iterate_notebooks() -> Iterable[str]:
     return notebooks_to_test
 
 
-def _get_all_notebooks(directory: Path = ROOT_DIRECTORY) -> Iterable[str]:
-    for root, _, files in os.walk(directory):
+@lru_cache
+def iterate_notebook_names() -> list[str]:
+    return list(map(os.path.basename, iterate_notebooks()))
+
+
+@lru_cache
+def _get_all_notebooks(
+    directory: Path = ROOT_DIRECTORY, suffix: str = ".ipynb"
+) -> list[str]:
+    return [
+        file
+        for root, _, files in os.walk(directory)
+        for file in files
+        if file.endswith(suffix)
+    ]
+
+
+def should_run_notebook(notebook_name: str) -> bool:
+    return notebook_name in iterate_notebook_names()
+
+
+def should_skip_notebook(notebook_name: str) -> bool:
+    return not should_run_notebook(notebook_name)
+
+
+@lru_cache
+def resolve_notebook_path(notebook_name: str) -> str:
+    notebook_name_lower = notebook_name.lower()
+    if not notebook_name_lower.endswith(".ipynb"):
+        notebook_name_lower += ".ipynb"
+
+    for root, _, files in os.walk(ROOT_DIRECTORY):
         for file in files:
-            if file.endswith(".ipynb"):
-                yield os.path.join(root, file)
+            if file.lower() == notebook_name_lower:
+                return os.path.join(root, file)
+    raise LookupError(f"Notebook not found: {notebook_name}")