| ############################################################################ |
| # Copyright (C) SchedMD LLC. |
| ############################################################################ |
| import _pytest |
| |
| # import inspect |
| import logging |
| import os |
| |
| # import pwd |
| import pytest |
| import re |
| |
| # import shutil |
| import sys |
| |
| # from pathlib import Path |
| |
| sys.path.append(sys.path[0] + "/lib") |
| import atf |
| |
| |
| # Add test description (docstring) as a junit property |
| def pytest_itemcollected(item): |
| node = item.obj |
| desc = node.__doc__.strip() if node.__doc__ else node.__name__ |
| if desc: |
| item.user_properties.append(("description", desc)) |
| |
| |
| def pytest_addoption(parser): |
| config_group = parser.getgroup("config mode") |
| config_group.addoption( |
| "--auto-config", |
| action="store_true", |
| help="the slurm configuration will be altered as needed by the test", |
| ) |
| config_group.addoption( |
| "--local-config", |
| action="store_false", |
| dest="auto_config", |
| help="the slurm configuration will not be altered", |
| ) |
| parser.addoption( |
| "--no-color", |
| action="store_true", |
| dest="no_color", |
| help="the pytest logs won't include colors", |
| ) |
| parser.addoption( |
| "--allow-slurmdbd-modify", |
| action="store_true", |
| dest="allow_slurmdbd_modify", |
| help="allow running in local-config even if require_accounting(modify=True)", |
| ) |
| |
| |
| def color_log_level(level: int, **color_kwargs): |
| # Adapted from deprecated py.io TerminalWriter source |
| # https://py.readthedocs.io/en/latest/_modules/py/_io/terminalwriter.html |
| _esctable = dict( |
| black=30, |
| red=31, |
| green=32, |
| yellow=33, |
| blue=34, |
| purple=35, |
| cyan=36, |
| white=37, |
| Black=40, |
| Red=41, |
| Green=42, |
| Yellow=43, |
| Blue=44, |
| Purple=45, |
| Cyan=46, |
| White=47, |
| bold=1, |
| light=2, |
| blink=5, |
| invert=7, |
| ) |
| |
| for handler in logging.getLogger().handlers: |
| if isinstance(handler, _pytest.logging.LogCaptureHandler): |
| formatter = handler.formatter |
| if match := formatter.LEVELNAME_FMT_REGEX.search(formatter._fmt): |
| levelname_fmt = match.group() |
| formatted_levelname = levelname_fmt % { |
| "levelname": logging.getLevelName(level) |
| } |
| |
| esc = [] |
| for option in color_kwargs: |
| esc.append(_esctable[option]) |
| |
| colorized_formatted_levelname = ( |
| "".join(["\x1b[%sm" % cod for cod in esc]) |
| + formatted_levelname |
| + "\x1b[0m" |
| ) |
| |
| formatter._level_to_fmt_mapping[level] = ( |
| formatter.LEVELNAME_FMT_REGEX.sub( |
| colorized_formatted_levelname, formatter._fmt |
| ) |
| ) |
| |
| |
| @pytest.fixture(scope="session", autouse=True) |
| def session_setup(request): |
| # Set the auto-config and other properties from the options |
| atf.properties["auto-config"] = request.config.getoption("--auto-config") |
| atf.properties["allow-slurmdbd-modify"] = request.config.getoption( |
| "--allow-slurmdbd-modify" |
| ) |
| if not request.config.getoption("--no-color"): |
| # Customize logging level colors |
| color_log_level(logging.CRITICAL, red=True, bold=True) |
| color_log_level(logging.ERROR, red=True) |
| color_log_level(logging.WARNING, yellow=True) |
| color_log_level(logging.INFO, green=True) |
| color_log_level(logging.NOTE, cyan=True) |
| color_log_level(logging.DEBUG, blue=True, bold=True) |
| color_log_level(logging.TRACE, purple=True, bold=True) |
| |
| |
| def update_tmp_path_exec_permissions(): |
| """ |
| For pytest versions 6+ the tmp path it uses no longer has |
| public exec permissions for dynamically created directories by default. |
| |
| This causes problems when trying to read temp files during tests as |
| users other than atf (ie slurm). The tests will fail with permission denied. |
| |
| To fix this we check and add the x bit to the public group on tmp |
| directories so the files inside can be read. Adding just 'read' is |
| not enough |
| |
| Bug 16568 |
| """ |
| |
| user_name = atf.get_user_name() |
| path = f"/tmp/pytest-of-{user_name}" |
| |
| if os.path.isdir(path): |
| os.chmod(path, 0o777) |
| for root, dirs, files in os.walk(path): |
| for d in dirs: |
| os.chmod(os.path.join(root, d), 0o777) |
| |
| |
| @pytest.fixture(scope="module", autouse=True) |
| def module_setup(request, tmp_path_factory): |
| atf.properties["slurm-started"] = False |
| atf.properties["slurmrestd-started"] = False |
| atf.properties["influxdb-started"] = False |
| atf.properties["configurations-modified"] = set() |
| atf.properties["orig-environment"] = dict(os.environ) |
| atf.properties["orig-pypath"] = list(sys.path) |
| atf.properties["forced_upgrade_setup"] = False |
| if "old-slurm-prefix" in atf.properties.keys(): |
| del atf.properties["old-slurm-prefix"] |
| if "new-slurm-prefix" in atf.properties.keys(): |
| del atf.properties["new-slurm-prefix"] |
| |
| # Ensure that slurm-spool-dir, slurm-tmpfs and nodes are set. |
| atf.properties["slurm-spool-dir"] = atf.get_config_parameter( |
| "SlurmdSpoolDir", live=False, quiet=True |
| ) |
| atf.properties["slurm-tmpfs"] = atf.get_config_parameter( |
| "TmpFS", live=False, quiet=True |
| ) |
| atf.properties["nodes"] = [] |
| |
| # Creating a module level tmp_path mimicking what tmp_path does |
| name = request.node.name |
| name = re.sub(r"[\W]", "_", name) |
| name = name[:30] |
| atf.properties["test_name"] = name |
| atf.module_tmp_path = tmp_path_factory.mktemp(name, numbered=True) |
| update_tmp_path_exec_permissions() |
| |
| # Module-level fixtures should run from within the module_tmp_path |
| os.chdir(atf.module_tmp_path) |
| |
| # Stop Slurm if using auto-config and Slurm is already running |
| if atf.properties["auto-config"] and atf.is_slurmctld_running(quiet=True): |
| logging.warning( |
| "Auto-config requires Slurm to be initially stopped but Slurm was found running. Stopping Slurm" |
| ) |
| atf.stop_slurm(quiet=True) |
| |
| if atf.properties["auto-config"]: |
| # Cleanup StateSaveLocation for auto-config |
| atf.properties["statesaveloc"] = atf.get_config_parameter( |
| "StateSaveLocation", live=False, quiet=True |
| ) |
| if os.path.exists(atf.properties["statesaveloc"]): |
| if os.path.exists(atf.properties["statesaveloc"] + name): |
| logging.warning( |
| f"Backup for StateSaveLocation already exists ({atf.properties['statesaveloc']+name}). Removing it." |
| ) |
| atf.run_command( |
| f"rm -rf {atf.properties['statesaveloc']+name}", |
| user="root", |
| quiet=True, |
| ) |
| atf.run_command( |
| f"mv {atf.properties['statesaveloc']} {atf.properties['statesaveloc']+name}", |
| user="root", |
| quiet=True, |
| ) |
| |
| # Create the required node directories for node0 |
| node_name = "node0" |
| spool_dir = atf.properties["slurm-spool-dir"].replace("%n", node_name) |
| tmpfs_dir = atf.properties["slurm-tmpfs"].replace("%n", node_name) |
| atf.properties["nodes"].append(node_name) |
| atf.run_command(f"sudo mkdir -p {spool_dir}", fatal=True, quiet=True) |
| atf.run_command(f"sudo mkdir -p {tmpfs_dir}", fatal=True, quiet=True) |
| |
| yield |
| |
| # Return to the folder from which pytest was executed |
| os.chdir(request.config.invocation_dir) |
| |
| # Teardown |
| module_teardown() |
| |
| |
| def module_teardown(): |
| failures = [] |
| |
| if atf.properties["auto-config"]: |
| if atf.properties["slurm-started"] is True: |
| # Cancel all jobs |
| if not atf.cancel_all_jobs(quiet=True): |
| failures.append("Not all jobs were successfully cancelled") |
| |
| # Stop Slurm if we started it |
| if not atf.stop_slurm(fatal=False, quiet=True): |
| failures.append("Not all Slurm daemons were successfully stopped") |
| |
| # Restore the Slurm database |
| atf.restore_accounting_database() |
| |
| # Restore StateSaveLocation for auto-config |
| atf.run_command( |
| f"rm -rf {atf.properties['statesaveloc']}", user="root", quiet=True |
| ) |
| if os.path.exists(atf.properties["statesaveloc"] + atf.properties["test_name"]): |
| atf.run_command( |
| f"mv {atf.properties['statesaveloc']+atf.properties['test_name']} {atf.properties['statesaveloc']}", |
| user="root", |
| quiet=True, |
| ) |
| |
| # Remove Nodes directories: |
| if "nodes" not in atf.properties: |
| atf.properties["nodes"] = ["node0"] |
| for node_name in atf.properties["nodes"]: |
| spool_dir = atf.properties["slurm-spool-dir"].replace("%n", node_name) |
| tmpfs_dir = atf.properties["slurm-tmpfs"].replace("%n", node_name) |
| atf.run_command(f"sudo rm -rf {spool_dir}", quiet=True) |
| atf.run_command(f"sudo rm -rf {tmpfs_dir}", quiet=True) |
| |
| # Restore upgrade setup |
| if atf.properties.get("forced_upgrade_setup"): |
| logging.debug("Restoring upgrade setup...") |
| if not os.path.exists(f"{atf.module_tmp_path}/upgrade-sbin"): |
| pytest.fail( |
| f"Can't restore upgrade setup, {atf.module_tmp_path}/upgrade-sbin doesn't exists." |
| ) |
| if not os.path.exists(f"{atf.module_tmp_path}/upgrade-bin"): |
| pytest.fail( |
| f"Can't restore upgrade setup, {atf.module_tmp_path}/upgrade-bin doesn't exists." |
| ) |
| atf.run_command( |
| f"sudo rm -rf {atf.properties['slurm-sbin-dir']} {atf.properties['slurm-bin-dir']}", |
| quiet=True, |
| fatal=True, |
| ) |
| atf.run_command( |
| f"sudo mv {atf.module_tmp_path}/upgrade-sbin {atf.properties['slurm-sbin-dir']}", |
| quiet=True, |
| fatal=True, |
| ) |
| atf.run_command( |
| f"sudo mv {atf.module_tmp_path}/upgrade-bin {atf.properties['slurm-bin-dir']}", |
| quiet=True, |
| fatal=True, |
| ) |
| |
| # Restore any backed up configuration files |
| for config in set(atf.properties["configurations-modified"]): |
| atf.restore_config_file(config) |
| |
| # Clean influxdb |
| if atf.properties["influxdb-started"]: |
| atf.request_influxdb(f"DROP DATABASE {atf.properties['influxdb_db']}") |
| else: |
| atf.cancel_jobs(atf.properties["submitted-jobs"]) |
| |
| # Restore the prior environment |
| os.environ.clear() |
| os.environ.update(atf.properties["orig-environment"]) |
| sys.path = atf.properties["orig-pypath"] |
| |
| if failures: |
| pytest.fail(failures[0]) |
| |
| |
| @pytest.fixture(scope="function", autouse=True) |
| def function_setup(request, monkeypatch, tmp_path): |
| # Log function docstring (test description) if present |
| if request.function.__doc__ is not None: |
| logging.info(request.function.__doc__) |
| |
| # Start each test inside the tmp_path |
| monkeypatch.chdir(tmp_path) |
| |
| |
| @pytest.fixture(scope="class", autouse=True) |
| def class_setup(request): |
| # Log class docstring (test description) if present |
| if request.cls.__doc__ is not None: |
| logging.info(request.cls.__doc__) |
| |
| |
| def pytest_keyboard_interrupt(excinfo): |
| """Called for keyboard interrupt""" |
| module_teardown() |
| |
| |
| @pytest.fixture(scope="module") |
| def mpi_program(module_setup): |
| """Create the MPI program from the mpi_program.c in scripts directory. |
| Returns the bin path of the mpi_program.""" |
| |
| # Check for MPI setup |
| atf.require_mpi("pmix", "mpicc") |
| |
| # Use the external C source file |
| src_path = atf.properties["testsuite_scripts_dir"] + "/mpi_program.c" |
| bin_path = os.getcwd() + "/mpi_program" |
| |
| # Compile the MPI program |
| atf.run_command(f"mpicc -o {bin_path} {src_path}", fatal=True) |
| |
| yield bin_path |
| |
| atf.run_command(f"rm -f {bin_path}", fatal=True) |
| |
| |
| @pytest.fixture(scope="module") |
| def use_memory_program(module_setup): |
| """ |
| Returns the bin path of a program that allocates a certain amount of MB for some seconds. |
| """ |
| |
| atf.require_tool("python3") |
| |
| src_path = atf.properties["testsuite_scripts_dir"] + "/use_memory_program.py" |
| bin_path = os.getcwd() + "/use_memory_program.py" |
| |
| # Ensure x permissions |
| atf.run_command(f"cp {src_path} {bin_path}") |
| atf.run_command(f"chmod a+x {bin_path}") |
| |
| yield bin_path |
| |
| atf.run_command(f"rm -f {bin_path}", fatal=True) |
| |
| |
| @pytest.fixture(scope="module") |
| def spank_fail_lib(module_setup): |
| """ |
| Returns the bin path of the spank .so that will fail if configured. |
| """ |
| |
| # The plugin uses ESPANK_NODE_FAILURE, so it needs to compile against 25.05+ |
| # It also needs to be built against the same version of slurmd and submit |
| # clients like sbatch |
| new_prefixes = False |
| if not atf.is_upgrade_setup(): |
| atf.require_version((25, 5), "config.h") |
| else: |
| slurmd_version = atf.get_version("sbin/slurmd") |
| sbatch_version = atf.get_version("bin/sbatch") |
| |
| if slurmd_version != sbatch_version: |
| pytest.skip( |
| f"We need to build SPANK against Slurm version of submit clients as sbatch {sbatch_version} and slurmd {slurmd_version}, but they diffear." |
| ) |
| |
| if slurmd_version < (25, 5): |
| pytest.skip( |
| f"This SPANK plugin needs a Slurm 25.05+, but slurmd version is {slurmd_version}" |
| ) |
| |
| if ( |
| atf.get_version("config.h", slurm_prefix=atf.properties["new-build-prefix"]) |
| == slurmd_version |
| ): |
| new_prefixes = True |
| elif ( |
| not atf.get_version( |
| "config.h", slurm_prefix=atf.properties["old-build-prefix"] |
| ) |
| == slurmd_version |
| ): |
| # This should never happen, slurmd should be one of those versions |
| pytest.fail( |
| "Unable to find build dir to match slurmd version {slurmd_version}" |
| ) |
| |
| src_path = atf.properties["testsuite_scripts_dir"] + "/spank_fail_test.c" |
| bin_path = os.getcwd() + "/spank_fail_test.so" |
| |
| atf.compile_against_libslurm( |
| src_path, bin_path, full=True, shared=True, new_prefixes=new_prefixes |
| ) |
| |
| yield bin_path |
| |
| atf.run_command(f"rm -f {bin_path}", fatal=True) |
| |
| |
| @pytest.fixture(scope="module") |
| def spank_tmp_lib(module_setup): |
| """ |
| Compiles a SPANK plugin that will write files in a /tmp directory. |
| Returns the tmp_spank dir and the bin path of the spank .so that will write |
| files in the tmp_spank dir if configured. |
| """ |
| |
| # The plugin uses ESPANK_NODE_FAILURE, so it needs to compile against 25.05+ |
| # It also needs to be built against the same version of slurmd and submit |
| # clients like sbatch |
| new_prefixes = False |
| if atf.is_upgrade_setup(): |
| slurmd_version = atf.get_version("sbin/slurmd") |
| sbatch_version = atf.get_version("bin/sbatch") |
| |
| if slurmd_version != sbatch_version: |
| pytest.skip( |
| f"We need to build SPANK against Slurm version of submit clients as sbatch {sbatch_version} and slurmd {slurmd_version}, but they diffear." |
| ) |
| if ( |
| atf.get_version("config.h", slurm_prefix=atf.properties["new-build-prefix"]) |
| == slurmd_version |
| ): |
| new_prefixes = True |
| elif ( |
| not atf.get_version( |
| "config.h", slurm_prefix=atf.properties["old-build-prefix"] |
| ) |
| == slurmd_version |
| ): |
| # This should never happen, slurmd should be one of those versions |
| pytest.fail( |
| "Unable to find build dir to match slurmd version {slurmd_version}" |
| ) |
| |
| src_path = atf.properties["testsuite_scripts_dir"] + "/spank_tmp_plugin.c" |
| bin_path = os.getcwd() + "/spank_tmp_plugin.so" |
| |
| atf.compile_against_libslurm( |
| src_path, bin_path, full=True, shared=True, new_prefixes=new_prefixes |
| ) |
| |
| tmp_spank = "/tmp/spank" |
| atf.run_command(f"mkdir -p {tmp_spank}", fatal=True) |
| |
| yield tmp_spank, bin_path |
| |
| atf.run_command(f"rm -f {bin_path}", fatal=True) |
| atf.run_command(f"rm -rf {tmp_spank}", fatal=True) |
