diff --git a/changelogs/fragments/ansible-test-atexit.yml b/changelogs/fragments/ansible-test-atexit.yml new file mode 100644 index 00000000000..e531da47503 --- /dev/null +++ b/changelogs/fragments/ansible-test-atexit.yml @@ -0,0 +1,2 @@ +minor_changes: + - ansible-test - Use a context manager to perform cleanup at exit instead of using the built-in ``atexit`` module. diff --git a/test/lib/ansible_test/_internal/__init__.py b/test/lib/ansible_test/_internal/__init__.py index ee24a852c2e..35584746b57 100644 --- a/test/lib/ansible_test/_internal/__init__.py +++ b/test/lib/ansible_test/_internal/__init__.py @@ -43,6 +43,7 @@ from .data import ( from .util_common import ( CommonConfig, + ExitHandler, ) from .cli import ( @@ -59,6 +60,12 @@ from .config import ( def main(cli_args: t.Optional[list[str]] = None) -> None: + """Wrapper around the main program function to invoke cleanup functions at exit.""" + with ExitHandler.context(): + main_internal(cli_args) + + +def main_internal(cli_args: t.Optional[list[str]] = None) -> None: """Main program function.""" try: os.chdir(data_context().content.root) diff --git a/test/lib/ansible_test/_internal/commands/integration/cloud/__init__.py b/test/lib/ansible_test/_internal/commands/integration/cloud/__init__.py index eac9265a4f8..cf920bc9e1e 100644 --- a/test/lib/ansible_test/_internal/commands/integration/cloud/__init__.py +++ b/test/lib/ansible_test/_internal/commands/integration/cloud/__init__.py @@ -2,7 +2,6 @@ from __future__ import annotations import abc -import atexit import datetime import os import re @@ -28,6 +27,7 @@ from ....util import ( ) from ....util_common import ( + ExitHandler, ResultType, write_json_test_results, ) @@ -306,7 +306,7 @@ class CloudProvider(CloudBase): self.resource_prefix = self.ci_provider.generate_resource_prefix() self.resource_prefix = re.sub(r'[^a-zA-Z0-9]+', '-', self.resource_prefix)[:63].lower().rstrip('-') - atexit.register(self.cleanup) + ExitHandler.register(self.cleanup) def cleanup(self) -> None: """Clean up the cloud resource and any temporary configuration files after tests complete.""" diff --git a/test/lib/ansible_test/_internal/commands/sanity/validate_modules.py b/test/lib/ansible_test/_internal/commands/sanity/validate_modules.py index c394a97a258..d2b1c546d58 100644 --- a/test/lib/ansible_test/_internal/commands/sanity/validate_modules.py +++ b/test/lib/ansible_test/_internal/commands/sanity/validate_modules.py @@ -1,7 +1,6 @@ """Sanity test using validate-modules.""" from __future__ import annotations -import atexit import collections import contextlib import json @@ -38,6 +37,7 @@ from ...util import ( ) from ...util_common import ( + ExitHandler, process_scoped_temporary_directory, run_command, ResultType, @@ -242,7 +242,7 @@ class ValidateModulesTest(SanitySingleVersion): files = payload_config.files files.append((path, os.path.relpath(path, data_context().content.root))) - atexit.register(cleanup) + ExitHandler.register(cleanup) data_context().register_payload_callback(git_callback) make_dirs(os.path.dirname(path)) diff --git a/test/lib/ansible_test/_internal/containers.py b/test/lib/ansible_test/_internal/containers.py index bfc36434dd5..869f1fba75b 100644 --- a/test/lib/ansible_test/_internal/containers.py +++ b/test/lib/ansible_test/_internal/containers.py @@ -1,7 +1,6 @@ """High level functions for working with containers.""" from __future__ import annotations -import atexit import collections.abc as c import contextlib import enum @@ -20,6 +19,7 @@ from .util import ( ) from .util_common import ( + ExitHandler, named_temporary_file, ) @@ -225,7 +225,7 @@ def run_support_container( raise Exception(f'Container already defined: {name}') if not support_containers: - atexit.register(cleanup_containers, args) + ExitHandler.register(cleanup_containers, args) support_containers[name] = descriptor diff --git a/test/lib/ansible_test/_internal/coverage_util.py b/test/lib/ansible_test/_internal/coverage_util.py index 0af1cac4e97..ae64024939d 100644 --- a/test/lib/ansible_test/_internal/coverage_util.py +++ b/test/lib/ansible_test/_internal/coverage_util.py @@ -1,7 +1,6 @@ """Utility code for facilitating collection of code coverage when running tests.""" from __future__ import annotations -import atexit import dataclasses import os import sqlite3 @@ -34,6 +33,7 @@ from .data import ( ) from .util_common import ( + ExitHandler, intercept_python, ResultType, ) @@ -223,7 +223,7 @@ def get_coverage_config(args: TestConfig) -> str: temp_dir = '/tmp/coverage-temp-dir' else: temp_dir = tempfile.mkdtemp() - atexit.register(lambda: remove_tree(temp_dir)) + ExitHandler.register(lambda: remove_tree(temp_dir)) path = os.path.join(temp_dir, COVERAGE_CONFIG_NAME) diff --git a/test/lib/ansible_test/_internal/payload.py b/test/lib/ansible_test/_internal/payload.py index 10dde7b8b15..ab9739b4835 100644 --- a/test/lib/ansible_test/_internal/payload.py +++ b/test/lib/ansible_test/_internal/payload.py @@ -1,7 +1,6 @@ """Payload management for sending Ansible files and test content to other systems (VMs, containers).""" from __future__ import annotations -import atexit import os import stat import tarfile @@ -32,6 +31,7 @@ from .data import ( from .util_common import ( CommonConfig, + ExitHandler, ) # improve performance by disabling uid/gid lookups @@ -192,7 +192,7 @@ def create_temporary_bin_files(args: CommonConfig) -> tuple[tuple[str, str], ... temp_path = '/tmp/ansible-tmp-bin' else: temp_path = tempfile.mkdtemp(prefix='ansible', suffix='bin') - atexit.register(remove_tree, temp_path) + ExitHandler.register(remove_tree, temp_path) for name, dest in ANSIBLE_BIN_SYMLINK_MAP.items(): path = os.path.join(temp_path, name) diff --git a/test/lib/ansible_test/_internal/provisioning.py b/test/lib/ansible_test/_internal/provisioning.py index e7f0fd317e8..4710757bd0a 100644 --- a/test/lib/ansible_test/_internal/provisioning.py +++ b/test/lib/ansible_test/_internal/provisioning.py @@ -1,7 +1,6 @@ """Provision hosts for running tests.""" from __future__ import annotations -import atexit import collections.abc as c import dataclasses import functools @@ -27,6 +26,10 @@ from .util import ( type_guard, ) +from .util_common import ( + ExitHandler, +) + from .thread import ( WrappedThread, ) @@ -124,7 +127,7 @@ def prepare_profiles( raise PrimeContainers() - atexit.register(functools.partial(cleanup_profiles, host_state)) + ExitHandler.register(functools.partial(cleanup_profiles, host_state)) def provision(profile: HostProfile) -> None: """Provision the given profile.""" diff --git a/test/lib/ansible_test/_internal/pypi_proxy.py b/test/lib/ansible_test/_internal/pypi_proxy.py index 97663eadd1b..5380dd9b6a4 100644 --- a/test/lib/ansible_test/_internal/pypi_proxy.py +++ b/test/lib/ansible_test/_internal/pypi_proxy.py @@ -1,7 +1,6 @@ """PyPI proxy management.""" from __future__ import annotations -import atexit import os import urllib.parse @@ -23,6 +22,7 @@ from .util import ( ) from .util_common import ( + ExitHandler, process_scoped_temporary_file, ) @@ -128,7 +128,7 @@ def configure_target_pypi_proxy(args: EnvironmentConfig, profile: HostProfile, p run_playbook(args, inventory_path, 'pypi_proxy_prepare.yml', capture=True, variables=dict( pypi_endpoint=pypi_endpoint, pypi_hostname=pypi_hostname, force=force)) - atexit.register(cleanup_pypi_proxy) + ExitHandler.register(cleanup_pypi_proxy) def configure_pypi_proxy_pip(args: EnvironmentConfig, profile: HostProfile, pypi_endpoint: str, pypi_hostname: str) -> None: @@ -153,7 +153,7 @@ trusted-host = {1} if not args.explain: write_text_file(pip_conf_path, pip_conf, True) - atexit.register(pip_conf_cleanup) + ExitHandler.register(pip_conf_cleanup) def configure_pypi_proxy_easy_install(args: EnvironmentConfig, profile: HostProfile, pypi_endpoint: str) -> None: @@ -177,4 +177,4 @@ index_url = {0} if not args.explain: write_text_file(pydistutils_cfg_path, pydistutils_cfg, True) - atexit.register(pydistutils_cfg_cleanup) + ExitHandler.register(pydistutils_cfg_cleanup) diff --git a/test/lib/ansible_test/_internal/util_common.py b/test/lib/ansible_test/_internal/util_common.py index 79ff6c03ae3..bb76b949974 100644 --- a/test/lib/ansible_test/_internal/util_common.py +++ b/test/lib/ansible_test/_internal/util_common.py @@ -1,7 +1,6 @@ """Common utility code that depends on CommonConfig.""" from __future__ import annotations -import atexit import collections.abc as c import contextlib import json @@ -64,6 +63,39 @@ from .host_configs import ( CHECK_YAML_VERSIONS: dict[str, t.Any] = {} +class ExitHandler: + """Simple exit handler implementation.""" + _callbacks: list[tuple[t.Callable, tuple[t.Any, ...], dict[str, t.Any]]] = [] + + @staticmethod + def register(func: t.Callable, *args, **kwargs) -> None: + """Register the given function and args as a callback to execute during program termination.""" + ExitHandler._callbacks.append((func, args, kwargs)) + + @staticmethod + @contextlib.contextmanager + def context() -> t.Generator[None, None, None]: + """Run all registered handlers when the context is exited.""" + last_exception: BaseException | None = None + + try: + yield + finally: + queue = list(ExitHandler._callbacks) + + while queue: + func, args, kwargs = queue.pop() + + try: + func(*args, **kwargs) + except BaseException as ex: # pylint: disable=broad-exception-caught + last_exception = ex + display.fatal(f'Exit handler failed: {ex}') + + if last_exception: + raise last_exception + + class ShellScriptTemplate: """A simple substitution template for shell scripts.""" @@ -211,7 +243,7 @@ def process_scoped_temporary_file(args: CommonConfig, prefix: t.Optional[str] = else: temp_fd, path = tempfile.mkstemp(prefix=prefix, suffix=suffix) os.close(temp_fd) - atexit.register(lambda: os.remove(path)) + ExitHandler.register(lambda: os.remove(path)) return path @@ -222,7 +254,7 @@ def process_scoped_temporary_directory(args: CommonConfig, prefix: t.Optional[st path = os.path.join(tempfile.gettempdir(), f'{prefix or tempfile.gettempprefix()}{generate_name()}{suffix or ""}') else: path = tempfile.mkdtemp(prefix=prefix, suffix=suffix) - atexit.register(lambda: remove_tree(path)) + ExitHandler.register(lambda: remove_tree(path)) return path @@ -296,7 +328,7 @@ def get_injector_path() -> str: """Remove the temporary injector directory.""" remove_tree(injector_path) - atexit.register(cleanup_injector) + ExitHandler.register(cleanup_injector) return injector_path @@ -354,7 +386,7 @@ def get_python_path(interpreter: str) -> str: verified_chmod(python_path, MODE_DIRECTORY) if not PYTHON_PATHS: - atexit.register(cleanup_python_paths) + ExitHandler.register(cleanup_python_paths) PYTHON_PATHS[interpreter] = python_path @@ -364,7 +396,7 @@ def get_python_path(interpreter: str) -> str: def create_temp_dir(prefix: t.Optional[str] = None, suffix: t.Optional[str] = None, base_dir: t.Optional[str] = None) -> str: """Create a temporary directory that persists until the current process exits.""" temp_path = tempfile.mkdtemp(prefix=prefix or 'tmp', suffix=suffix or '', dir=base_dir) - atexit.register(remove_tree, temp_path) + ExitHandler.register(remove_tree, temp_path) return temp_path