ansible-test - Avoid using atexit module (#81096)

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.

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)

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."""

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))

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 json
@@ -19,6 +18,7 @@ from .util import (
 )
 
 from .util_common import (
+    ExitHandler,
     named_temporary_file,
 )
 
@@ -184,7 +184,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
 

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)
 

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)

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."""

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)

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