topology-controller: add BackupTopologyController

This controller provides automatic backup and restore of the topology
hosts.

pytest_mh/__init__.py

@@ -24,7 +24,7 @@
 )
 from ._private.plugin import MultihostPlugin, mh_fixture, pytest_addoption, pytest_configure
 from ._private.topology import Topology, TopologyDomain
-from ._private.topology_controller import TopologyController
+from ._private.topology_controller import BackupTopologyController, TopologyController
 
 __all__ = [
     "mh",
@@ -53,6 +53,7 @@
     "pytest_configure",
     "Topology",
     "TopologyController",
+    "BackupTopologyController",
     "TopologyDomain",
     "TopologyMark",
 ]

pytest_mh/_private/topology_controller.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
+from functools import partial, wraps
 from types import SimpleNamespace
 from typing import Any, Callable, Generic
 
 from .artifacts import MultihostArtifactsType, MultihostTopologyControllerArtifacts
 from .logging import MultihostLogger
 from .misc import OperationStatus, invoke_callback
-from .multihost import ConfigType, MultihostDomain, MultihostHost
+from .multihost import ConfigType, MultihostBackupHost, MultihostDomain, MultihostHost
 from .topology import Topology, TopologyDomain
 
 
@@ -343,3 +344,142 @@ def teardown(self, *args, **kwargs) -> None:
         Called after execution of each test.
         """
         pass
+
+
+class BackupTopologyController(TopologyController[ConfigType]):
+    """
+    Implements automatic backup and restore of all topology hosts that inherit
+    from :class:`MultihostBackupHost`.
+
+    The backup of all hosts is taken in :meth:`topology_setup`. It is expected
+    that this method is overridden by the user to setup the topology
+    environment. In such case, it is possible to call
+    ``super().topology_setup(**kwargs)`` at the end of the overridden function
+    or omit this call and store the backup in :attr:`backup_data` manually.
+
+    :meth:`teardown` restores the hosts to the backup taken in
+    :meth:`topology_setup`. This is done after each test, so each test starts
+    with clear topology environment.
+
+    When all tests for this topology are run, :meth:`topology_teardown` is
+    called and the hosts are restored to the original state which backup was
+    taken in :meth:`MultihostBackupHost.pytest_setup` so the environment is
+    fresh for the next topology.
+
+    .. note::
+
+        It is possible to decorate methods, usually the custom implementation of
+        :meth:`topology_setup` with :meth:`restore_vanilla_on_error`. This makes
+        sure that the hosts are reverted to the original state if any of the
+        setup calls fail.
+
+        .. code-block:: python
+
+            @BackupTopologyController.restore_vanilla_on_error
+            def topology_setup(self, *kwargs) -> None:
+                raise Exception("Hosts are automatically restored now.")
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        self.backup_data: dict[MultihostBackupHost, Any | None] = {}
+        """
+        Backup data. Dictionary with host as a key and backup as a value.
+        """
+
+    def restore(self, hosts: dict[MultihostBackupHost, Any | None]) -> None:
+        """
+        Restore given hosts to their given backup.
+
+        :param hosts: Dictionary (host, backup)
+        :type hosts: dict[MultihostBackupHost, Any  |  None]
+        :raises ExceptionGroup: If some hosts fail to restore.
+        """
+        errors = []
+        for host, backup_data in hosts.items():
+            if not isinstance(host, MultihostBackupHost):
+                continue
+
+            try:
+                host.restore(backup_data)
+            except Exception as e:
+                errors.append(e)
+
+        if errors:
+            raise ExceptionGroup("Some hosts failed to restore to original state", errors)
+
+    def restore_vanilla(self) -> None:
+        """
+        Restore to the original host state that is stored in the host object.
+
+        This backup was taken when pytest started and we want to revert to this
+        state when this topology is finished.
+        """
+        restore_data: dict[MultihostBackupHost, Any | None] = {}
+
+        for host in self.hosts:
+            if not isinstance(host, MultihostBackupHost):
+                continue
+
+            restore_data[host] = host.backup_data
+
+        self.restore(restore_data)
+
+    def topology_setup(self, *args, **kwargs) -> None:
+        """
+        Take backup of all topology hosts.
+        """
+        super().topology_setup(**kwargs)
+
+        for host in self.hosts:
+            if not isinstance(host, MultihostBackupHost):
+                continue
+
+            self.backup_data[host] = host.backup()
+
+    def topology_teardown(self, *args, **kwargs) -> None:
+        """
+        Remove all topology backups from the hosts and restore the hosts to the
+        original state before this topology.
+        """
+        try:
+            for host, backup_data in self.backup_data.items():
+                if not isinstance(host, MultihostBackupHost):
+                    continue
+
+                host.remove_backup(backup_data)
+        except Exception:
+            # This is not that important, we can just ignore
+            pass
+
+        self.restore_vanilla()
+
+    def teardown(self, *args, **kwargs) -> None:
+        """
+        Restore the host to the state created by this topology in
+        :meth:`topology_setup` after each test is finished.
+        """
+        self.restore(self.backup_data)
+
+    @staticmethod
+    def restore_vanilla_on_error(method):
+        """
+        Decorator. Restore all hosts to its original state if an exception
+        occurs during method execution.
+
+        :param method: Method to decorate.
+        :type method: Any setup or teardown callback.
+        :return: Decorated method.
+        :rtype: Callback
+        """
+
+        @wraps(method)
+        def wrapper(self: BackupTopologyController, *args, **kwargs):
+            try:
+                return self._invoke_with_args(partial(method, self))
+            except Exception:
+                self.restore_vanilla()
+                raise
+
+        return wrapper