diff --git a/pytest_mh/__init__.py b/pytest_mh/__init__.py index c310d4b..35a1e30 100644 --- a/pytest_mh/__init__.py +++ b/pytest_mh/__init__.py @@ -32,7 +32,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", @@ -66,6 +66,7 @@ "pytest_configure", "Topology", "TopologyController", + "BackupTopologyController", "TopologyDomain", "TopologyMark", ] diff --git a/pytest_mh/_private/topology_controller.py b/pytest_mh/_private/topology_controller.py index 8bf12cc..b21e3a4 100644 --- a/pytest_mh/_private/topology_controller.py +++ b/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 @@ -349,3 +350,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