Merge pull request #43 from ma-sadeghi/fix_ipynb_bug_redo

Fix ipynb bug [REDO]

CHANGELOG.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Added
 - PR template
 - Test against 3.10
+- Compatibility with IPython Notebooks
 
 ## [0.5.0] - 2021-11-09
 ### Added

README.md

@@ -61,6 +61,12 @@ Lists the slowest tests based on the information stored in the test durations fi
 ## Interactions with other pytest plugins
 * [`pytest-random-order`](https://github.com/jbasko/pytest-random-order): ⚠️ The **default settings** of that plugin (setting only `--random-order` to activate it) are **incompatible** with `pytest-split`. Test selection in the groups happens after randomization, potentially causing some tests to be selected in several groups and others not at all. Instead, a global random seed needs to be computed before running the tests (for example using `$RANDOM` from the shell) and that single seed then needs to be used for all groups by setting the `--random-order-seed` option.
 
+* [`nbval`](https://github.com/computationalmodelling/nbval): `pytest-split` could, in principle, break up a single IPython Notebook into different test groups. This most likely causes broken up pieces to fail (for the very least, package `import`s are usually done at Cell 1, and so, any broken up piece that doesn't contain Cell 1 will certainly fail). To avoid this, after splitting step is done, test groups are reorganized based on a simple algorithm illustrated in the following cartoon:
+
+![image](https://user-images.githubusercontent.com/14086031/145830494-07afcaf0-5a0f-4817-b9ee-f84a459652a8.png)
+
+where the letters (A to E) refer to individual IPython Notebooks, and the numbers refer to the corresponding cell number.
+
 ## Splitting algorithms
 The plugin supports multiple algorithms to split tests into groups.
 Each algorithm makes different tradeoffs, but generally `least_duration` should give more balanced groups.

src/pytest_split/ipynb_compatibility.py

@@ -0,0 +1,62 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import List
+
+    from pytest_split.algorithms import TestGroup
+
+
+def ensure_ipynb_compatibility(group: "TestGroup", items: list) -> None:
+    """
+    Ensures that group doesn't contain partial IPy notebook cells.
+
+    ``pytest-split`` might, in principle, break up the cells of an
+    IPython notebook into different test groups, in which case the tests
+    most likely fail (for starters, libraries are imported in Cell 0, so
+    all subsequent calls to the imported libraries in the following cells
+    will raise ``NameError``).
+
+    """
+    if not group.selected or not _is_ipy_notebook(group.selected[0].nodeid):
+        return
+
+    item_node_ids = [item.nodeid for item in items]
+
+    # Deal with broken up notebooks at the beginning of the test group
+    first = group.selected[0].nodeid
+    siblings = _find_sibiling_ipynb_cells(first, item_node_ids)
+    if first != siblings[0]:
+        for item in list(group.selected):
+            if item.nodeid in siblings:
+                group.deselected.append(item)
+                group.selected.remove(item)
+
+    if not group.selected or not _is_ipy_notebook(group.selected[-1].nodeid):
+        return
+
+    # Deal with broken up notebooks at the end of the test group
+    last = group.selected[-1].nodeid
+    siblings = _find_sibiling_ipynb_cells(last, item_node_ids)
+    if last != siblings[-1]:
+        for item in list(group.deselected):
+            if item.nodeid in siblings:
+                group.deselected.remove(item)
+                group.selected.append(item)
+
+
+def _find_sibiling_ipynb_cells(
+    ipynb_node_id: str, item_node_ids: "List[str]"
+) -> "List[str]":
+    """
+    Returns all sibiling IPyNb cells given an IPyNb cell nodeid.
+    """
+    fpath = ipynb_node_id.split("::")[0]
+    return [item for item in item_node_ids if fpath in item]
+
+
+def _is_ipy_notebook(node_id: str) -> bool:
+    """
+    Returns True if node_id is an IPython notebook, otherwise False.
+    """
+    fpath = node_id.split("::")[0]
+    return fpath.endswith(".ipynb")

src/pytest_split/plugin.py

@@ -7,6 +7,7 @@
 from _pytest.reports import TestReport
 
 from pytest_split import algorithms
+from pytest_split.ipynb_compatibility import ensure_ipynb_compatibility
 
 if TYPE_CHECKING:
     from typing import Dict, List, Optional, Union
@@ -16,6 +17,7 @@
     from _pytest.config.argparsing import Parser
     from _pytest.main import ExitCode
 
+
 # Ugly hack for freezegun compatibility: https://github.com/spulec/freezegun/issues/286
 STORE_DURATIONS_SETUP_AND_TEARDOWN_THRESHOLD = 60 * 10  # seconds
 
@@ -165,6 +167,8 @@ def pytest_collection_modifyitems(
         groups = algo(splits, items, self.cached_durations)
         group = groups[group_idx - 1]
 
+        ensure_ipynb_compatibility(group, items)
+
         items[:] = group.selected
         config.hook.pytest_deselected(items=group.deselected)
 

tests/test_ipynb.py

@@ -0,0 +1,82 @@
+from collections import namedtuple
+
+import pytest
+
+from pytest_split.algorithms import Algorithms
+from pytest_split.ipynb_compatibility import ensure_ipynb_compatibility
+
+item = namedtuple("item", "nodeid")
+
+
+class TestIPyNb:
+    @pytest.mark.parametrize("algo_name", ["duration_based_chunks"])
+    def test_ensure_ipynb_compatibility(self, algo_name):
+        durations = {
+            "temp/nbs/test_1.ipynb::Cell 0": 1,
+            "temp/nbs/test_1.ipynb::Cell 1": 1,
+            "temp/nbs/test_1.ipynb::Cell 2": 1,
+            "temp/nbs/test_2.ipynb::Cell 0": 3,
+            "temp/nbs/test_2.ipynb::Cell 1": 5,
+            "temp/nbs/test_2.ipynb::Cell 2": 1,
+            "temp/nbs/test_2.ipynb::Cell 3": 4,
+            "temp/nbs/test_3.ipynb::Cell 0": 5,
+            "temp/nbs/test_3.ipynb::Cell 1": 1,
+            "temp/nbs/test_3.ipynb::Cell 2": 1,
+            "temp/nbs/test_3.ipynb::Cell 3": 2,
+            "temp/nbs/test_3.ipynb::Cell 4": 1,
+            "temp/nbs/test_4.ipynb::Cell 0": 1,
+            "temp/nbs/test_4.ipynb::Cell 1": 1,
+            "temp/nbs/test_4.ipynb::Cell 2": 3,
+        }
+        items = [item(x) for x in durations.keys()]
+        algo = Algorithms[algo_name].value
+        groups = algo(splits=3, items=items, durations=durations)
+
+        assert groups[0].selected == [
+            item(nodeid="temp/nbs/test_1.ipynb::Cell 0"),
+            item(nodeid="temp/nbs/test_1.ipynb::Cell 1"),
+            item(nodeid="temp/nbs/test_1.ipynb::Cell 2"),
+            item(nodeid="temp/nbs/test_2.ipynb::Cell 0"),
+            item(nodeid="temp/nbs/test_2.ipynb::Cell 1"),
+        ]
+        assert groups[1].selected == [
+            item(nodeid="temp/nbs/test_2.ipynb::Cell 2"),
+            item(nodeid="temp/nbs/test_2.ipynb::Cell 3"),
+            item(nodeid="temp/nbs/test_3.ipynb::Cell 0"),
+            item(nodeid="temp/nbs/test_3.ipynb::Cell 1"),
+        ]
+        assert groups[2].selected == [
+            item(nodeid="temp/nbs/test_3.ipynb::Cell 2"),
+            item(nodeid="temp/nbs/test_3.ipynb::Cell 3"),
+            item(nodeid="temp/nbs/test_3.ipynb::Cell 4"),
+            item(nodeid="temp/nbs/test_4.ipynb::Cell 0"),
+            item(nodeid="temp/nbs/test_4.ipynb::Cell 1"),
+            item(nodeid="temp/nbs/test_4.ipynb::Cell 2"),
+        ]
+
+        ensure_ipynb_compatibility(groups[0], items)
+        assert groups[0].selected == [
+            item(nodeid="temp/nbs/test_1.ipynb::Cell 0"),
+            item(nodeid="temp/nbs/test_1.ipynb::Cell 1"),
+            item(nodeid="temp/nbs/test_1.ipynb::Cell 2"),
+            item(nodeid="temp/nbs/test_2.ipynb::Cell 0"),
+            item(nodeid="temp/nbs/test_2.ipynb::Cell 1"),
+            item(nodeid="temp/nbs/test_2.ipynb::Cell 2"),
+            item(nodeid="temp/nbs/test_2.ipynb::Cell 3"),
+        ]
+
+        ensure_ipynb_compatibility(groups[1], items)
+        assert groups[1].selected == [
+            item(nodeid="temp/nbs/test_3.ipynb::Cell 0"),
+            item(nodeid="temp/nbs/test_3.ipynb::Cell 1"),
+            item(nodeid="temp/nbs/test_3.ipynb::Cell 2"),
+            item(nodeid="temp/nbs/test_3.ipynb::Cell 3"),
+            item(nodeid="temp/nbs/test_3.ipynb::Cell 4"),
+        ]
+
+        ensure_ipynb_compatibility(groups[2], items)
+        assert groups[2].selected == [
+            item(nodeid="temp/nbs/test_4.ipynb::Cell 0"),
+            item(nodeid="temp/nbs/test_4.ipynb::Cell 1"),
+            item(nodeid="temp/nbs/test_4.ipynb::Cell 2"),
+        ]