Merge pull request #8219 from bluetech/setupstate-refactor

runner: refactor SetupState

2 files changed, 115 insertions, 68 deletions

diff --git a/src/_pytest/fixtures.py b/src/_pytest/fixtures.py
index 43a40a864..269369642 100644
--- a/src/_pytest/fixtures.py
+++ b/src/_pytest/fixtures.py
@@ -372,6 +372,7 @@ def _fill_fixtures_impl(function: "Function") -> None:
         fi = fm.getfixtureinfo(function.parent, function.obj, None)
         function._fixtureinfo = fi
         request = function._request = FixtureRequest(function, _ispytest=True)
+        fm.session._setupstate.prepare(function)
         request._fillfixtures()
         # Prune out funcargs for jstests.
         newfuncargs = {}
@@ -543,8 +544,8 @@ class FixtureRequest:
         self._addfinalizer(finalizer, scope=self.scope)
 
     def _addfinalizer(self, finalizer: Callable[[], object], scope) -> None:
-        item = self._getscopeitem(scope)
-        item.addfinalizer(finalizer)
+        node = self._getscopeitem(scope)
+        node.addfinalizer(finalizer)
 
     def applymarker(self, marker: Union[str, MarkDecorator]) -> None:
         """Apply a marker to a single test function invocation.
diff --git a/src/_pytest/runner.py b/src/_pytest/runner.py
index 844e41f80..7bb92cecf 100644
--- a/src/_pytest/runner.py
+++ b/src/_pytest/runner.py
@@ -33,8 +33,10 @@ from _pytest.nodes import Collector
 from _pytest.nodes import Item
 from _pytest.nodes import Node
 from _pytest.outcomes import Exit
+from _pytest.outcomes import OutcomeException
 from _pytest.outcomes import Skipped
 from _pytest.outcomes import TEST_OUTCOME
+from _pytest.store import StoreKey
 
 if TYPE_CHECKING:
     from typing_extensions import Literal
@@ -103,7 +105,7 @@ def pytest_sessionstart(session: "Session") -> None:
 
 
 def pytest_sessionfinish(session: "Session") -> None:
-    session._setupstate.teardown_all()
+    session._setupstate.teardown_exact(None)
 
 
 def pytest_runtest_protocol(item: Item, nextitem: Optional[Item]) -> bool:
@@ -175,7 +177,7 @@ def pytest_runtest_call(item: Item) -> None:
 
 def pytest_runtest_teardown(item: Item, nextitem: Optional[Item]) -> None:
     _update_current_test_var(item, "teardown")
-    item.session._setupstate.teardown_exact(item, nextitem)
+    item.session._setupstate.teardown_exact(nextitem)
     _update_current_test_var(item, None)
 
 
@@ -401,88 +403,132 @@ def pytest_make_collect_report(collector: Collector) -> CollectReport:
 
 
 class SetupState:
-    """Shared state for setting up/tearing down test items or collectors."""
+    """Shared state for setting up/tearing down test items or collectors
+    in a session.
 
-    def __init__(self):
-        self.stack: List[Node] = []
-        self._finalizers: Dict[Node, List[Callable[[], object]]] = {}
+    Suppose we have a collection tree as follows:
 
-    def addfinalizer(self, finalizer: Callable[[], object], colitem) -> None:
-        """Attach a finalizer to the given colitem."""
-        assert colitem and not isinstance(colitem, tuple)
-        assert callable(finalizer)
-        # assert colitem in self.stack  # some unit tests don't setup stack :/
-        self._finalizers.setdefault(colitem, []).append(finalizer)
+    <Session session>
+        <Module mod1>
+            <Function item1>
+        <Module mod2>
+            <Function item2>
 
-    def _pop_and_teardown(self):
-        colitem = self.stack.pop()
-        self._teardown_with_finalization(colitem)
+    The SetupState maintains a stack. The stack starts out empty:
 
-    def _callfinalizers(self, colitem) -> None:
-        finalizers = self._finalizers.pop(colitem, None)
-        exc = None
-        while finalizers:
-            fin = finalizers.pop()
-            try:
-                fin()
-            except TEST_OUTCOME as e:
-                # XXX Only first exception will be seen by user,
-                #     ideally all should be reported.
-                if exc is None:
-                    exc = e
-        if exc:
-            raise exc
+        []
 
-    def _teardown_with_finalization(self, colitem) -> None:
-        self._callfinalizers(colitem)
-        colitem.teardown()
-        for colitem in self._finalizers:
-            assert colitem in self.stack
+    During the setup phase of item1, prepare(item1) is called. What it does
+    is:
 
-    def teardown_all(self) -> None:
-        while self.stack:
-            self._pop_and_teardown()
-        for key in list(self._finalizers):
-            self._teardown_with_finalization(key)
-        assert not self._finalizers
+        push session to stack, run session.setup()
+        push mod1 to stack, run mod1.setup()
+        push item1 to stack, run item1.setup()
 
-    def teardown_exact(self, item, nextitem) -> None:
-        needed_collectors = nextitem and nextitem.listchain() or []
-        self._teardown_towards(needed_collectors)
+    The stack is:
 
-    def _teardown_towards(self, needed_collectors) -> None:
-        exc = None
-        while self.stack:
-            if self.stack == needed_collectors[: len(self.stack)]:
-                break
-            try:
-                self._pop_and_teardown()
-            except TEST_OUTCOME as e:
-                # XXX Only first exception will be seen by user,
-                #     ideally all should be reported.
-                if exc is None:
-                    exc = e
-        if exc:
-            raise exc
+        [session, mod1, item1]
+
+    While the stack is in this shape, it is allowed to add finalizers to
+    each of session, mod1, item1 using addfinalizer().
+
+    During the teardown phase of item1, teardown_exact(item2) is called,
+    where item2 is the next item to item1. What it does is:
+
+        pop item1 from stack, run its teardowns
+        pop mod1 from stack, run its teardowns
+
+    mod1 was popped because it ended its purpose with item1. The stack is:
+
+        [session]
+
+    During the setup phase of item2, prepare(item2) is called. What it does
+    is:
+
+        push mod2 to stack, run mod2.setup()
+        push item2 to stack, run item2.setup()
 
-    def prepare(self, colitem) -> None:
-        """Setup objects along the collector chain to the test-method."""
+    Stack:
 
-        # Check if the last collection node has raised an error.
+        [session, mod2, item2]
+
+    During the teardown phase of item2, teardown_exact(None) is called,
+    because item2 is the last item. What it does is:
+
+        pop item2 from stack, run its teardowns
+        pop mod2 from stack, run its teardowns
+        pop session from stack, run its teardowns
+
+    Stack:
+
+        []
+
+    The end!
+    """
+
+    def __init__(self) -> None:
+        # Maps node -> the node's finalizers.
+        # The stack is in the dict insertion order.
+        self.stack: Dict[Node, List[Callable[[], object]]] = {}
+
+    _prepare_exc_key = StoreKey[Union[OutcomeException, Exception]]()
+
+    def prepare(self, item: Item) -> None:
+        """Setup objects along the collector chain to the item."""
+        # If a collector fails its setup, fail its entire subtree of items.
+        # The setup is not retried for each item - the same exception is used.
         for col in self.stack:
-            if hasattr(col, "_prepare_exc"):
-                exc = col._prepare_exc  # type: ignore[attr-defined]
-                raise exc
+            prepare_exc = col._store.get(self._prepare_exc_key, None)
+            if prepare_exc:
+                raise prepare_exc
 
-        needed_collectors = colitem.listchain()
+        needed_collectors = item.listchain()
         for col in needed_collectors[len(self.stack) :]:
-            self.stack.append(col)
+            assert col not in self.stack
+            self.stack[col] = [col.teardown]
             try:
                 col.setup()
             except TEST_OUTCOME as e:
-                col._prepare_exc = e  # type: ignore[attr-defined]
+                col._store[self._prepare_exc_key] = e
                 raise e
 
+    def addfinalizer(self, finalizer: Callable[[], object], node: Node) -> None:
+        """Attach a finalizer to the given node.
+
+        The node must be currently active in the stack.
+        """
+        assert node and not isinstance(node, tuple)
+        assert callable(finalizer)
+        assert node in self.stack, (node, self.stack)
+        self.stack[node].append(finalizer)
+
+    def teardown_exact(self, nextitem: Optional[Item]) -> None:
+        """Teardown the current stack up until reaching nodes that nextitem
+        also descends from.
+
+        When nextitem is None (meaning we're at the last item), the entire
+        stack is torn down.
+        """
+        needed_collectors = nextitem and nextitem.listchain() or []
+        exc = None
+        while self.stack:
+            if list(self.stack.keys()) == needed_collectors[: len(self.stack)]:
+                break
+            node, finalizers = self.stack.popitem()
+            while finalizers:
+                fin = finalizers.pop()
+                try:
+                    fin()
+                except TEST_OUTCOME as e:
+                    # XXX Only first exception will be seen by user,
+                    #     ideally all should be reported.
+                    if exc is None:
+                        exc = e
+        if exc:
+            raise exc
+        if nextitem is None:
+            assert not self.stack
+
 
 def collect_one_node(collector: Collector) -> CollectReport:
     ihook = collector.ihook