multihost: add MultihostBackupHost and BackupTopologyController

pytest_mh/__init__.py

@@ -9,6 +9,7 @@
 from ._private.fixtures import MultihostFixture, mh
 from ._private.marks import KnownTopologyBase, KnownTopologyGroupBase, TopologyMark
 from ._private.multihost import (
+    MultihostBackupHost,
     MultihostConfig,
     MultihostDomain,
     MultihostHost,
@@ -23,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",
@@ -34,6 +35,7 @@
     "MultihostDomain",
     "MultihostFixture",
     "MultihostHost",
+    "MultihostBackupHost",
     "MultihostHostArtifacts",
     "MultihostItemData",
     "MultihostOSFamily",
@@ -51,6 +53,7 @@
     "pytest_configure",
     "Topology",
     "TopologyController",
+    "BackupTopologyController",
     "TopologyDomain",
     "TopologyMark",
 ]

pytest_mh/_private/multihost.py

@@ -4,8 +4,8 @@
 from collections import deque
 from contextlib import contextmanager
 from functools import wraps
-from pathlib import Path
-from typing import TYPE_CHECKING, Any, Generator, Generic, Self, Type, TypeVar
+from pathlib import Path, PurePath
+from typing import TYPE_CHECKING, Any, Generator, Generic, Self, Sequence, Type, TypeVar
 
 import pytest
 
@@ -632,6 +632,164 @@ def get_connection(self, shell: Shell) -> Connection:
                 raise ValueError(f"Unknown connection type: {conn_type}!")
 
 
+class MultihostBackupHost(MultihostHost[DomainType], ABC):
+    """
+    Abstract class implementing automatic backup and restore for a host.
+
+    A backup of the host is created once when pytest starts and the host is
+    restored automatically (unless disabled) when a test run is finished.
+
+    If the backup data is stored as :class:`~pathlib.PurePath` or a sequence of
+    :class:`~pathlib.PurePath`, the file is automatically removed from the host
+    when all tests are finished. Otherwise no action is done -- it is possible
+    to overwrite :meth:`remove_backup` to clean up your data if needed.
+
+    It is required to implement :meth:`start`, :meth:`stop`, :meth:`backup` and
+    :meth:`restore`. The :meth:`start` method is called in :meth:`pytest_setup`
+    unless ``auto_start`` is set to False and the implementation of this method
+    may raise ``NotImplementedError`` which will be ignored.
+
+    By default, the host is reverted when each test run is finished. This may
+    not always be desirable and can be disabled via ``auto_restore`` parameter
+    of the constructor.
+    """
+
+    def __init__(self, *args, auto_start: bool = True, auto_restore: bool = True, **kwargs) -> None:
+        """
+        :param auto_start: Automatically start service before taking the first
+            backup.
+        :type auto_restore: bool, optional
+        :param auto_restore: If True, the host is automatically restored to the
+            backup state when a test is finished in :meth:`teardown`, defaults
+            to True
+        :type auto_restore: bool, optional
+        """
+        super().__init__(*args, **kwargs)
+
+        self.backup_data: PurePath | Sequence[PurePath] | Any | None = None
+        """Backup data of vanilla state of this host."""
+
+        self._backup_auto_start: bool = auto_start
+        """
+        If True, the host is automatically started prior taking the first
+        backup.
+        """
+
+        self._backup_auto_restore: bool = auto_restore
+        """
+        If True, the host is automatically restored to the backup state when a
+        test is finished in :meth:`teardown`.
+        """
+
+    def pytest_setup(self) -> None:
+        """
+        Start the services via :meth:`start` and take a backup by calling
+        :meth:`backup`.
+        """
+        # Make sure required services are running
+        if self._backup_auto_start:
+            try:
+                self.start()
+            except NotImplementedError:
+                pass
+
+        # Create backup of initial state
+        self.backup_data = self.backup()
+
+    def pytest_teardown(self) -> None:
+        """
+        Remove backup files from the host (calls :meth:`remove_backup`).
+        """
+        self.remove_backup(self.backup_data)
+
+    def teardown(self) -> None:
+        """
+        Restore the host from the backup by calling :meth:`restore`.
+        """
+        if self._backup_auto_restore:
+            self.restore(self.backup_data)
+
+        super().teardown()
+
+    def remove_backup(self, backup_data: PurePath | Sequence[PurePath] | Any | None) -> None:
+        """
+        Remove backup data from the host.
+
+        If backup_data is not :class:`~pathlib.PurePath` or a sequence of
+        :class:`~pathlib.PurePath`, this will not have any effect. Otherwise,
+        the paths are removed from the host.
+
+        :param backup_data: Backup data.
+        :type backup_data: PurePath | Sequence[PurePath] | Any | None
+        """
+        if backup_data is None:
+            return
+
+        if isinstance(backup_data, PurePath):
+            backup_data = [backup_data]
+
+        if isinstance(backup_data, Sequence):
+            only_paths = True
+            for item in backup_data:
+                if not isinstance(item, PurePath):
+                    only_paths = False
+                    break
+
+            if only_paths:
+                if isinstance(self.conn.shell, Powershell):
+                    for item in backup_data:
+                        path = str(item)
+                        self.conn.exec(["Remove-Item", "-Force", "-Recurse", path])
+                else:
+                    for item in backup_data:
+                        path = str(item)
+                        self.conn.exec(["rm", "-fr", path])
+
+    @abstractmethod
+    def start(self) -> None:
+        """
+        Start required services.
+
+        :raises NotImplementedError: If start operation is not supported.
+        """
+        pass
+
+    @abstractmethod
+    def stop(self) -> None:
+        """
+        Stop required services.
+
+        :raises NotImplementedError: If stop operation is not supported.
+        """
+        pass
+
+    @abstractmethod
+    def backup(self) -> PurePath | Sequence[PurePath] | Any | None:
+        """
+        Backup backend data.
+
+        Returns directory or file path where the backup is stored (as
+        :class:`~pathlib.PurePath` or sequence of :class:`~pathlib.PurePath`) or
+        any Python data relevant for the backup. This data is passed to
+        :meth:`restore` which will use this information to restore the host to
+        its original state.
+
+        :return: Backup data.
+        :rtype: PurePath | Sequence[PurePath] | Any | None
+        """
+        pass
+
+    @abstractmethod
+    def restore(self, backup_data: Any | None) -> None:
+        """
+        Restore data from the backup.
+
+        :param backup_data: Backup data.
+        :type backup_data: PurePath | Sequence[PurePath] | Any | None
+        """
+        pass
+
+
 HostType = TypeVar("HostType", bound=MultihostHost)
 
 

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