##############################################################################
# ATF (Automated Testing Framework) python testsuite module
# Copyright (C) SchedMD LLC.
##############################################################################
import collections
import errno

# import glob
import logging
import math
import os
import pwd
import pathlib
import pytest
import re
import shutil
import stat
import socket
import subprocess
import sys
import time
import traceback
import copy

# slurmrestd
import requests
import signal

import json
import jsondiff
import jsonpatch
import yaml

# This module will be (un)imported in require_openapi_generator()
openapi_client = None
import importlib
import threading

##############################################################################
# ATF module functions
##############################################################################

default_command_timeout = 60
default_polling_timeout = 45
default_sql_cmd_timeout = 120

PERIODIC_TIMEOUT = 30


def get_open_port():
    """Finds an open port.

    Warning: Race conditions abound so be ready to retry calling function;

    Example:
        >>> while not some_test(port):
        >>>     port = get_open_port()

    Shamelessly based on:
    https://stackoverflow.com/questions/2838244/get-open-tcp-port-in-python
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind(("", 0))
    s.listen(1)
    port = s.getsockname()[1]
    s.close()
    return port


def node_range_to_list(node_expression):
    """Converts a node range expression into a list of node names.

    Example:
        >>> node_range_to_list('node[1,3-5]')
        ['node1', 'node3', 'node4', 'node5']
    """

    node_list = []
    output = run_command_output(
        f"scontrol show hostnames {node_expression}", fatal=True, quiet=True
    )
    for line in output.rstrip().splitlines():
        node_list.append(line)
    return node_list


def node_list_to_range(node_list):
    """Converts a list of node names to a node range expression.

    Example:
        >>> node_list_to_range(['node1', 'node3', 'node4', 'node5'])
        'node[1,3-5]'
    """

    return run_command_output(
        f"scontrol show hostlistsorted {','.join(node_list)}", fatal=True, quiet=True
    ).rstrip()


def range_to_list(range_expression):
    """Converts an integer range expression into a list of integers.

    Example:
        >>> range_to_list('1,3-5')
        [1, 3, 4, 5]
    """

    return list(map(int, node_range_to_list(f"[{range_expression}]")))


def list_to_range(numeric_list):
    """Converts a list of integers to an integer range expression.

    Example:
        >>> list_to_range([1, 3, 4, 5])
        '1,3-5'
    """

    node_range_expression = node_list_to_range(map(str, numeric_list))
    return re.sub(r"^\[(.*)\]$", r"\1", node_range_expression)


def run_command(
    command,
    fatal=False,
    timeout=default_command_timeout,
    quiet=False,
    chdir=None,
    user=None,
    input=None,
    xfail=False,
    env_vars=None,
    background=False,
):
    """Executes a command and returns a dictionary result.

    Args:
       command (string): The command to execute. The command is run within a
           bash subshell, so pipes, redirection, etc. are performed.
       fatal (boolean): If True, a non-zero exit code (or zero if combined
           with xfail) will result in the test failing.
       timeout (integer): If the command does not exit before timeout number
           of seconds, this function will return with an exit code of 110.
       quiet (boolean): If True, logging is performed at the TRACE log level.
       chdir (directory): Change to the specified directory before executing
           the command.
       user (user name): Run the command as the specified user. This requires
           the invoking user to have unprompted sudo rights.
       input (string): The specified input is supplied to the command as stdin.
       xfail (boolean): If True, the command is expected to fail.
       env_vars (string): A string to set environmental variables that is
            prepended to the command when run.
       background (boolean): If True, runs command in background and returns
            immediately. Process can be accessed via results["process"]. Cannot
            be used with fatal, xfail, timeout, capture_output, or check.

    Returns:
        A dictionary containing the following keys:
            start_time: epoch start time
            duration: number of seconds the command ran for
            exit_code: exit code for the command
            stdout: command stdout as a string
            stderr: command stderr as a string

    Example:
        >>> run_command('ls -l', fatal=True)
        {'command': 'ls -l', 'start_time': 1712268971.532, 'duration': 0.007, 'exit_code': 0, 'stdout': 'total 124\n-rw-rw-r-- 1 slurm slurm 118340 Apr  4 22:15 atf.py\n-rw-rw-r-- 1 slurm slurm    498 Apr  4 22:09 test.py\n-rw-rw-r-- 1 slurm slurm   1013 Apr  4 22:09 python_script.py\n', 'stderr': ''}

        >>> run_command('ls /non/existent/path', xfail=True)
        {'command': 'ls /non/existent/path', 'start_time': 1712269123.2, 'duration': 0.005, 'exit_code': 2, 'stdout': '', 'stderr': "ls: cannot access '/non/existent/path': No such file or directory\n"}

        >>> run_command('sleep 5', timeout=2)
        {'command': 'sleep 5', 'start_time': 1712269157.113, 'duration': 2.0, 'exit_code': 110, 'stdout': '', 'stderr': ''}
    """

    additional_run_kwargs = {}
    if chdir is not None:
        additional_run_kwargs["cwd"] = chdir
    if input is not None:
        additional_run_kwargs["input"] = input
    if timeout is not None:
        additional_run_kwargs["timeout"] = timeout

    if quiet:
        log_command_level = logging.TRACE
        log_details_level = logging.TRACE
    else:
        log_command_level = logging.NOTE
        log_details_level = logging.DEBUG

    if env_vars is not None:
        command = env_vars.strip() + " " + command

    # If user is not specified but test-user is set, then set user to test-user
    if not user and properties["test-user-set"]:
        user = properties["test-user"]

    start_time = time.time()
    invocation_message = "Running command"
    if user is not None:
        invocation_message += f" as user {user}"
    invocation_message += f": {command}"
    logging.log(log_command_level, invocation_message)
    try:
        if user is not None:
            if not properties["sudo-rights"]:
                pytest.skip(
                    "This test requires the test user to have unprompted sudo rights",
                    allow_module_level=True,
                )
            # Use su to honor ulimits, specially core
            cmd = [
                "sudo",
                "--preserve-env=PATH",
                "-u",
                user,
                "/bin/bash",
                "-lc",
                command,
            ]
        else:
            cmd = command

        # Run in background mode and return before completion if requested
        if background:
            # Create kwargs for Popen, excluding unsupported run() params
            background_run_kwargs = additional_run_kwargs.copy()
            # Check that parameters that don't make sense haven't been specified
            if background_run_kwargs["timeout"] != default_command_timeout:
                pytest.fail("Unsupported: timeout cannot be used with background=True")
            del background_run_kwargs["timeout"]
            if fatal:
                pytest.fail("Unsupported: fatal cannot be used with background=True")
            if xfail:
                pytest.fail("Unsupported: xfail cannot be used with background=True")
            if "input" in background_run_kwargs:
                input_str = background_run_kwargs.pop("input")
            else:
                input_str = None

            # Run command differently depending on if using sudo for security
            if isinstance(cmd, list):
                process = subprocess.Popen(
                    cmd,
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    universal_newlines=True,
                    **background_run_kwargs,
                )
            else:
                process = subprocess.Popen(
                    cmd,
                    shell=True,
                    executable="/bin/bash",
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    universal_newlines=True,
                    **background_run_kwargs,
                )

            # Handle input if provided
            if input_str is not None:
                process.stdin.write(input_str)
                process.stdin.close()

            # Return early with process object
            results = {
                "command": command,
                "start_time": float(int(start_time * 1000)) / 1000,
                "duration": float(int((time.time() - start_time) * 1000)) / 1000,
                "exit_code": None,
                "stdout": "",
                "stderr": "",
                "process": process,
            }
            return results

        # Synchronous code path
        if isinstance(cmd, list):
            cp = subprocess.run(
                cmd, capture_output=True, text=True, **additional_run_kwargs
            )
        else:
            cp = subprocess.run(
                cmd,
                shell=True,
                executable="/bin/bash",
                capture_output=True,
                text=True,
                **additional_run_kwargs,
            )
        end_time = time.time()
        duration = end_time - start_time
        exit_code = cp.returncode
        stdout = cp.stdout
        stderr = cp.stderr
    except subprocess.TimeoutExpired as e:
        duration = e.timeout
        exit_code = errno.ETIMEDOUT
        # These are byte objects, not strings
        stdout = e.stdout.decode("utf-8") if e.stdout else ""
        stderr = e.stderr.decode("utf-8") if e.stderr else ""

    if input is not None:
        logging.log(log_details_level, f"Command input: {input}")
    logging.log(log_details_level, f"Command exit code: {exit_code}")
    logging.log(log_details_level, f"Command stdout: {stdout}")
    logging.log(log_details_level, f"Command stderr: {stderr}")
    logging.log(log_details_level, "Command duration: %.03f seconds", duration)

    message = ""
    if exit_code == errno.ETIMEDOUT:
        message = f'Command "{command}" timed out after {duration} seconds'
    elif exit_code != 0 and not xfail:
        message = f'Command "{command}" failed with rc={exit_code}'
    elif exit_code == 0 and xfail:
        message = f'Command "{command}" was expected to fail but succeeded'
    if (exit_code != 0 and not xfail) or (exit_code == 0 and xfail):
        if stderr != "" or stdout != "":
            message += ":"
        if stderr != "":
            message += f" {stderr}"
        if stdout != "":
            message += f" {stdout}"

    if message != "":
        message = message.rstrip()
        if fatal:
            pytest.fail(message)
        elif not quiet:
            logging.warning(message)

    results = {}
    results["command"] = command
    results["start_time"] = float(int(start_time * 1000)) / 1000
    results["duration"] = float(int(duration * 1000)) / 1000
    results["exit_code"] = exit_code
    results["stdout"] = stdout
    results["stderr"] = stderr

    return results


def run_command_error(command, **run_command_kwargs):
    """Executes a command and returns the standard error.

    This function accepts the same arguments as run_command.

    Args:
        command (string): The command to execute. The command is run within a
            bash subshell, so pipes, redirection, etc. are performed.

    Returns:
        The standard error (stderr) output of the command as a string.

    Example:
        >>> run_command_error('ls /non/existent/path')
        "ls: cannot access '/non/existent/path': No such file or directory\n"

        >>> run_command_error('grep foo /etc/passwd', quiet=True)
        ''

        >>> run_command_error('echo error message >&2', xfail=True)
        'error message\n'
    """

    results = run_command(command, **run_command_kwargs)

    return results["stderr"]


def run_command_output(command, **run_command_kwargs):
    """Executes a command and returns the standard output.

    This function accepts the same arguments as run_command.

    Args:
        command (string): The command to execute. The command is run within a
            Bash subshell, so pipes, redirection, etc. are performed.

    Returns:
        The standard output (stdout) of the command as a string.

    Example:
        >>> run_command_output('ls')
        'file1.txt\nfile2.txt\nscript.py\n'

        >>> run_command_output('echo Hello, World!')
        'Hello, World!\n'

        >>> run_command_output('grep foo /etc/passwd', xfail=True)
        ''
    """

    results = run_command(command, **run_command_kwargs)

    return results["stdout"]


def run_command_exit(command, **run_command_kwargs):
    """Executes a command and returns the exit code.

    This function accepts the same arguments as run_command.

    Args:
        command (string): The command to execute. The command is run within a
            Bash subshell, so pipes, redirection, etc. are performed.

    Returns:
        The exit code of the command as an integer.

    Example:
        >>> run_command_exit('ls')
        0

        >>> run_command_exit('grep foo /etc/passwd', xfail=True)
        1

        >>> run_command_exit('sleep 5', timeout=2)
        110
    """

    results = run_command(command, **run_command_kwargs)

    return results["exit_code"]


def repeat_until(
    callable,
    condition,
    timeout=default_polling_timeout,
    poll_interval=None,
    xfail=False,
    fatal=False,
):
    """Repeats a callable until a condition is met or it times out.

    The callable returns an object that the condition operates on.

    Args:
        callable (callable): Repeatedly called until the condition is met or
            the timeout is reached.
        condition (callable): A callable object that returns a boolean. This
            function will return True when the condition call returns True.
        timeout (integer): If timeout number of seconds expires before the
            condition is met, return False.
        poll_interval (float): Number of seconds to wait between condition
            polls. This may be a decimal fraction. The default poll interval
            depends on the timeout used, but varies between .1 and 1 seconds.
        xfail (boolean): If True, a timeout is expected.
        fatal (boolean): If True, the test will fail if condition is not met
            (or if condition is met with xfail).

    Returns:
        True if the condition is met by the timeout, False otherwise.

    Example:
        >>> repeat_until(lambda : random.randint(1,10), lambda n: n == 5, timeout=30, poll_interval=1)
        True
    """

    begin_time = time.time()

    if poll_interval is None:
        if timeout <= 5:
            poll_interval = 0.1
        elif timeout <= 10:
            poll_interval = 0.2
        else:
            poll_interval = 1

    condition_met = False
    while time.time() < begin_time + timeout:
        if condition(callable()):
            condition_met = True
            break
        time.sleep(poll_interval)

    if not xfail and not condition_met:
        if fatal:
            pytest.fail(f"Condition was not met within the {timeout} second timeout")
        else:
            logging.warning(
                f"Condition was not met within the {timeout} second timeout"
            )
    elif xfail and condition_met:
        if fatal:
            pytest.fail(
                f"Condition was met within the {timeout} second timeout and wasn't expected"
            )
        else:
            logging.warning(
                f"Condition was met within the {timeout} second timeout and wasn't expected"
            )

    return condition_met


def repeat_command_until(command, condition, quiet=True, **repeat_until_kwargs):
    """Repeats a command until a condition is met or it times out.

    This function accepts the same arguments as repeat_until.

    Args:
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        True if the condition is met by the timeout, False otherwise.

    Example:
        >>> repeat_command_until("scontrol ping", lambda results: re.search(r'is UP', results['stdout']))
        True
    """

    return repeat_until(
        lambda: run_command(command, quiet=quiet), condition, **repeat_until_kwargs
    )


def pids_from_exe(executable):
    """Finds process IDs (PIDs) of running processes with given executable name.

    Args:
        executable (string): The name of the executable for which to find running processes.

    Returns:
        A list of integer process IDs (PIDs) for running processes that match the given
        executable name.

    Example:
        >>> pids_from_exe('/usr/bin/python3')
        [12345, 67890]

        >>> pids_from_exe('/usr/sbin/sshd')
        [54321]

        >>> pids_from_exe('/bin/non-existent')
        []
    """

    # We have to elevate privileges here, but forking off thousands of sudo
    # commands is expensive, so we will sudo a dynamic bash script for speed
    script = f"""cd /proc
for pid in `ls -d1 [0-9]*`;
do if [[ "$(readlink $pid/exe)" = "{executable}" ]];
   then echo $pid;
   fi
done"""
    pids = []
    output = run_command_output(script, user="root", quiet=True)
    for line in output.rstrip().splitlines():
        pids.append(int(line))
    return pids


def is_slurmrestd_running():
    """Checks if slurmrestd is running.
    Needs to be run after the related properties are set.
    """
    # TODO: We could check also if the required plugins/parsers in properties
    #       are in the returned specs, but the format still depends on the version.
    #       Once v0.0.39 is removed, we could add the extra check.
    return repeat_until(
        lambda: request_slurmrestd("openapi/v3"), lambda r: r.status_code == 200
    )


def is_slurmctld_running(quiet=False):
    """Checks whether slurmctld is running.

    Args:
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        True if the slurmctld is running, False otherwise.

    Example:
        >>> is_slurmctld_running()
        True

        >>> is_slurmctld_running(quiet=True)
        True
    """

    # Check whether slurmctld is running
    if re.search(r"is UP", run_command_output("scontrol ping", quiet=quiet)):
        return True

    return False


def gcore(component, pid=None, sbin=True):
    """Generates a gcore file for all pids running of a given Slurm component.

    The gcore file will be save at slurm-logs-dir, where all the logs and kernel
    coredumps should be generated too.

    Args:
        component (string): The name of the component. E.g. slurmctld, slurmdbd...
        pid (integer): The specific PID of the process to gcore. All PIDs of component are gcored by default...
        sbin: If True search for pids related to slurm-sbin-dir, or slurm-bin-dir otherwise.

    Returns:
        None
    """
    # Ensure that slurm-logs-dir is set.
    if "slurm-logs-dir" not in properties:
        properties["slurm-logs-dir"] = os.path.dirname(
            get_config_parameter("SlurmctldLogFile", live=False, quiet=True)
        )

    if sbin:
        prefix = properties["slurm-sbin-dir"]
    else:
        prefix = properties["slurm-bin-dir"]

    pids = pids_from_exe(f"{prefix}/{component}")

    if pid:
        if pid not in pids:
            logging.warning(
                f"Requested PID {pid} is not in the obtained PIDs ({pids}), but using it anyway"
            )
        pids = [pid]

    if not pids:
        logging.warning(f"Process {prefix}/{component} not found")
    logging.debug(f"Getting gcores for PIDs: {pids}")
    for pid in pids:
        run_command(
            f"sudo gcore -o {properties['slurm-logs-dir']}/{component}.core {pid}"
        )


def start_slurmctld(clean=False, quiet=False, also_slurmds=False):
    """Starts the Slurm controller daemon (slurmctld).

    This function may only be used in auto-config mode.

    Args:
        clean (boolean): If True, clears previous slurmctld state.
        quiet (boolean): If True, logging is performed at the TRACE log level.
        also_slurmds (boolean): If True, also start all required slurmds.

    Returns:
        None

    Example:
        >>> start_slurmctld()  # Start slurmctld with default settings
        >>> start_slurmctld(clean=True, quiet=True)  # Start slurmctld with clean state and quiet logging
    """

    if not properties["auto-config"]:
        require_auto_config("wants to start slurmctld")

    logging.debug("Starting slurmctld...")

    if not is_slurmctld_running(quiet=quiet):
        # Start slurmctld
        command = f"{properties['slurm-sbin-dir']}/slurmctld"
        if clean:
            command += " -c -i"
        results = run_command(command, user=properties["slurm-user"], quiet=quiet)
        if results["exit_code"] != 0:
            pytest.fail(
                f"Unable to start slurmctld (rc={results['exit_code']}): {results['stderr']}"
            )

        # Verify that slurmctld is running
        if not repeat_command_until(
            "scontrol ping", lambda results: re.search(r"is UP", results["stdout"])
        ):
            logging.warning(
                "scontrol ping is not responding, trying to get slurmctld core file..."
            )
            gcore("slurmctld")
            pytest.fail("Slurmctld is not running")
        else:
            logging.debug("Slurmctld started successfully")
    else:
        logging.warning("Slurmctld was already started")

    if also_slurmds:
        # Build list of slurmds
        slurmd_list = []
        output = run_command_output(
            f"perl -nle 'print $1 if /^NodeName=(\\S+)/' {properties['slurm-config-dir']}/slurm.conf",
            user=properties["slurm-user"],
            quiet=quiet,
        )
        if not output:
            pytest.fail("Unable to determine the slurmd node names")
        for node_name_expression in output.rstrip().split("\n"):
            if node_name_expression != "DEFAULT":
                slurmd_list.extend(node_range_to_list(node_name_expression))

        # (Multi)Slurmds
        for slurmd_name in slurmd_list:
            logging.debug(f"Starting slurmd for {slurmd_name}...")
            # Check whether slurmd is running
            slurmd_pgrep = run_command(
                f"pgrep -f 'slurmd -N {slurmd_name}'", quiet=quiet
            )
            if slurmd_pgrep["exit_code"] != 0:
                # Start slurmd
                results = run_command(
                    f"{properties['slurm-sbin-dir']}/slurmd -N {slurmd_name}",
                    user="root",
                    quiet=quiet,
                )
                if results["exit_code"] != 0:
                    pytest.fail(
                        f"Unable to start slurmd -N {slurmd_name} (rc={results['exit_code']}): {results['stderr']}"
                    )

                # Verify that the slurmd is running
                if (
                    run_command_exit(f"pgrep -f 'slurmd -N {slurmd_name}'", quiet=quiet)
                    != 0
                ):
                    pytest.fail(f"Slurmd -N {slurmd_name} is not running")
            else:
                logging.warning(f"slurmd for {slurmd_name} already running")
                logging.warning(f"slurmd_pgrep['stdout']: {slurmd_pgrep['stdout']}")
                logging.warning(f"slurmd_pgrep['stderr']: {slurmd_pgrep['stderr']}")
                logging.warning(
                    f"slurmd_pgrep['exit_code']: {slurmd_pgrep['exit_code']}"
                )

        # Verify that the slurmd is registered correctly
        timeout = 30 + 5 * len(slurmd_list)
        if not repeat_until(
            lambda: get_nodes(quiet=True),
            lambda nodes: all(nodes[name]["state"] == ["IDLE"] for name in slurmd_list),
            timeout=timeout,
        ):
            nodes = get_nodes(quiet=True)
            non_idle = [
                name for name in slurmd_list if nodes[name]["state"] != ["IDLE"]
            ]
            logging.warning(
                f"Getting the core files of the still not IDLE slurmds ({non_idle})"
            )
            for node in non_idle:
                pid = run_command_output(
                    f"pgrep -f 'slurmd -N {slurmd_name}'", quiet=quiet
                ).strip()
                gcore("slurmd", pid)
            pytest.fail(f"Some nodes are not IDLE: {non_idle}")

        logging.debug(f"All nodes are IDLE: {slurmd_list}")


def start_slurmdbd(clean=False, quiet=False):
    """Starts the Slurm DB daemon (slurmdbd).

    This function may only be used in auto-config mode.

    Args:
        clean (boolean): If True, clears previous slurmdbd state.
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        None
    """
    if not properties["auto-config"]:
        require_auto_config("wants to start slurmdbd")

    logging.debug("Starting slurmdbd...")

    if (
        run_command_exit(
            "sacctmgr show cluster", user=properties["slurm-user"], quiet=quiet
        )
        == 0
    ):
        logging.warning("slurmdbd was already running")
    else:
        # Start slurmdbd
        results = run_command(
            f"{properties['slurm-sbin-dir']}/slurmdbd",
            user=properties["slurm-user"],
            quiet=quiet,
        )
        if results["exit_code"] != 0:
            pytest.fail(
                f"Unable to start slurmdbd (rc={results['exit_code']}): {results['stderr']}"
            )

        # Verify that slurmdbd is running
        if not repeat_command_until(
            "sacctmgr show cluster", lambda results: results["exit_code"] == 0
        ):
            logging.warning(
                "sacctmgr show cluster is not responding, trying to get slurmdbd core file..."
            )
            gcore("slurmdbd")
            pytest.fail("Slurmdbd is not running")
        else:
            logging.debug("Slurmdbd started successfully")


def start_slurm(clean=False, quiet=False):
    """Starts all applicable Slurm daemons.

    This function examines the Slurm configuration files to determine which daemons
    need to be started.

    This function may only be used in auto-config mode.

    Args:
        clean (boolean): If True, clears previous slurmctld state.
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        None

    Example:
        >>> start_slurm()  # Start all Slurm daemons with default settings
        >>> start_slurm(clean=True, quiet=True)  # Start all Slurm daemons with clean state and quiet logging
    """

    if not properties["auto-config"]:
        require_auto_config("wants to start slurm")

    # Determine whether slurmdbd should be included
    if (
        get_config_parameter("AccountingStorageType", live=False, quiet=quiet)
        == "accounting_storage/slurmdbd"
    ):
        start_slurmdbd(clean, quiet)

    # Remove unnecessary default node0 from config to avoid being used or reserved
    output = run_command_output(
        f"cat {properties['slurm-config-dir']}/slurm.conf",
        user=properties["slurm-user"],
        quiet=quiet,
    )
    if len(re.findall(r"NodeName=", output)) > 1:
        run_command(
            f"sed -i '/NodeName=node0 /d' {properties['slurm-config-dir']}/slurm.conf",
            user=properties["slurm-user"],
            quiet=quiet,
        )

    # Start slurmctld
    start_slurmctld(clean, quiet, also_slurmds=True)

    # Start slurmrestd if required
    if properties["slurmrestd-started"]:
        start_slurmrestd()


def stop_slurmctld(quiet=False, also_slurmds=False):
    """Stops the Slurm controller daemon (slurmctld).

    This function may only be used in auto-config mode.

    Args:
        also_slurmds (boolean): If True, also stop all slurmds.
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        None

    Example:
        >>> stop_slurmctld()  # Stop slurmctld with default logging
        >>> stop_slurmctld(quiet=True)  # Stop slurmctld with quiet logging
    """

    rc = None
    failures = []

    if not properties["auto-config"]:
        require_auto_config("wants to stop slurmctld")

    # Stop slurmctld
    command = "scontrol shutdown"
    if not also_slurmds:
        command += " slurmctld"
        logging.debug("Stopping slurmctld...")
    else:
        logging.debug("Stopping slurmctld and slurmds...")

    results = run_command(command, user=properties["slurm-user"], quiet=quiet)
    if results["exit_code"] != 0:
        failures.append(f"Command {command} failed with rc={results['exit_code']}")

    # Verify that slurmctld is not running
    if not repeat_until(
        lambda: pids_from_exe(f"{properties['slurm-sbin-dir']}/slurmctld"),
        lambda pids: len(pids) == 0,
        timeout=default_polling_timeout * 2,
    ):
        failures.append("Slurmctld is still running")
        logging.warning("Getting the core files of the still running slurmctld")
        gcore("slurmctld")
    else:
        logging.debug("No slurmctld is running.")

    if also_slurmds:
        if get_version("sbin/slurmd") < (24, 11):
            # FIXED: t20764.
            slurmd_list = []
            output = run_command_output(
                f"perl -nle 'print $1 if /^NodeName=(\\S+)/' {properties['slurm-config-dir']}/slurm.conf",
                quiet=quiet,
            )
            if not output:
                failures.append("Unable to determine the slurmd node names")
            else:
                for node_name_expression in output.rstrip().split("\n"):
                    if node_name_expression != "DEFAULT":
                        slurmd_list.extend(node_range_to_list(node_name_expression))

            for slurmd_name in slurmd_list:
                run_command(
                    f"sudo systemctl stop {slurmd_name}_slurmstepd.scope", quiet=quiet
                )

        # Verify that slurmds are not running
        if not repeat_until(
            lambda: pids_from_exe(f"{properties['slurm-sbin-dir']}/slurmd"),
            lambda pids: len(pids) == 0,
            timeout=default_polling_timeout * 2,
        ):
            logging.warning("Some slurmd was not stopped.")
            logging.warning(
                "Sometimes slurm may be just rebooting or similar when shutdown was run"
            )

            logging.debug("Shutting down slurmds manually")
            for pid in pids_from_exe(f"{properties['slurm-sbin-dir']}/slurmd"):
                run_command(f"kill {pid}", user="root")

            if not repeat_until(
                lambda: pids_from_exe(f"{properties['slurm-sbin-dir']}/slurmd"),
                lambda pids: len(pids) == 0,
                timeout=default_polling_timeout * 2,
            ):
                failures.append("Some slurmds are still running")
                logging.warning("Getting the core files of the still running slurmds")
                gcore("slurmd")
            else:
                logging.debug("No slurmd is running.")
        else:
            logging.debug("No slurmd is running.")

    if failures:
        rc = failures

    return rc


def stop_slurmdbd(quiet=False):
    """Stops the Slurm DB daemon (slurmdbd).

    This function may only be used in auto-config mode.

    Args:
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        None
    """

    rc = None
    failures = []

    if not properties["auto-config"]:
        require_auto_config("wants to stop slurmdbd")

    logging.debug("Stopping slurmdbd...")

    # Stop slurmdbd
    results = run_command(
        "sacctmgr shutdown", user=properties["slurm-user"], quiet=quiet
    )
    if results["exit_code"] != 0:
        failures.append(
            f"Command \"sacctmgr shutdown\" failed with rc={results['exit_code']}"
        )

    # Verify that slurmdbd is not running (we might have to wait for rollups to complete)
    if not repeat_until(
        lambda: pids_from_exe(f"{properties['slurm-sbin-dir']}/slurmdbd"),
        lambda pids: len(pids) == 0,
        timeout=default_polling_timeout * 2,
    ):
        failures.append("Slurmdbd is still running")
        logging.warning("Getting the core files of the still running slurmdbd")
        gcore("slurmdbd")
    else:
        logging.debug("No slurmdbd is running.")

    if failures:
        rc = failures

    return rc


def stop_slurm(fatal=True, quiet=False):
    """Stops all applicable Slurm daemons.

    This function examines the Slurm configuration files to determine which daemons
    need to be stopped.

    This function may only be used in auto-config mode.

    Args:
        fatal (boolean): If True, a failure to stop all daemons will result in the
            test failing.
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        True if all Slurm daemons were stopped, False otherwise.

    Example:
        >>> stop_slurm()  # Stop all Slurm daemons with default settings
        True

        >>> stop_slurm(fatal=False, quiet=True)  # Stop all Slurm daemons with non-fatal failures and quiet logging
        False
    """

    failures = []

    if not properties["auto-config"]:
        require_auto_config("wants to stop slurm")

    # Determine whether slurmdbd should be included
    if (
        get_config_parameter("AccountingStorageType", live=False, quiet=quiet)
        == "accounting_storage/slurmdbd"
    ):
        err = stop_slurmdbd(quiet)
        if err:
            failures.extend(err)

    # Stop slurmctld and slurmds
    err = stop_slurmctld(quiet=quiet, also_slurmds=True)
    if err:
        failures.extend(err)

    # Stop slurmrestd if was started
    if properties["slurmrestd-started"]:
        properties["slurmrestd"].send_signal(signal.SIGINT)
        try:
            properties["slurmrestd"].wait(timeout=60)
        except Exception:
            properties["slurmrestd"].kill()
        properties["slurmrestd_log"].close()

    if failures:
        for fail in failures:
            logging.warning(fail)
        if fatal:
            pytest.fail(failures[0])

        return False
    else:
        return True


def restart_slurmctld(clean=False, quiet=False):
    """Restarts the Slurm controller daemon (slurmctld).

    This function may only be used in auto-config mode.

    Args:
        clean (boolean): If True, clears previous slurmctld state.
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        None

    Example:
        >>> restart_slurmctld()  # Restart slurmctld with default settings
        >>> restart_slurmctld(clean=True, quiet=True)  # Restart slurmctld with clean state and quiet logging
    """

    stop_slurmctld(quiet=quiet)
    start_slurmctld(clean=clean, quiet=quiet)


def restart_slurm(clean=False, quiet=False):
    """Restarts all applicable Slurm daemons.

    This function may only be used in auto-config mode.

    Args:
        clean (boolean): If True, clears previous slurmctld state.
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        None

    Example:
        >>> restart_slurm()  # Restart all Slurm daemons with default settings
        >>> restart_slurm(clean=True, quiet=True)  # Restart all Slurm daemons with clean state and quiet logging
    """

    stop_slurm(quiet=quiet)
    start_slurm(clean=clean, quiet=quiet)


def require_slurm_running():
    """Ensures that the Slurm daemons are running.

    In local-config mode, the test is skipped if Slurm is not running.
    In auto-config mode, Slurm is started if necessary.

    In order to avoid multiple restarts of Slurm (in auto-config), this function
    should be called at the end of the setup preconditions.

    Args:
        None

    Returns:
        None

    Example:
        >>> require_slurm_running()  # Ensure Slurm is running or start it in auto-config mode
    """

    global nodes

    if properties["auto-config"]:
        if not is_slurmctld_running(quiet=True):

            # Check and report the mixed components
            versions = dict()
            versions["slurmdbd"] = get_version("sbin/slurmdbd")
            versions["slurmctld"] = get_version("sbin/slurmctld")
            versions["slurmd"] = get_version("sbin/slurmd")
            versions["scontrol"] = get_version("bin/scontrol")
            if len(set(versions.values())) == 1:
                logging.info(
                    f"Starting Slurm with all components in the same version: {versions['slurmctld']}"
                )
            else:
                logging.info(f"Starting Slurm in a mixed version setup: {versions}")

            properties["slurm-started"] = True
            start_slurm(clean=True, quiet=True)
    else:
        if not is_slurmctld_running(quiet=True):
            pytest.skip(
                "This test requires slurm to be running", allow_module_level=True
            )

    # As a side effect, build up initial nodes dictionary
    nodes = get_nodes(quiet=True)

    config = get_config()
    properties["slurmctld_host"] = config["SlurmctldHost[0]"].split("(")[0]
    properties["slurmctld_port"] = config["SlurmctldPort"]


def is_upgrade_setup(
    old_slurm_prefix="/opt/slurm-old",
    new_slurm_prefix="/opt/slurm-new",
    old_build_prefix="",
    new_build_prefix="",
    old_source_prefix="",
    new_source_prefix="",
    force_old=False,
):
    """
    Return True if we have two Slurms configured in the system.
    """

    if not os.path.exists(old_slurm_prefix):
        logging.debug(f"Old prefix {old_slurm_prefix} not exists.")
        return False

    if not os.path.exists(new_slurm_prefix):
        logging.debug(f"New prefix {new_slurm_prefix} not exists.")
        return False

    # Add the right properties
    setup_upgrades(
        old_slurm_prefix,
        new_slurm_prefix,
        old_build_prefix,
        new_build_prefix,
        old_source_prefix,
        new_source_prefix,
        force_old,
    )
    return True


def require_upgrades(
    old_slurm_prefix="/opt/slurm-old",
    new_slurm_prefix="/opt/slurm-new",
    old_build_prefix="",
    new_build_prefix="",
    old_source_prefix="",
    new_source_prefix="",
    force_old=True,
):
    """Checks if has two different versions installed.

    If they are not, skip.
    """
    if not properties["auto-config"]:
        require_auto_config("to change/upgrade Slurm setup")

    if not is_upgrade_setup(
        old_slurm_prefix,
        new_slurm_prefix,
        old_build_prefix,
        new_build_prefix,
        old_source_prefix,
        new_source_prefix,
        force_old,
    ):
        pytest.skip("This test needs an upgrade setup")

    # Double-check that old_version <= new_version
    old_version = get_version(slurm_prefix=old_slurm_prefix)
    new_version = get_version(slurm_prefix=new_slurm_prefix)
    if old_version > new_version:
        pytest.skip(
            f"Old version ({old_version}) has to be older than new version ({new_version})"
        )
    logging.info(f"Required upgrade setup found: {old_version} and {new_version}")


def setup_upgrades(
    old_slurm_prefix="/opt/slurm-old",
    new_slurm_prefix="/opt/slurm-new",
    old_build_prefix="",
    new_build_prefix="",
    old_source_prefix="",
    new_source_prefix="",
    force_old=False,
):
    """
    Adds the necessary atf.properties[] with the old/new paths.
    If force_old is specified itt also update the links pointing to the old
    paths, and they will be restored in the global teardown.
    """
    # TODO: We should use slurm-new(-build) instead of slurm-git(-build)
    if old_build_prefix == "":
        old_build_prefix = properties["slurm-build-dir"]
    if new_build_prefix == "":
        new_build_prefix = f"{properties['slurm-build-dir']}/../slurm-git-build"
    if old_source_prefix == "":
        old_source_prefix = properties["slurm-source-dir"]
    if new_source_prefix == "":
        new_source_prefix = f"{properties['slurm-source-dir']}/../slurm-git"

    properties["old-slurm-prefix"] = old_slurm_prefix
    properties["new-slurm-prefix"] = new_slurm_prefix

    properties["old-build-prefix"] = old_build_prefix
    properties["new-build-prefix"] = new_build_prefix

    properties["old-source-prefix"] = old_source_prefix
    properties["new-source-prefix"] = new_source_prefix

    properties["forced_upgrade_setup"] = force_old

    if force_old:
        logging.debug(
            "Setting bin/ and sbin/ pointing to old version and saving a backup..."
        )
        run_command(
            f"sudo mv {properties['slurm-sbin-dir']} {module_tmp_path}/upgrade-sbin",
            quiet=True,
            fatal=True,
        )
        run_command(
            f"sudo mv {properties['slurm-bin-dir']} {module_tmp_path}/upgrade-bin",
            quiet=True,
            fatal=True,
        )
        run_command(
            f"sudo mkdir {properties['slurm-sbin-dir']} {properties['slurm-bin-dir']}",
            quiet=True,
            fatal=True,
        )
        run_command(
            f"sudo ln -s {properties['old-slurm-prefix']}/sbin/* {properties['slurm-sbin-dir']}/",
            quiet=True,
            fatal=True,
        )
        run_command(
            f"sudo ln -s {properties['old-slurm-prefix']}/bin/* {properties['slurm-bin-dir']}/",
            quiet=True,
            fatal=True,
        )


def upgrade_component(component, new_version=True):
    """Upgrades a component creating the required links, and restarts it if necessary.

    This function needs require_upgrades() to be already run to work properly.

    Args:
        component (string): The bin/ or sbin/ component of Slurm to check.
        new_version (boolean): Set it false to downgrade to the older version instead.

    Returns:
        A tuple representing the version. E.g. (25.05.0).
    """

    if (
        "old-slurm-prefix" not in properties.keys()
        or "new-slurm-prefix" not in properties.keys()
    ):
        pytest.fail("To upgrade_components() first we need to call require_upgrades()")

    if not os.path.exists(f"{properties['slurm-prefix']}/{component}"):
        pytest.fail(f"Unknown or not existing {component}")

    if new_version:
        upgrade_prefix = properties["new-slurm-prefix"]
    else:
        upgrade_prefix = properties["old-slurm-prefix"]

    # Stop components when necessary
    if component == "sbin/slurmdbd":
        stop_slurmdbd()
    elif component == "sbin/slurmctld":
        stop_slurmctld()

    run_command(
        f"sudo rm -f {properties['slurm-prefix']}/{component}",
        quiet=True,
        fatal=True,
    )
    run_command(
        f"sudo ln -s {upgrade_prefix}/{component} {properties['slurm-prefix']}/{component}",
        quiet=True,
        fatal=True,
    )

    # Restart components when necessary
    if component == "sbin/slurmdbd":
        start_slurmdbd()
    elif component == "sbin/slurmctld":
        start_slurmctld()


def get_version(component="sbin/slurmctld", slurm_prefix=""):
    """Returns the version of the Slurm component as a tuple.

    It calls the component with -V and converts the output into a tuple.

    Args:
        component (string): The bin/ or sbin/ component of Slurm to check.
                            It also supports "config.h" to obtain the VERSION in the header.
        slurm_prefix (string): The path where the component is. By default the defined in testsuite.conf.
                               If component is "config.h", then it's the build dir.

    Returns:
        A tuple representing the version. E.g. (25.05.0).
    """
    if component == "config.h":
        if slurm_prefix == "":
            slurm_prefix = properties["slurm-build-dir"]
        header = pathlib.Path(f"{slurm_prefix}/config.h")
        if not header.exists():
            pytest.fail("Unable to access to config.h to get Slurm version")

        version_str = re.search(
            r'#define\s+VERSION\s+"([^"]+)"', header.read_text()
        ).group(1)

    else:
        if slurm_prefix == "":
            slurm_prefix = f"{properties['slurm-sbin-dir']}/.."

        version_str = (
            run_command_output(
                f"{slurm_prefix}/{component} -V", quiet=True, user="root"
            )
            .strip()
            .replace("slurm ", "")
        )

    return tuple(int(part) if part.isdigit() else 0 for part in version_str.split("."))


def require_version(
    min_version,
    component="sbin/slurmctld",
    slurm_prefix="",
    reason=None,
    max_version=None,
):
    """Checks if the component is at least the required version, or skips.

    Args:
        min_version (tuple): The tuple representing the minimal version required, version should be >=.
        max_version (tuple): The tuple representing the maximal version required, version should be <.
        component (string): The bin/ or sbin/ component of Slurm to check.
        slurm_prefix (string): The path where the component is. By default the defined in testsuite.conf.
        reason (string): The reason why the version of the component is required.
    """
    component_version = get_version(component, slurm_prefix)
    if component_version >= min_version and (
        max_version is None or component_version < max_version
    ):
        return

    if not reason:
        reason = f"The version of {component} is {component_version}, required is {min_version}"
        if max_version:
            reason += f" up to {max_version} (not included)"
    pytest.skip(reason)


def request_slurmctld(request):
    """
    Returns the slurmctld response of a given request.
    It needs slurmctld >= 25.11 listening HTTP requests.
    """

    return requests.get(
        f"http://{properties['slurmctld_host']}:{properties['slurmctld_port']}/{request}",
    )


def request_slurmrestd(request):
    """Returns the slurmrestd response of a given request.
    It needs slurmrestd to be running (see require_slurmrestd())
    """
    return requests.get(
        f"{properties['slurmrestd_url']}/{request}",
        headers=properties["slurmrestd-headers"],
    )


def get_deprecated_openapi_spec_patch(openapi_spec, undeprecate=False):
    """
    Returns a JsonPatch object that can be applied to the openapi_spec to
    fully deprecate it, as we should do for older specs in newer Slurm
    versions.

    If undeprecate is set to True, then it does the opposite, it returns a
    JsonPatch to remove all deprecate marks of the openapi_specs. This
    should NOT be used in general, but as we generated some openapi_specs
    with newer Slurm versions, we may need to undeprecate them (see v41).
    """
    ops = []

    def _escape_json_pointer(s: str) -> str:
        return s.replace("~", "~0").replace("/", "~1")

    for raw_path, path_obj in openapi_spec.get("paths", {}).items():
        escaped_path = _escape_json_pointer(raw_path)

        for method, operation in path_obj.items():
            if method.lower() not in (
                "get",
                "post",
                "put",
                "delete",
                "patch",
                "head",
                "options",
            ):
                continue

            # 1. Mark the operation itself as deprecated
            if not undeprecate:
                ops.append(
                    {
                        "op": "add",
                        "path": f"/paths/{escaped_path}/{method}/deprecated",
                        "value": True,
                    }
                )
            else:
                ops.append(
                    {
                        "op": "remove",
                        "path": f"/paths/{escaped_path}/{method}/deprecated",
                    }
                )

            # 2. Mark every parameter as deprecated
            params = operation.get("parameters", [])
            for i, _ in enumerate(params):
                if not undeprecate:
                    ops.append(
                        {
                            "op": "add",
                            "path": f"/paths/{escaped_path}/{method}/parameters/{i}/deprecated",
                            "value": True,
                        }
                    )
                else:
                    ops.append(
                        {
                            "op": "remove",
                            "path": f"/paths/{escaped_path}/{method}/parameters/{i}/deprecated",
                        }
                    )

    return jsonpatch.JsonPatch(ops)


def assert_openapi_spec_eq(spec_a, spec_b):
    """
    Asserts that two OpenAPI specifications are equal, ignoring some info and
    description fields.

    Note that atf.properties["openapi_spec"] saves the specs used to generate
    the python openapi_client from slurmrestd and we have some json specs in
    the testsuite_data_dir that can be read with the openapi_spec feature.
    Usually we want to assert both are equal.

    Args:
        spec_a: one loaded json spec
        spec_b: the other loaded json spec

    Returns:
        None, but a pytest assert call.
    """

    # Copy specs so we can change some things
    _spec_a = copy.deepcopy(spec_a)
    _spec_b = copy.deepcopy(spec_b)

    # Avoid checking versions that change with each new release
    _spec_a["info"]["x-slurm"] = None
    _spec_b["info"]["x-slurm"] = None

    _spec_a["info"]["version"] = None
    _spec_b["info"]["version"] = None

    _spec_a["tags"] = None
    _spec_b["tags"] = None

    # Recursively remove all description fields
    def _strip_descriptions(oas):
        """Recursively removes all description fields in oas
        Returns:
            oas with all description fields set to None
        """

        if type(oas) is dict:
            for key, value in oas.items():
                if key.casefold() == "description".casefold():
                    oas[key] = None
                else:
                    oas[key] = _strip_descriptions(value)
        elif type(oas) is list:
            for value in oas:
                value = _strip_descriptions(value)

        return oas

    _spec_a = _strip_descriptions(_spec_a)
    _spec_b = _strip_descriptions(_spec_b)

    d = jsondiff.diff(_spec_a, _spec_b)
    assert (
        len(d) == 0
    ), f"OpenAPI specs should be equal(equivalent, but these differences were found: {d}"


def require_openapi_generator(version="7.3.0"):
    """Generates an OpenAPI client using OpenAPI-Generator, or skips if not available (even in auto-config).
    It needs slurmrestd to be running (see require_slurmrestd()).
    It also sets the necessary OPENAPI_GENERATOR_VERSION and JAVA_OPTS
    environment variables.
    Args:
        version (string): the required version.

    Returns:
        None
    """

    # Require specific testing version
    os.environ["OPENAPI_GENERATOR_VERSION"] = version

    # Work around: https://github.com/OpenAPITools/openapi-generator/issues/13684
    os.environ["JAVA_OPTS"] = (
        "--add-opens java.base/java.util=ALL-UNNAMED --add-opens java.base/java.lang=ALL-UNNAMED"
    )

    ogc_version = (
        run_command_output("openapi-generator-cli version").strip().split("\n")[-1]
    )
    if ogc_version != version:
        pytest.skip(
            f"test requires openapi-generator-cli version {version} (not {ogc_version})",
            allow_module_level=True,
        )

    # allow pointing to an existing OpenAPI generated client
    opath = module_tmp_path
    if "SLURM_TESTSUITE_OPENAPI_CLIENT" in os.environ:
        opath = os.environ["SLURM_TESTSUITE_OPENAPI_CLIENT"]

    pyapi_path = f"{opath}/pyapi/"
    spec_path = f"{opath}/openapi.json"

    # Always create path if needed
    os.makedirs(opath, exist_ok=True)

    if not os.path.exists(spec_path):
        r = request_slurmrestd("openapi/v3")
        if r.status_code != 200:
            pytest.fail(f"Error requesting openapi specs from slurmrestd: {r}")

        properties["openapi_spec"] = json.loads(r.text)
        if properties["openapi_spec"] is None:
            pytest.fail(f"Error parsing OpenAPI specs from slurmrestd: {r.text}")

        with open(spec_path, "w") as f:
            f.write(r.text)
            f.close()

    if not os.path.exists(pyapi_path):
        run_command(
            f"openapi-generator-cli generate -i '{spec_path}' -g python-pydantic-v1 --strict-spec=true -o '{pyapi_path}'",
            fatal=True,
            timeout=60,
        )

    sys.path.insert(0, pyapi_path)

    # Re-import openapi_client
    # Regular import doesn't work if was already imported by another test.
    global openapi_client
    module_name = "openapi_client"
    module_prefix = module_name + "."
    for mod in list(sys.modules):
        if mod == module_name or mod.startswith(module_prefix):
            del sys.modules[mod]
    openapi_client = importlib.import_module(module_name)
    importlib.reload(openapi_client)

    properties["openapi_config"] = openapi_client.Configuration()
    properties["openapi_config"].host = properties["slurmrestd_url"]
    properties["openapi_config"].access_token = properties["slurmrestd-headers"][
        "X-SLURM-USER-TOKEN"
    ]


def openapi_util():
    """
    Returns a UtilApi client from OpenAPI.
    It needs require_openapi_generator() to be run first.
    """
    return openapi_client.UtilApi(
        openapi_client.ApiClient(properties["openapi_config"])
    )


def openapi_slurm():
    """
    Returns a SlurmApi client from OpenAPI.
    It needs require_openapi_generator() to be run first.
    """
    return openapi_client.SlurmApi(
        openapi_client.ApiClient(properties["openapi_config"])
    )


def openapi_slurmdb():
    """
    Returns a SlurmdbApi client from OpenAPI.
    It needs require_openapi_generator() to be run first.
    """
    return openapi_client.SlurmdbApi(
        openapi_client.ApiClient(properties["openapi_config"])
    )


def get_config(live=True, source="slurm", quiet=False, delimiter="="):
    """Returns the Slurm configuration as a dictionary.

    Args:
        live (boolean):
            If True, the configuration information is obtained via
            a query to the relevant Slurm daemon (e.g., scontrol show config).
            If False, the configuration information is obtained by directly
            parsing the relevant Slurm configuration file (e.g. slurm.conf).
        source (string):
            If live is True, source should be either scontrol or sacctmgr.
            If live is False, source should be the name of the config file
            without the .conf prefix (e.g. slurmdbd).
        quiet (boolean): If True, logging is performed at the TRACE log level.
        delimiter (string): The delimiter between the parameter name and the value.

    Returns:
        A dictionary comprised of the parameter names and their values.
        For parameters that can have multiple lines and subparameters,
        the dictionary value will be a dictionary of dictionaries.

    Example:
        >>> get_config()
        {'AccountingStorageBackupHost': '(null)', 'AccountingStorageEnforce': 'none', 'AccountingStorageHost': 'localhost', 'AccountingStorageExternalHost': '(null)', 'AccountingStorageParameters': '(null)', 'AccountingStoragePort': '0', 'AccountingStorageTRES': 'cpu,mem,energy,node,billing,fs/disk,vmem,pages', ...}
        >>> get_config(live=False, source='slurm')
        {'SlurmctldHost': 'nathan-atf-docstrings', 'SlurmUser': 'slurm', 'SlurmctldLogFile': '/var/slurm/log/slurmctld.log', 'SlurmctldPidFile': '/var/slurm/run/slurmctld.pid', 'SlurmctldDebug': 'debug3', 'SlurmdLogFile': '/var/slurm/log/slurmd.%n.log', 'SlurmdPidFile': '/var/slurm/run/slurmd.%n.pid', ...}
        >>> get_config(live=True, source='scontrol', quiet=True)
        {'AccountingStorageBackupHost': '(null)', 'AccountingStorageEnforce': 'none', 'AccountingStorageHost': 'localhost', 'AccountingStorageExternalHost': '(null)', 'AccountingStorageParameters': '(null)', 'AccountingStoragePort': '0', 'AccountingStorageTRES': 'cpu,mem,energy,node,billing,fs/disk,vmem,pages', ...}
    """

    slurm_dict = {}

    if live:
        if source == "slurm" or source == "controller" or source == "scontrol":
            command = "scontrol"
        elif source == "slurmdbd" or source == "dbd" or source == "sacctmgr":
            command = "sacctmgr"
        else:
            pytest.fail(f"Invalid live source value ({source})")

        output = run_command_output(f"{command} show config", fatal=True, quiet=quiet)

        for line in output.splitlines():
            if match := re.search(rf"^\s*(\S+)\s*{re.escape(delimiter)}\s*(.*)$", line):
                slurm_dict[match.group(1)] = match.group(2).rstrip()
    else:
        config = source
        config_file = f"{properties['slurm-config-dir']}/{config}.conf"

        # We might be looking for parameters in a config file that has not
        # been created yet. If so, we just want this to return an empty dict
        output = run_command_output(
            f"cat {config_file}", user=properties["slurm-user"], quiet=quiet
        )
        for line in output.splitlines():
            if match := re.search(
                rf"^\s*(\S+?)\s*{re.escape(delimiter)}\s*(.*)$", line
            ):
                parameter_name, parameter_value = (
                    match.group(1),
                    match.group(2).rstrip(),
                )
                if parameter_name.lower() in [
                    "downnodes",
                    "name",
                    "nodename",
                    "nodeset",
                    "partitionname",
                    "switchname",
                ]:
                    instance_name, subparameters = parameter_value.split(" ", 1)
                    subparameters_dict = {}
                    for subparameter_name, subparameter_value in re.findall(
                        rf" *([^= ]+) *{re.escape(delimiter)} *([^ ]+)", subparameters
                    ):
                        # Reformat the value if necessary
                        if is_integer(subparameter_value):
                            subparameter_value = int(subparameter_value)
                        elif is_float(subparameter_value):
                            subparameter_value = float(subparameter_value)
                        elif subparameter_value == "(null)":
                            subparameter_value = None
                        subparameters_dict[subparameter_name] = subparameter_value
                    if parameter_name not in slurm_dict:
                        slurm_dict[parameter_name] = {}
                    slurm_dict[parameter_name][instance_name] = subparameters_dict
                else:
                    # Reformat the value if necessary
                    if is_integer(parameter_value):
                        parameter_value = int(parameter_value)
                    elif is_float(parameter_value):
                        parameter_value = float(parameter_value)
                    elif parameter_value == "(null)":
                        parameter_value = None
                    slurm_dict[parameter_name] = parameter_value

    return slurm_dict


def get_config_parameter(name, default=None, **get_config_kwargs):
    """Obtains the value for a Slurm configuration parameter.

    This function accepts the same arguments as get_config.

    Args:
        name (string): The parameter name.
        default (string or None): This value is returned if the parameter
            is not found.

    Returns:
        The value of the specified parameter, or the default if not found.

    Example:
        >>> get_config_parameter('JobAcctGatherFrequency')
        '30'
        >>> get_config_parameter('MaxJobCount', default='10000')
        '10000'
        >>> get_config_parameter('partitionname', default='debug', live=True, source='scontrol')
        'debug'
    """

    config_dict = get_config(**get_config_kwargs)

    # Convert keys to lower case so we can do a case-insensitive search
    lower_dict = dict(
        (key.casefold(), str(value).casefold()) for key, value in config_dict.items()
    )

    if name.casefold() in lower_dict:
        return lower_dict[name.casefold()]
    else:
        return default


def config_parameter_includes(name, value, **get_config_kwargs):
    """Checks whether a configuration parameter includes a specific value.

    When a parameter may contain a comma-separated list of values, this
    function can be used to determine whether a specific value is within
    the list.

    This function accepts the same arguments as get_config.

    Args:
        name (string): The parameter name.
        value (string): The value you are looking for.

    Returns:
        True if the specified string value is found within the parameter
        value list, False otherwise.

    Example:
        >>> config_parameter_includes('SlurmdParameters', 'config_overrides')
        False
    """

    config_dict = get_config(**get_config_kwargs)

    # Convert keys to lower case so we can do a case-insensitive search
    lower_dict = dict((key.casefold(), value) for key, value in config_dict.items())

    if name.casefold() in lower_dict and value.lower() in map(
        str.lower, lower_dict[name.casefold()].split(",")
    ):
        return True
    else:
        return False


def set_config_parameter(
    parameter_name,
    parameter_value,
    source="slurm",
    restart=False,
    delimiter="=",
):
    """Sets the value of the specified configuration parameter.

    This function modifies the specified slurm configuration file and
    reconfigures slurm (or restarts slurm if restart=True). A backup
    is automatically created and the original configuration is restored
    after the test completes.

    This function may only be used in auto-config mode.

    Args:
        parameter_name (string): The parameter name.
        parameter_value (string): The parameter value.
            Use a value of None to unset a parameter.
        source (string): Name of the config file without the .conf prefix.
        restart (boolean): If True and slurm is running, slurm will be
            restarted rather than reconfigured.
        delimiter (string): The delimiter between the parameter name and the value.

    Note:
        When setting a complex parameter (one which may be repeated and has
        its own subparameters, such as with nodes, partitions and gres),
        the parameter_value should be a dictionary of dictionaries.

    Returns:
        None

    Example:
        >>> set_config_parameter("ClusterName", "cluster1")
        >>> set_config_parameter("required", "/tmp/spank_plugin.so", source="plugstack", delimiter=" ")
    """

    if not properties["auto-config"]:
        require_auto_config("wants to modify parameters")

    if source == "dbd":
        config = "slurmdbd"
    else:
        config = source

    config_file = f"{properties['slurm-config-dir']}/{config}.conf"

    # Remove all matching parameters and append the new parameter
    lines = []
    output = run_command_output(
        f"cat {config_file}", user=properties["slurm-user"], quiet=True
    )
    for line in output.splitlines():
        if not re.search(rf"(?i)^\s*{parameter_name}\s*{re.escape(delimiter)}", line):
            lines.append(line)
    if isinstance(parameter_value, dict):
        for instance_name in parameter_value:
            line = f"{parameter_name}{delimiter}{instance_name}"
            for subparameter_name, subparameter_value in parameter_value[
                instance_name
            ].items():
                line += f" {subparameter_name}{delimiter}{subparameter_value}"
            lines.append(line)
    elif parameter_value is not None:
        lines.append(f"{parameter_name}{delimiter}{parameter_value}")
    input = "\n".join(lines)
    run_command(
        f"cat > {config_file}",
        input=input,
        user=properties["slurm-user"],
        fatal=True,
        quiet=True,
    )

    slurmctld_running = is_slurmctld_running(quiet=True)

    # Remove clustername state file if we aim to change the cluster name
    if parameter_name.lower() == "clustername":
        state_save_location = get_config_parameter(
            "StateSaveLocation", live=slurmctld_running, quiet=True
        )
        run_command(
            f"rm -f {state_save_location}/clustername",
            user=properties["slurm-user"],
            quiet=True,
        )

    # Reconfigure (or restart) slurm controller if it is already running
    if slurmctld_running:
        if source != "slurm" or parameter_name.lower() in [
            "accountingstoragetype",
            "rebootprogram",
        ]:
            restart_slurm(quiet=True)
        elif restart or parameter_name.lower() in [
            "authtype",
            "controlmach",
            "plugindir",
            "statesavelocation",
            "slurmctldhost",
            "slurmctldport",
            "slurmdport",
        ]:
            restart_slurmctld(quiet=True)
        else:
            run_command(
                "scontrol reconfigure", user=properties["slurm-user"], quiet=True
            )


def add_config_parameter_value(name, value, source="slurm"):
    """Appends a value to configuration parameter list.

    When a parameter may contain a comma-separated list of values, this
    function can be used to add a value to the list.

    This function may only be used in auto-config mode.

    Args:
        name (string): The parameter name.
        value (string): The value to add.
        source (string): Name of the config file without the .conf prefix.

    Returns:
        None

    Example:
        >>> add_config_parameter_value('SlurmdParameters', 'config_overrides')
    """

    if config_parameter_includes(name, value, live=False, quiet=True, source=source):
        return

    original_value_string = get_config_parameter(
        name, live=False, quiet=True, source=source
    )
    if original_value_string is None:
        set_config_parameter(name, value, source=source)
    else:
        value_list = original_value_string.split(",")
        value_list.append(value)
        set_config_parameter(name, ",".join(value_list), source=source)


def remove_config_parameter_value(name, value, source="slurm"):
    """Removes a value from a configuration parameter list.

    When a parameter may contain a comma-separated list of values, this
    function can be used to remove a value from the list.

    This function may only be used in auto-config mode.

    Args:
        name (string): The parameter name.
        value (string): The value to remove.
        source (string): Name of the config file without the .conf prefix.

    Returns:
        None

    Example:
        >>> remove_config_parameter_value('SlurmdParameters', 'config_overrides')
    """
    if not config_parameter_includes(
        name, value, live=False, quiet=True, source=source
    ):
        return

    value_list = get_config_parameter(
        name, live=False, quiet=True, source=source
    ).split(",")
    value_list.remove(value.casefold())
    if value_list:
        set_config_parameter(name, ",".join(value_list), source=source)
    else:
        set_config_parameter(name, None, source=source)


def is_tool(tool):
    """Returns True if the tool is found in PATH.

    Args:
        tool (string): The name of the tool to check for in the PATH environment variable.

    Returns:
        True if the tool is found in PATH, False otherwise.

    Example:
        >>> is_tool('ls')
        True
        >>> is_tool('uninstalled-tool-name')
        False
    """

    from shutil import which

    return which(tool) is not None


def require_tool(tool):
    """Skips if the supplied tool is not found.

    Args:
        tool (string): The name of the tool to check for in the PATH environment variable.

    Returns:
        None

    Example:
        >>> require_tool('ls')
        >>> require_tool('uninstalled-tool-name')
    """

    if not is_tool(tool):
        msg = f"This test requires '{tool}' and it was not found"
        pytest.skip(msg, allow_module_level=True)


def require_mpi(mpi_option="pmix", mpi_compiler="mpicc"):
    """Skips if we cannot use the --mpi=mpi_option or the mpi_compiler is not available".

    Args:
        mpi_option (string): The value to use with --mpi when submitting jobs.
        mpi_compiler (string): The required compiler in the system.

    Returns:
        None
    """

    require_tool(mpi_compiler)
    output = run_command_output("srun --mpi=list", fatal=True)
    if re.search(rf"plugin versions available: .*{mpi_option}", output) is None:
        pytest.skip(
            f"This test needs to be able to use --mpi={mpi_option}",
            allow_module_level=True,
        )


def require_influxdb(influx_client="influx", jobacct_gather="jobacct_gather/cgroup"):
    """Require the influx client available and the right config.
    With auto-config the default values would work. Without you may need to setup the acct_gather values in the testsuite.conf file.

    Args:
        influx_client (string): The name of the influxdb client. Default: influx.
        jobacct_gather (string): To use influxdb we need to set some JobAcctGatherType. Default: cgroup.

    Returns:
        None
    """

    require_tool(influx_client)

    require_config_parameter("JobAcctGatherType", jobacct_gather)
    require_config_parameter("AcctGatherProfileType", "acct_gather_profile/influxdb")

    require_config_parameter(
        "ProfileInfluxDBHost",
        f"{properties['influxdb_host']}:{properties['influxdb_port']}",
        source="acct_gather",
    )
    require_config_parameter(
        "ProfileInfluxDBDatabase", properties["influxdb_db"], source="acct_gather"
    )
    require_config_parameter("ProfileInfluxDBDefault", "ALL", source="acct_gather")
    require_config_parameter("ProfileInfluxDBRTPolicy", "autogen", source="acct_gather")

    request_influxdb(f"CREATE DATABASE {properties['influxdb_db']}")
    properties["influxdb-started"] = True


def request_influxdb(query):
    """Send a query using the influx client. Requires to run require_influxdb first.

    Args:
        query (string): The query to send,

    Returns:
        The stdout returned by the influxdb client.
        If the clients fails somehow, this function will fatal.
    """
    return run_command_output(
        f"influx -host {properties['influxdb_host']} -port {properties['influxdb_port']} -database {properties['influxdb_db']} -format column -execute \"{query}\"",
        fatal=True,
    )


def require_whereami():
    """Compiles the whereami.c program to be used by tests.

    This function installs the whereami program.  To get the
    correct output, TaskPlugin is required in the slurm.conf
    file before slurm starts up.
    ex: TaskPlugin=task/cgroup,task/affinity

    The file will be installed in the testsuite/python/lib/scripts
    directory where the whereami.c file is located

    Args:
        None

    Returns:
        None

    Examples:
        >>> atf.require_whereami()
        >>> print('\nwhereami is located at', atf.properties['whereami'])
        >>> output = atf.run_command(f"srun {atf.properties['whereami']}",
        >>>     user=atf.properties['slurm-user'])
    """
    require_config_parameter("TaskPlugin", "task/cgroup,task/affinity")

    # If the file already exists and we don't need to recompile
    dest_file = f"{properties['testsuite_scripts_dir']}/whereami"
    if os.path.isfile(dest_file):
        properties["whereami"] = dest_file
        return

    source_file = f"{properties['testsuite_scripts_dir']}/whereami.c"
    if not os.path.isfile(source_file):
        pytest.skip("Could not find whereami.c!", allow_module_level=True)

    run_command(
        f"gcc {source_file} -o {dest_file}", fatal=True, user=properties["slurm-user"]
    )
    properties["whereami"] = dest_file


def require_config_file(
    filename,
    content,
):
    """
    Ensures that a configuration file exists with the required content.

    In local-config mode, the test is skipped if the required config file
    doesn't exists or is not equivalent.
    In auto-config mode, save current config file if exists and creates the
    required one.

    Note that it cannot be used with some base config files like slurm.conf,
    use require_config_parameter() for those config files instead.

    Supported extensions for local-config are .conf and .yaml.

    Args:
        filename (string): The filename of the config path (e.g. "resources.yaml").
        content (string): The desired content of the config file.

    Returns:
        None
        It will fail or skip the test if required config file cannot be set.
    """

    # Some config files need to be modified, cannot be created from scratch
    if filename in {
        "slurm.conf",
        "slurmdbd.conf",
        "gres.conf",
        "cgroup.conf",
        "testsuite.conf",
    }:
        pytest.fail(f"Use require_config_param() for file {filename}")

    file_path = f"{properties['slurm-config-dir']}/{filename}"

    # If not auto-condig:
    if not properties["auto-config"]:
        supported_ext = {".conf", ".yaml"}
        file, ext = os.path.splitext(filename)
        match ext:
            case ".yaml":
                desired_yaml = yaml.safe_load(content)
                local_yaml = None
                if os.path.exists(file_path):
                    with open(file_path, "r", encoding="utf-8") as f:
                        local_yaml = yaml.safe_load(f)
                if desired_yaml != local_yaml:
                    pytest.skip(f"{filename} is not as required by the test")
                else:
                    logging.debug(f"Found compatible required {filename}")
                    return
            case ".conf":
                # TODO: Ignore blank or commented lines
                with open(file_path, "r", encoding="utf-8") as f:
                    local_content = f.read()
                if content != local_content:
                    pytest.skip(f"{filename} is not as required by the test")
                else:
                    logging.debug(f"Found compatible required {filename}")
                    return
            case _:
                pytest.fail(f"Supported config files are only {supported_ext}")

        pytest.fail(
            "Internal error: function should already skip, fail or return in local-config"
        )

    # In auto-config
    run_command(
        f"cat > {file_path}", input=content, user=properties["slurm-user"], fatal=True
    )


def require_config_parameter(
    parameter_name,
    parameter_value,
    condition=None,
    source="slurm",
    skip_message=None,
    delimiter="=",
):
    """Ensures that a configuration parameter has the required value.

    In local-config mode, the test is skipped if the required configuration is not set.
    In auto-config mode, sets the required configuration value if necessary.

    Args:
        parameter_name (string): The parameter name.
        parameter_value (string): The target parameter value.
        condition (callable): If there is a range of acceptable values, a
            condition can be specified to test whether the current parameter
            value is sufficient. If not, the target parameter_value will be
            used (or the test will be skipped in the case of local-config mode).
        source (string): Name of the config file without the .conf prefix.
        skip_message (string): Message to be displayed if in local-config mode
            and parameter not present.
        delimiter (string): The delimiter between the parameter name and the value.

    Note:
        When requiring a complex parameter (one which may be repeated and has
        its own subparameters, such as with nodes, partitions and gres),
        the parameter_value should be a dictionary of dictionaries. See the
        fourth example for multi-line parameters.

    Returns:
        None

    Examples:
        >>> require_config_parameter('SelectType', 'select/cons_tres')
        >>> require_config_parameter('SlurmdTimeout', 5, lambda v: v <= 5)
        >>> require_config_parameter('Name', {'gpu': {'File': '/dev/tty0'}, 'mps': {'Count': 100}}, source='gres')
        >>> require_config_parameter("PartitionName", {"primary": {"Nodes": "ALL"}, "dynamic1": {"Nodes": "ns1"}, "dynamic2": {"Nodes": "ns2"}, "dynamic3": {"Nodes": "ns1,ns2"}})
    """

    if isinstance(parameter_value, dict):
        tmp1_dict = dict()
        for k1, v1 in parameter_value.items():
            tmp2_dict = dict()
            for k2, v2 in v1.items():
                if isinstance(v2, str):
                    tmp2_dict[k2.casefold()] = v2.casefold()
                else:
                    tmp2_dict[k2.casefold()] = v2
            tmp1_dict[k1.casefold()] = tmp2_dict

        parameter_value = tmp1_dict
    elif isinstance(parameter_value, str):
        parameter_value = parameter_value.casefold()

    observed_value = get_config_parameter(
        parameter_name, live=False, source=source, quiet=True, delimiter=delimiter
    )

    condition_satisfied = False
    if condition is None:
        # condition = lambda observed, desired: observed == desired
        if observed_value == str(parameter_value):
            condition_satisfied = True
    else:
        if condition(observed_value):
            condition_satisfied = True

    if not condition_satisfied:
        if properties["auto-config"]:
            set_config_parameter(
                parameter_name, parameter_value, source=source, delimiter=delimiter
            )
        else:
            if skip_message is None:
                skip_message = f"This test requires the {parameter_name} parameter to be {parameter_value} (but it is {observed_value})"
            pytest.skip(skip_message, allow_module_level=True)


def require_config_parameter_includes(name, value, source="slurm"):
    """Ensures that a configuration parameter list contains the required value.

    In local-config mode, the test is skipped if the configuration parameter
    list does not include the required value.
    In auto-config mode, adds the required value to the configuration parameter
    list if necessary.

    Args:
        name (string): The parameter name.
        value (string): The value we want to be in the list.
        source (string): Name of the config file without the .conf prefix.

    Returns:
        None

    Example:
        >>> require_config_parameter_includes('SlurmdParameters', 'config_overrides')
    """

    if properties["auto-config"]:
        add_config_parameter_value(name, value, source=source)
    else:
        if not config_parameter_includes(
            name, value, source=source, live=False, quiet=True
        ):
            pytest.skip(
                f"This test requires the {name} parameter to include {value}",
                allow_module_level=True,
            )


def require_config_parameter_excludes(name, value, source="slurm"):
    """Ensures that a configuration parameter list does not contain a value.

    In local-config mode, the test is skipped if the configuration parameter
    includes the specified value. In auto-config mode, removes the specified
    value from the configuration parameter list if necessary.

    Args:
        name (string): The parameter name.
        value (string): The value we do not want to be in the list.
        source (string): Name of the config file without the .conf prefix.

    Returns:
        None

    Example:
        >>> require_config_parameter_excludes('SlurmdParameters', 'config_overrides')
    """
    if properties["auto-config"]:
        remove_config_parameter_value(name, value, source=source)
    else:
        if config_parameter_includes(
            name, value, source=source, live=False, quiet=True
        ):
            pytest.skip(
                f"This test requires the {name} parameter to exclude {value}",
                allow_module_level=True,
            )


def require_tty(number):
    """Creates a TTY device file if it does not exist.

    Args:
        number (integer): The number of the TTY device.

    Returns:
        None

    Example:
        >>> require_tty(1)
        >>> require_tty(2)
    """

    tty_file = f"/dev/tty{number}"
    if not os.path.exists(tty_file):
        run_command(
            f"mknod -m 666 {tty_file} c 4 {number}", user="root", fatal=True, quiet=True
        )


## Use this to create an entry in gres.conf and create an associated tty
# def require_gres_device(name):
#
#    gres_value = get_config_parameter('Name', live=False, source='gres', quiet=True)
#    if gres_value is None or name not in gres_value:
#        if not properties['auto-config']:
#            pytest.skip(f"This test requires a '{name}' gres device to be defined in gres.conf", allow_module_level=True)
#        else:
#            require_tty(0)
#            require_config_parameter('Name', {name: {'File': '/dev/tty0'}}, source='gres')


def require_auto_config(reason=""):
    """Ensures that auto-config mode is being used.

    This function skips the test if auto-config mode is not enabled.

    Args:
        reason (string): Augments the skip reason with a context-specific
            explanation for why the auto-config mode is needed by the test.

    Returns:
        None

    Example:
        >>> require_auto_config("wants to set the Epilog")
    """

    if not properties["auto-config"]:
        message = "This test requires auto-config to be enabled"
        if reason != "":
            message += f" ({reason})"
        pytest.skip(message, allow_module_level=True)


def require_accounting(modify=False):
    """Ensures that Slurm accounting is configured.

    In local-config mode, the test is skipped if Slurm accounting is not
    configured. In auto-config mode, configures Slurm accounting if necessary.

    Args:
        modify (boolean): If True, this indicates to the ATF that the test
            will modify the accounting database (e.g. adding accounts, etc).
            A database backup is automatically created and the original dump
            is restored after the test completes.

    Returns:
        None

    Example:
        >>> require_accounting()
        >>> require_accounting(modify=True)
    """

    if properties["auto-config"]:
        if (
            get_config_parameter("AccountingStorageType", live=False, quiet=True)
            != "accounting_storage/slurmdbd"
        ):
            set_config_parameter("AccountingStorageType", "accounting_storage/slurmdbd")
        if modify:
            backup_accounting_database()
    else:
        if modify and not properties["allow-slurmdbd-modify"]:
            require_auto_config("wants to modify the accounting database")
        elif (
            get_config_parameter("AccountingStorageType", live=False, quiet=True)
            != "accounting_storage/slurmdbd"
        ):
            pytest.skip(
                "This test requires accounting to be configured",
                allow_module_level=True,
            )


def require_slurmrestd(openapi_plugins, data_parsers):
    properties["openapi_plugins"] = openapi_plugins
    properties["data_parsers"] = data_parsers

    if properties["auto-config"]:
        properties["slurmrestd-started"] = True
    elif "SLURM_TESTSUITE_SLURMRESTD_URL" in os.environ:
        properties["slurmrestd_url"] = os.environ["SLURM_TESTSUITE_SLURMRESTD_URL"]

        # Setup auth token
        setup_slurmrestd_headers()

        # Check version is the expected one
        if not is_slurmrestd_running():
            pytest.skip(
                f"This test needs slurmrestd running in SLURM_TESTSUITE_SLURMRESTD_URL but cannot connect with {os.environ['SLURM_TESTSUITE_SLURMRESTD_URL']}",
                allow_module_level=True,
            )
    else:
        pytest.skip(
            "This test requires to start slurmrestd or SLURM_TESTSUITE_SLURMRESTD_URL",
            allow_module_level=True,
        )


def start_slurmrestd():
    os.environ["SLURM_JWT"] = "daemon"
    port = None
    attempts = 0

    if "slurm-logs-dir" not in properties:
        properties["slurm-logs-dir"] = os.path.dirname(
            get_config_parameter("SlurmctldLogFile", live=False, quiet=True)
        )

    properties["slurmrestd_log"] = open(
        f"{properties['slurm-logs-dir']}/slurmrestd.log", "w"
    )
    if not properties["slurmrestd_log"]:
        pytest.fail(
            f"Unable to open slurmrestd log: {properties['slurm-logs-dir']}/slurmrestd.log"
        )

    while not port and attempts < 15:
        if "slurmrestd_port" not in properties or attempts > 0:
            port = get_open_port()
        else:
            port = properties["slurmrestd_port"]
        attempts += 1
        args = [
            "slurmrestd",
            "-a",
            "jwt",
            "-s",
            properties["openapi_plugins"],
        ]
        if properties["data_parsers"] is not None:
            args.extend(["-d", properties["data_parsers"]])

        args.append(f"localhost:{port}")
        logging.debug(f"Trying to start slurmrestd: {args}")

        properties["slurmrestd"] = subprocess.Popen(
            args,
            stdin=subprocess.DEVNULL,
            stdout=properties["slurmrestd_log"],
            stderr=properties["slurmrestd_log"],
        )
        s = None

        for i in range(100):
            if properties["slurmrestd"].poll():
                break

            try:
                s = socket.create_connection(("localhost", port))
                break
            except Exception as e:
                logging.debug(f"Unable to connect to port {port}: {e}")
            time.sleep(1)

        if s:
            s.close()
            break

        logging.debug(f"slurmrestd accepting on port {port} but is still running")
        properties["slurmrestd"].kill()
        properties["slurmrestd"].wait()
        port = None

    if not port:
        pytest.fail(f"Unable start slurmrestd after trying {attempts} different ports")

    del os.environ["SLURM_JWT"]

    properties["slurmrestd_port"] = port
    properties["slurmrestd_url"] = f"http://localhost:{port}/"

    # Setup auth token
    setup_slurmrestd_headers()

    # Check slurmrestd is up
    if not is_slurmrestd_running():
        pytest.fail("Slurmrestd not responding")


def setup_slurmrestd_headers():
    # Create the headers with the token to connect later
    token = (
        run_command_output("scontrol token lifespan=600", fatal=True)
        .replace("SLURM_JWT=", "")
        .replace("\n", "")
    )
    if token == "":
        logging.warning("unable to get auth/jwt token")

    properties["slurmrestd-headers"] = {
        "X-SLURM-USER-NAME": get_user_name(),
        "X-SLURM-USER-TOKEN": token,
    }


def get_user_name():
    """Returns the username of the current user.

    Args:
        None

    Returns:
        The username of the current user.

    Example:
        >>> get_user_name()
        'john_doe'
    """
    return pwd.getpwuid(os.getuid()).pw_name


def cancel_jobs(
    job_list,
    timeout=default_polling_timeout,
    poll_interval=0.1,
    fatal=False,
    quiet=False,
):
    """Cancels a list of jobs and waits for them to complete.

    Args:
        job_list (list): A list of job ids to cancel. All 0s will be ignored.
        timeout (integer): Number of seconds to wait for jobs to be done before
            timing out.
        poll_interval (float): Number of seconds to wait between job state
            polls.
        fatal (boolean): If True, a timeout will result in the test failing.
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        True if all jobs were successfully cancelled and completed within
        the timeout period, False otherwise.

    Example:
        >>> cancel_jobs([1234, 5678], timeout=60, fatal=True)
        True
        >>> cancel_jobs([9876, 5432], timeout=30, fatal=False)
        False
    """

    # Filter list to ignore job_ids being 0
    job_list = [i for i in job_list if i != 0]
    job_list_string = " ".join(str(i) for i in job_list)

    if job_list_string == "":
        return True

    run_command(f"scancel {job_list_string}", fatal=fatal, quiet=quiet)

    for job_id in job_list:
        status = wait_for_job_state(
            job_id,
            "DONE",
            timeout=timeout,
            poll_interval=poll_interval,
            fatal=fatal,
            quiet=quiet,
        )
        if not status:
            if fatal:
                pytest.fail(
                    f"Job ({job_id}) was not cancelled within the {timeout} second timeout"
                )
            return status

    return True


def cancel_all_jobs(
    timeout=default_polling_timeout, poll_interval=0.1, fatal=False, quiet=False
):
    """Cancels all jobs by the test user and waits for them to be cancelled.

    Args:
        fatal (boolean): If True, a timeout will result in the test failing.
        timeout (integer): If timeout number of seconds expires before the
            jobs are verified to be cancelled, fail.
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        True if all jobs were successfully cancelled and completed within
        the timeout period, False otherwise.

    Example:
        >>> cancel_all_jobs(timeout=60, fatal=True)
        True
        >>> cancel_all_jobs(timeout=30, fatal=False)
        False
    """

    user_name = properties["test-user"]

    run_command(f"scancel -u {user_name}", fatal=fatal, quiet=quiet)

    return repeat_command_until(
        f"squeue -u {user_name} --noheader",
        lambda results: results["stdout"] == "",
        timeout=timeout,
        poll_interval=poll_interval,
        fatal=fatal,
        quiet=quiet,
    )


def is_integer(value):
    try:
        int(value)
    except ValueError:
        return False
    else:
        return True


def is_float(value):
    try:
        float(value)
    except ValueError:
        return False
    else:
        return True


def get_nodes(live=True, quiet=False, **run_command_kwargs):
    """Returns the node configuration as a dictionary of dictionaries.

    If the live argument is not True, the dictionary will contain the literal
    configuration values and the DEFAULT node will be separately indexed.

    Args:
        live (boolean):
            If True, the node configuration information is obtained via a query
            to the slurmctld daemon (e.g. scontrol show config).
            If False, the node configuration information is obtained by
            directly parsing the Slurm configuration file (slurm.conf).
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        A dictionary of dictionaries where the first level keys are the node
        names and with their values being a dictionary of configuration
        parameters for the respective node.

    Example:
        >>> get_nodes()
        {'node1': {'NodeName': 'node1', 'Arch': 'x86_64', 'CoresPerSocket': 1, 'CPUAlloc': 0, 'CPUEfctv': 1, 'CPUTot': 1, 'CPULoad': 91.4, 'AvailableFeatures': None, 'ActiveFeatures': None, 'Gres': None, 'NodeAddr': 'slurm-host', 'NodeHostName': 'slurm-host', 'Port': 6821, 'Version': '24.05.0-0rc1', ...}}
        >>> get_nodes(live=False, quiet=True)
        {'node1': {'NodeHostname': 'slurm-host', 'Port': 6821, 'NodeName': 'node1'}}
    """

    nodes_dict = {}

    if live:
        # TODO: Remove extra debug info for t22858 instead of fatal
        result = run_command(
            "scontrol show nodes --json -F",
            fatal=False,
            quiet=quiet,
            **run_command_kwargs,
        )
        if result["exit_code"]:
            logging.debug(
                "Fatal failure of 'scontrol show nodes', probably due 'Unable to contact slurm controller' (t22858)"
            )
            logging.debug("Getting gcore from slurmctld before fatal.")
            gcore("slurmctld")
            pytest.fail(
                f"Command 'scontrol show nodes -oF' failed with rc={result['exit_code']}: {result['stderr']}"
            )
        output = result["stdout"]
        node_dict_json = json.loads(output)["nodes"]

        node_dict = {}
        for node in node_dict_json:
            # Add the node dictionary to the nodes dictionary
            nodes_dict[node["name"]] = node
    else:
        # Get the config dictionary
        config_dict = get_config(live=False, quiet=quiet)

        # Convert keys to lower case so we can do a case-insensitive search
        lower_config_dict = dict(
            (key.lower(), value) for key, value in config_dict.items()
        )

        # DEFAULT will be included separately
        if "nodename" in lower_config_dict:
            for node_expression, node_expression_dict in lower_config_dict[
                "nodename"
            ].items():
                port_expression = (
                    node_expression_dict["Port"]
                    if "Port" in node_expression_dict
                    else ""
                )

                # Break up the node expression and port expression into lists
                node_list = node_range_to_list(node_expression)
                port_list = range_to_list(port_expression)

                # Iterate over the nodes in the expression
                for node_index in range(len(node_list)):
                    node_name = node_list[node_index]
                    # Add the parameters to the temporary node dictionary
                    node_dict = dict(node_expression_dict)
                    node_dict["NodeName"] = node_name
                    if node_index < len(port_list):
                        node_dict["Port"] = int(port_list[node_index])
                    # Add the node dictionary to the nodes dictionary
                    nodes_dict[node_name] = node_dict

    return nodes_dict


def get_node_parameter(node_name, parameter_name, default=None, live=True):
    """Obtains the value for a node configuration parameter.

    Args:
        node_name (string): The node name.
        parameter_name (string): The parameter name.
        default (string or None): This value is returned if the parameter
            is not found.
        live (boolean):
            If True, the node configuration information is obtained via a query
            to the slurmctld daemon (e.g. scontrol show config).
            If False, the node configuration information is obtained by
            directly parsing the Slurm configuration file (slurm.conf).

    Returns:
        The value of the specified node parameter, or the default if not found.

    Example:
        >>> get_node_parameter('node1', 'State')
        'IDLE'
        >>> get_node_parameter('node2', 'RealMemory', default=4)
        1
        >>> get_node_parameter('node3', 'Partitions', default='primary', live=False)
        'primary'
    """

    nodes_dict = get_nodes(live=live, quiet=True)

    if node_name in nodes_dict:
        node_dict = nodes_dict[node_name]
    else:
        pytest.fail(f"Node ({node_name}) was not found in the node configuration")

    if parameter_name in node_dict:
        return node_dict[parameter_name]
    else:
        return default


def set_node_parameter(node_name, new_parameter_name, new_parameter_value):
    """Sets the value of the specified node configuration parameter.

    This function sets a node property for the specified node and restarts
    the relevant slurm daemons. A backup is automatically created and the
    original configuration is restored after the test completes.

    This function may only be used in auto-config mode.

    Args:
        node_name (string): The node name.
        new_parameter_name (string): The parameter name.
        new_parameter_value (string): The parameter value.
            Use a value of None to unset a node parameter.

    Returns:
        None

    Example:
        >>> set_node_parameter('node1', 'Features', 'f1')
    """

    if not properties["auto-config"]:
        require_auto_config("wants to modify node parameters")

    config_file = f"{properties['slurm-config-dir']}/slurm.conf"

    # Read the original slurm.conf into a list of lines
    output = run_command_output(
        f"cat {config_file}", user=properties["slurm-user"], quiet=True
    )
    original_config_lines = output.splitlines()
    new_config_lines = original_config_lines.copy()

    # Locate the node among the various NodeName definitions
    found_node_name = False
    for line_index in range(len(original_config_lines)):
        line = original_config_lines[line_index]

        words = re.split(r" +", line.strip())
        if len(words) < 1:
            continue
        if words[0][0] == "#":
            continue
        parameter_name, parameter_value = words[0].split("=", 1)
        if parameter_name.lower() != "nodename":
            continue

        # We found a NodeName line. Read in the node parameters
        node_expression = parameter_value
        port_expression = ""
        original_node_parameters = collections.OrderedDict()
        for word in words[1:]:
            parameter_name, parameter_value = word.split("=", 1)
            if parameter_name.lower() == "port":
                port_expression = parameter_value
            else:
                original_node_parameters[parameter_name] = parameter_value

        node_list = node_range_to_list(node_expression)
        port_list = range_to_list(port_expression)

        # Determine whether our node is included in this node expression
        if node_name in node_list:
            # Yes. We found the node
            found_node_name = True
            node_index = node_list.index(node_name)

            # Delete the original node definition
            new_config_lines.pop(line_index)

            # Add the modified definition for the specified node
            modified_node_parameters = original_node_parameters.copy()
            if new_parameter_value is None:
                if new_parameter_name in modified_node_parameters:
                    del modified_node_parameters[new_parameter_name]
            else:
                modified_node_parameters[new_parameter_name] = new_parameter_value
            modified_node_line = f"NodeName={node_name}"
            if node_index < len(port_list):
                modified_node_line += f" Port={port_list[node_index]}"
            for parameter_name, parameter_value in modified_node_parameters.items():
                modified_node_line += f" {parameter_name}={parameter_value}"
            new_config_lines.insert(line_index, modified_node_line)

            # If the node was part of an aggregate node definition
            node_list.remove(node_name)
            if len(node_list):
                # Write the remainder of the aggregate using the original attributes
                remainder_node_expression = node_list_to_range(node_list)
                remainder_node_line = f"NodeName={remainder_node_expression}"
                if node_index < len(port_list):
                    port_list.pop(node_index)
                if port_list:
                    remainder_port_expression = list_to_range(port_list)
                    remainder_node_line += f" Port={remainder_port_expression}"
                for parameter_name, parameter_value in original_node_parameters.items():
                    remainder_node_line += f" {parameter_name}={parameter_value}"
                new_config_lines.insert(line_index, remainder_node_line)

            break

    if not found_node_name:
        pytest.fail(
            f"Invalid node specified in set_node_parameter(). Node ({node_name}) does not exist"
        )

    # Write the config file back out with the modifications
    new_config_string = "\n".join(new_config_lines)
    run_command(
        f"echo '{new_config_string}' > {config_file}",
        user=properties["slurm-user"],
        fatal=True,
        quiet=True,
    )

    # Restart slurm controller if it is already running
    if is_slurmctld_running(quiet=True):
        restart_slurm(quiet=True)


def get_reservations(quiet=False, **run_command_kwargs):
    """Returns the reservations as a dictionary of dictionaries.

    Args:
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        A dictionary of dictionaries where the first level keys are the
        reservation names and with their values being a dictionary of
        configuration parameters for the respective reservation.

    Example:
        >>> get_reservations()
        {'resv1': {'ReservationName': 'resv1', 'StartTime': '2024-04-06T00:00:17', 'EndTime': '2024-04-06T00:15:17', 'Duration': '00:15:00', 'Nodes': 'node1', 'NodeCnt': 1, 'CoreCnt': 1, 'Features': None, 'PartitionName': 'primary', 'Flags': 'DAILY,REPLACE_DOWN,PURGE_COMP=00:01:00', 'TRES': 'cpu=1', 'Users': 'atf', 'Groups': None, 'Accounts': None, 'Licenses': None, 'State': 'INACTIVE', 'BurstBuffer': None, 'MaxStartDelay': None}}
    """

    resvs_dict = {}
    resv_dict = {}

    output = run_command_output(
        "scontrol show reservations -o", fatal=True, quiet=quiet, **run_command_kwargs
    )
    for line in output.splitlines():
        if line == "":
            continue

        while match := re.search(r"^ *([^ =]+)=(.*?)(?= +[^ =]+=| *$)", line):
            parameter_name, parameter_value = match.group(1), match.group(2)

            # Remove the consumed parameter from the line
            line = re.sub(r"^ *([^ =]+)=(.*?)(?= +[^ =]+=| *$)", "", line)

            # Reformat the value if necessary
            if is_integer(parameter_value):
                parameter_value = int(parameter_value)
            elif is_float(parameter_value):
                parameter_value = float(parameter_value)
            elif parameter_value == "(null)":
                parameter_value = None

            # Add it to the temporary resv dictionary
            resv_dict[parameter_name] = parameter_value

        # Add the resv dictionary to the resvs dictionary
        resvs_dict[resv_dict["ReservationName"]] = resv_dict

        # Clear the resv dictionary for use by the next resv
        resv_dict = {}

    return resvs_dict


def get_reservation_parameter(resv_name, parameter_name, default=None):
    """Obtains the value for a reservation configuration parameter.

    Args:
        resv_name (string): The reservation name.
        parameter_name (string): The parameter name.
        default (string or None): This value is returned if the parameter
            is not found.

    Returns:
        The value of the specified reservation parameter, or the default if not found.

    Example:
        >>> get_reservation_parameter('resv1', 'PartitionName')
        'primary'
        >>> get_reservation_parameter('resv2', 'EndTime', default='2024-04-06T00:15:17')
        '2024-04-06T00:15:17'
    """

    resvs_dict = get_reservations()

    if resv_name in resvs_dict:
        resv_dict = resvs_dict[resv_name]
    else:
        pytest.fail(f"reservation ({resv_name}) was not found")

    if parameter_name in resv_dict:
        return resv_dict[parameter_name]
    else:
        return default


def is_super_user():
    """Checks if the current user is a super user.

    Args:
        None

    Returns:
        True if the current user is a super user, False otherwise.

    Example:
        >>> is_super_user()
        False
    """

    uid = os.getuid()

    if uid == 0:
        return True

    user = pwd.getpwuid(uid)[0]
    if get_config_parameter("SlurmUser") == user:
        return True

    return False


def require_sudo_rights():
    """Skips the test if the test user does not have unprompted sudo privileges.

    Args:
        None

    Returns:
        None

    Example:
        >>> require_sudo_rights()
    """
    if not properties["sudo-rights"]:
        pytest.skip(
            "This test requires the test user to have unprompted sudo privileges",
            allow_module_level=True,
        )


def submit_job_sbatch(sbatch_args='--wrap "sleep 60"', **run_command_kwargs):
    """Submits a job using sbatch and returns the job id.

    The submitted job will automatically be cancelled when the test ends.

    Args:
        sbatch_args (string): The arguments to sbatch.

    Returns:
        The job id.

    Note:
        The background parameter is ignored and always treated as False, since
        sbatch is already non-blocking and the function needs to capture the
        job ID from output.

    Example:
        >>> submit_job_sbatch('--wrap "echo Hello"')
        1234
        >>> submit_job_sbatch('-J myjob --wrap "echo World"', fatal=True)
        5678
    """

    # Since sbatch is non-blocking anyway, make sure we're not in the background
    run_command_kwargs["background"] = False

    output = run_command_output(f"sbatch {sbatch_args}", **run_command_kwargs)

    if match := re.search(r"Submitted \S+ job (\d+)", output):
        job_id = int(match.group(1))
        properties["submitted-jobs"].append(job_id)
        return job_id
    else:
        return 0


# Returns results
def run_job(srun_args, **run_command_kwargs):
    """Runs a job using srun and returns the run_command results dictionary.

    If the srun command times out, it will automatically be cancelled when the
    test ends.

    Args:
        srun_args (string): The arguments to srun.

    Returns:
        The srun run_command results dictionary.

    Example:
        >>> run_job('-n 4 --output=output.txt ./my_executable', timeout=60)
        {'command': 'srun -n 4 --output=output.txt ./my_executable', 'start_time': 1712276708.827, 'duration': 60.0, 'exit_code': 110, 'stdout': '', 'stderr': 'srun: Requested partition configuration not available now\nsrun: job 15 queued and waiting for resources\n'}
        >>> run_job('-n 1 --error=error.txt ./my_executable', fatal=True)
        {'command': 'srun -n 1 --error=error.txt my_executable', 'start_time': 1712277798.016, 'duration': 0.06, 'exit_code': 0, 'stdout': 'foo bar\n', 'stderr': ''}
    """

    results = run_command(f"srun {srun_args}", **run_command_kwargs)

    return results


# Return exit code
def run_job_exit(srun_args, **run_command_kwargs):
    """Runs a job using srun and returns the exit code.

    If the srun command times out, it will automatically be cancelled when the
    test ends.

    Args:
        srun_args (string): The arguments to srun.

    Returns:
        The exit code from srun.

    Example:
        >>> run_job_exit('-n 4 --output=output.txt ./my_executable', timeout=60)
        2
        >>> run_job_exit('-n 2 --error=error.txt ./my_executable', fatal=True)
        0
    """

    results = run_job(srun_args, **run_command_kwargs)

    return results["exit_code"]


# Return output
def run_job_output(srun_args, **run_command_kwargs):
    """Runs a job using srun and returns the standard output.

    If the srun command times out, it will automatically be cancelled when the
    test ends.

    Args:
        srun_args (string): The arguments to srun.

    Returns:
        The standard output from srun.

    Example:
        >>> run_job_output('-n 4 ./my_executable', timeout=60)
        'stdout of the command'
        >>> run_job_output('-n 2 --output=output.txt ./my_executable', fatal=True)
        ''
    """

    results = run_job(srun_args, **run_command_kwargs)

    return results["stdout"]


# Return error
def run_job_error(srun_args, **run_command_kwargs):
    """Runs a job using srun and returns the standard error.

    If the srun command times out, it will automatically be cancelled when the
    test ends.

    Args:
        srun_args (string): The arguments to srun.

    Returns:
        The standard error from srun.

    Example:
        >>> run_job_error('-n 4 --output=output.txt ./my_executable', timeout=60)
        'stderr of the command'
        >>> run_job_error('-n 200000 ./my_executable', fatal=True) # Will automatically fail the test due to resources
    """

    results = run_job(srun_args, **run_command_kwargs)

    return results["stderr"]


# Return job id
def submit_job_srun(srun_args, **run_command_kwargs):
    """Runs a job using srun and returns the job id.

    This function obtains the job id by adding the -v option to srun and parsing
    out the job id. If the srun command times out, it will automatically be
    cancelled when the test ends.

    Args:
        srun_args (string): The arguments to srun.

    Returns:
        The job id (int) or 0 if failed.

    Example:
        >>> submit_job_srun('-n 4 --output=output.txt ./my_executable')
        12345
        >>> submit_job_srun('-n 2 --output=output.txt ./my_executable', timeout=60)
        67890
        >>> submit_job_srun('-n 2 --output=output.txt ./my_executable', background=True)
        34567
    """
    verbose_arg = "-v"

    if run_command_kwargs.get("background", False):
        # Use timeout for finding job id
        timeout = run_command_kwargs.pop("timeout", 30)
        # Extract fatal/xfail to handle at job submission level
        fatal = run_command_kwargs.pop("fatal", False)
        xfail = run_command_kwargs.pop("xfail", False)

        if get_version("bin/srun") < (25, 11):
            # Older versions may not print the jobid with -v until job is scheduled
            verbose_arg = "--begin=now+1"
            logging.warning(
                "Using {verbose_arg} to be able to extract jobid when using srun background"
            )

    results = run_command(f"srun {verbose_arg} {srun_args}", **run_command_kwargs)

    if run_command_kwargs.get("background", False):
        job_id = None
        start_time = time.time()

        while time.time() - start_time < timeout:
            line = _readline_with_timeout(results["process"].stderr, timeout=5)
            if line is None:
                # Timeout or EOF
                if results["process"].poll() is not None:
                    break
                continue

            if match := re.search(r"job(id)?[ =](\d+)", line, re.IGNORECASE):
                job_id = int(match.group(2))
                break

            if results["process"].poll() is not None:
                break
    else:
        if match := re.search(r"job(id)?[ =](\d+)", results["stderr"], re.IGNORECASE):
            job_id = int(match.group(2))
        else:
            job_id = None

    if job_id is not None:
        properties["submitted-jobs"].append(job_id)

    # Handle fatal/xfail for background jobs
    if run_command_kwargs.get("background", False):
        if fatal:
            if xfail and job_id is not None:
                pytest.fail(
                    f'Command "srun -v {srun_args}" was expected to fail but got job id {job_id}'
                )
            elif not xfail and job_id is None:
                pytest.fail(f'Command "srun -v {srun_args}" failed to get job id')

    return job_id if job_id else 0


def _readline_with_timeout(stream, timeout=default_command_timeout):
    """Read a line from a stream with a timeout.

    Args:
        stream: The stream to read from (e.g., process.stderr)
        timeout: Maximum seconds to wait for a line

    Returns:
        The line read from the stream, or None if timeout/EOF
    """
    result = [None]

    def read_line():
        try:
            result[0] = stream.readline()
        except Exception:
            pass

    thread = threading.Thread(target=read_line)
    thread.daemon = True
    thread.start()
    thread.join(timeout)

    if thread.is_alive():
        # Timeout occurred
        return None
    return result[0]


# Return job id (command should not be interactive/shell)
def submit_job_salloc(salloc_args, **run_command_kwargs):
    """Submits a job using salloc and returns the job id.

    The submitted job will automatically be cancelled when the test ends.

    Args:
        salloc_args (string): The arguments to salloc.

    Returns:
        The job id (int) or 0 if failed.

    Example:
        >>> submit_job_salloc('-N 1 -t 60 --output=output.txt ./my_executable')
        12345
        >>> submit_job_salloc('-N 2 -t 120 --output=output.txt ./my_executable', timeout=60)
        67890
        >>> submit_job_salloc('-N 2 -t 120 --output=output.txt ./my_executable', background=True)
        34567
    """
    # If background=True, use timeout for finding job id
    if run_command_kwargs.get("background", False):
        timeout = run_command_kwargs.pop("timeout", 30)
        # Extract fatal/xfail to handle at job submission level
        fatal = run_command_kwargs.pop("fatal", False)
        xfail = run_command_kwargs.pop("xfail", False)

    results = run_command(f"salloc {salloc_args}", **run_command_kwargs)

    if run_command_kwargs.get("background", False):
        job_id = None
        start_time = time.time()

        while time.time() - start_time < timeout:
            line = _readline_with_timeout(results["process"].stderr, timeout=5)
            if line is None:
                # Timeout or EOF
                if results["process"].poll() is not None:
                    break
                continue

            if match := re.search(r"job allocation (\d+)", line):
                job_id = int(match.group(1))
                break

            if results["process"].poll() is not None:
                break
    else:
        if match := re.search(r"Granted job allocation (\d+)", results["stderr"]):
            job_id = int(match.group(1))
        else:
            job_id = None

    if job_id is not None:
        properties["submitted-jobs"].append(job_id)

    # Handle fatal/xfail for background jobs
    if run_command_kwargs.get("background", False):
        if fatal:
            if xfail and job_id is not None:
                pytest.fail(
                    f'Command "salloc {salloc_args}" was expected to fail but got job id {job_id}'
                )
            elif not xfail and job_id is None:
                pytest.fail(f'Command "salloc {salloc_args}" failed to get job id')

    return job_id if job_id else 0


# Return job id
def submit_job(command, job_param, job, *, wrap_job=True, **run_command_kwargs):
    """Submits a job using the given command and returns the job id.

    Args:
        command (string): The command to submit the job (salloc, srun, sbatch).
        job_param (string): The arguments to the job.
        job (string): The command or job file to be executed.
        wrap_job (boolean): If True, the job will be wrapped when the command is 'sbatch'.

    Returns:
        The job id.

    Example:
        >>> submit_job('salloc', '-N 1 -t 60', './my_executable', quiet=True)
        12345
        >>> submit_job('srun', '-N 2 -t 120', './my_executable', quiet=True)
        67890
        >>> submit_job('sbatch', '-N 1 -t 60', './my_script.sh', wrap_job=False, quiet=True)
        23456
    """

    # Make sure command is a legal command to run a job
    assert command in [
        "salloc",
        "srun",
        "sbatch",
    ], f"Invalid command '{command}'. Should be salloc, srun, or sbatch."

    if command == "salloc":
        return submit_job_salloc(f"{job_param} {job}", **run_command_kwargs)
    elif command == "srun":
        return submit_job_srun(f"{job_param} {job}", **run_command_kwargs)
    elif command == "sbatch":
        # If the job should be wrapped, do so before submitting
        if wrap_job:
            job = f"--wrap '{job}'"
        return submit_job_sbatch(f"{job_param} {job}", **run_command_kwargs)


def run_job_nodes(srun_args, **run_command_kwargs):
    """Runs a job using srun and returns the allocated node list.

    This function obtains the job id by adding the -v option to srun and parsing
    out the allocated node list. If the srun command times out, it will
    automatically be cancelled when the test ends.

    Args:
        srun_args (string): The arguments to srun.

    Returns:
        The allocated node list for the job.

    Example:
        >>> run_job_nodes('-N 2 hostname', quiet=True)
        ['node001', 'node002']
        >>> run_job_nodes('-N 1 --exclude=node001 hostname', quiet=True)
        ['node002']
    """

    results = run_command(f"srun -v {srun_args}", **run_command_kwargs)
    node_list = []
    if results["exit_code"] == 0:
        if match := re.search(r"jobid \d+: nodes\(\d+\):`([^']+)'", results["stderr"]):
            node_list = node_range_to_list(match.group(1))
    return node_list


def get_qos(name=None, **run_command_kwargs):
    """Returns the QOSes in the system as a dictionary of dictionaries.

    Args:
        name: The name of a specific QOS of which to get parameters.

    Returns:
        A dictionary of dictionaries where the first level keys are the QOSes names
        and with their values being a dictionary of configuration parameters for
        the respective QOS.
    """

    command = "sacctmgr --json show qos"
    if id is not None:
        command += f" {name}"
    output = run_command_output(command, fatal=True, quiet=True, **run_command_kwargs)

    qos_dict = {}
    qos_list = json.loads(output)["qos"]
    for qos in qos_list:
        qos_dict[qos["name"]] = qos

    return qos_dict


def get_jobs(job_id=None, dbd=False, **run_command_kwargs):
    """Returns the jobs in the system as a dictionary of dictionaries.

    Args:
        job_id (integer): The id of a specific job of which to get parameters.
        dbd (boolean): If True, obtain the jobs from sacct instead of scontrol.

    Returns:
        A dictionary of dictionaries where the first level keys are the job ids
        and with their values being a dictionary of configuration parameters for
        the respective job.

    Example:
        >>> get_jobs()
        {38: {'JobId': 38, 'JobName': 'wrap', 'UserId': 'atf(1002)', 'GroupId': 'atf(1002)', ...},
         39: {'JobId': 39, 'JobName': 'wrap', 'UserId': 'atf(1002)', 'GroupId': 'atf(1002)', ...}}
        >>> get_jobs(job_id='12345', quiet=True)
        {12345: {'JobId': '12345', 'JobName': 'foo.sh', 'UserId': 'test(1003)', ...}}
    """

    jobs_dict = {}

    if dbd:
        command = "sacct --json -X"
        if job_id is not None:
            command += f" -j {job_id}"
        output = run_command_output(
            command, fatal=True, quiet=True, **run_command_kwargs
        )

        jobs_list = json.loads(output)["jobs"]
        for job in jobs_list:
            jobs_dict[job["job_id"]] = job

    else:
        command = "scontrol -d -o show jobs"
        if job_id is not None:
            command += f" {job_id}"
        # TODO: Remove extra debug info for t22858 instead of fatal
        result = run_command(command, fatal=False, **run_command_kwargs)
        if result["exit_code"]:
            logging.debug(
                f"Fatal failure of {command}, probably due 'Unable to contact slurm controller' (t22858)"
            )
            logging.debug("Getting gcore from slurmctld before fatal.")
            gcore("slurmctld")
            pytest.fail(
                f"Command {command} failed with rc={result['exit_code']}: {result['stderr']}"
            )

        output = result["stdout"]
        job_dict = {}
        for line in output.splitlines():
            if line == "":
                continue

            while match := re.search(r"^ *([^ =]+)=(.*?)(?= +[^ =]+=| *$)", line):
                param_name, param_value = match.group(1), match.group(2)

                # Remove the consumed parameter from the line
                line = re.sub(r"^ *([^ =]+)=(.*?)(?= +[^ =]+=| *$)", "", line)

                # Reformat the value if necessary
                if is_integer(param_value):
                    param_value = int(param_value)
                elif is_float(param_value):
                    param_value = float(param_value)
                elif param_value == "(null)":
                    param_value = None

                # Add it to the temporary job dictionary
                job_dict[param_name] = param_value

            # Add the job dictionary to the jobs dictionary
            if job_dict:
                jobs_dict[job_dict["JobId"]] = job_dict

                # Clear the job dictionary for use by the next job
                job_dict = {}

    return jobs_dict


def get_steps(job_id, **run_command_kwargs):
    """Returns the steps of a given job_id as a dictionary of dictionaries.

    Args:
        job_id (string): The specific JobId to retrieve information for.

    Returns:
        A dictionary of dictionaries where the first level keys are the step ids
        and with their values being a dictionary of configuration parameters for
        the respective step.

    Example:
        >>> get_steps(44)
        {'44.batch': {'StepId': '44.batch', 'UserId': 1002, 'StartTime': '2024-04-05T01:17:53', ...},
         '44.0': {'StepId': 44.0, 'UserId': 1002, 'StartTime': '2024-04-05T01:17:54', ...}}
    """

    steps_dict = {}
    step_dict = {}

    command = f"scontrol -d -o show steps {job_id}"
    result = run_command(command, **run_command_kwargs)

    if result["exit_code"]:
        logging.debug("scontrol command failed, no steps returned")
        return step_dict

    output = result["stdout"]
    for line in output.splitlines():
        if line == "":
            continue

        while match := re.search(r"^ *([^ =]+)=(.*?)(?= +[^ =]+=| *$)", line):
            param_name, param_value = match.group(1), match.group(2)

            # Remove the consumed parameter from the line
            line = re.sub(r"^ *([^ =]+)=(.*?)(?= +[^ =]+=| *$)", "", line)

            # Reformat the value if necessary
            if is_integer(param_value):
                param_value = int(param_value)
            elif is_float(param_value) and param_name != "StepId":
                param_value = float(param_value)
            elif param_value == "(null)":
                param_value = None

            # Add it to the temporary step dictionary
            step_dict[param_name] = param_value

        # Add the step dictionary to the steps dictionary
        if step_dict:
            steps_dict[str(step_dict["StepId"])] = step_dict

            # Clear the step dictionary for use by the next step
            step_dict = {}

    return steps_dict


def get_job(job_id, quiet=False):
    """Returns job information for a specific job as a dictionary.

    Args:
        job_id (integer): The id of the job for which information is requested.
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        A dictionary containing parameters for the specified job. If the job id
        is not found, an empty dictionary is returned.

    Example:
        >>> get_job(51)
        {'JobId': 51, 'JobName': 'wrap', 'UserId': 'atf(1002)', 'GroupId': 'atf(1002)', ...}
        >>> get_job(182, quiet=True)
        {'JobId': 182, 'JobName': 'foo.sh', 'UserId': 'atf(1002)', 'GroupId': 'atf(1002)', ...}
    """

    jobs_dict = get_jobs(job_id, quiet=quiet)

    return jobs_dict[job_id] if job_id in jobs_dict else {}


def get_job_parameter(job_id, parameter_name, default=None, quiet=False):
    """Returns the value of a specific parameter for a given job.

    Args:
        job_id (integer): The id of the job for which the parameter value is requested.
        parameter_name (string): The name of the parameter whose value is to be obtained.
        default (string or None): The value to be returned if the parameter is not found.
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        The value of the specified job parameter, or the default value if the
        parameter is not found.

    Example:
        >>> get_job_parameter(12345, 'UserId')
        'atf(1002)'
        >>> get_job_parameter(67890, 'Partition', default='normal', quiet=True)
        'primary'
    """

    jobs_dict = get_jobs(quiet=quiet)

    if job_id in jobs_dict:
        job_dict = jobs_dict[job_id]
    else:
        pytest.fail(f"Job ({job_id}) was not found in the job configuration")

    if parameter_name in job_dict:
        return job_dict[parameter_name]
    else:
        return default


def get_job_id_from_array_task(array_job_id, array_task_id, fatal=False, quiet=True):
    """Returns the raw job id of a task of a job array.

    Args:
        array_job_id (integer): The id of the job array.
        array_task_id (integer): The id of the task of the job array.
        fatal (boolean): If True, fails if the raw job id is not found in the system.
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        The raw job id of the given task of a job array, or 0 if not found.

    Example:
        >>> get_job_id_from_array_task(234, 2)
        241
    """

    jobs_dict = get_jobs(quiet=quiet)
    for job_id, job_values in jobs_dict.items():
        if (
            job_values["ArrayJobId"] == array_job_id
            and job_values["ArrayTaskId"] == array_task_id
        ):
            return job_id

    if fatal:
        pytest.fail(f"{array_job_id}_{array_task_id} was not found in the system")

    return 0


def get_step_parameter(step_id, parameter_name, default=None, quiet=False):
    """Returns the value of a specific parameter for a given step.

    Args:
        step_id (string): The id of the step for which the parameter value is requested.
        parameter_name (string): The name of the parameter whose value is to be obtained.
        default (string or None): The value to be returned if the parameter is not found.
        quiet (boolean): If True, logging is performed at the TRACE log level.

    Returns:
        The value of the specified step parameter, or the default value if the
        parameter is not found.

    Example:
        >>> get_step_parameter("45.0", 'NodeList')
        'node[2,4]'
        >>> get_step_parameter("60.1", 'Partition', default='normal', quiet=True)
        'primary'
    """

    steps_dict = get_steps(step_id, quiet=quiet)

    if step_id not in steps_dict:
        logging.debug(f"Step ({step_id}) was not found in the step list")
        return default

    step_dict = steps_dict[step_id]
    if parameter_name in step_dict:
        return step_dict[parameter_name]
    else:
        return default


def wait_for_node_state_any(
    nodename,
    desired_node_states,
    timeout=default_polling_timeout,
    poll_interval=None,
    fatal=False,
    reverse=False,
):
    """Wait for any of the specified node states to be reached.

    Polls the node state every poll interval seconds, waiting up to the timeout
    for the specified node state to be reached.

    Args:
        nodename (string): The name of the node whose state is being monitored.
        desired_node_states (iterable): The states that the node is expected to reach.
        timeout (integer): The number of seconds to wait before timing out.
        poll_interval (float): Number of seconds between node state polls.
        fatal (boolean): If True, a timeout will cause the test to fail.
        reverse (boolean): If True, wait for the node to lose the desired state.

    Returns:
        Boolean value indicating whether the node ever reached the desired state.

    Example:
        >>> wait_for_node_state_any('node1', ['IDLE', 'ALLOCATED'], timeout=60, poll_interval=5)
        True
        >>> wait_for_node_state_any('node2', ['DOWN'], timeout=30, fatal=True)
        False
    """

    state_set = frozenset(desired_node_states)

    def any_overlap(state):
        return bool(state_set & set(state)) != reverse

    # Wrapper for the repeat_until command to do all our state checking for us
    repeat_until(
        lambda: get_node_parameter(nodename, "state"),
        any_overlap,
        timeout=timeout,
        poll_interval=poll_interval,
        fatal=fatal,
    )

    return any_overlap(get_node_parameter(nodename, "state"))


def wait_for_node_state(
    nodename,
    desired_node_state,
    timeout=default_polling_timeout,
    poll_interval=None,
    fatal=False,
    reverse=False,
):
    """Wait for a specified node state to be reached.

    Polls the node state every poll interval seconds, waiting up to the timeout
    for the specified node state to be reached.

    Args:
        nodename (string): The name of the node whose state is being monitored.
        desired_node_state (string): The state that the node is expected to reach.
        timeout (integer): The number of seconds to wait before timing out.
        poll_interval (float): Number of seconds between node state polls.
        fatal (boolean): If True, a timeout will cause the test to fail.
        reverse (boolean): If True, wait for the node to lose the desired state.

    Returns:
        Boolean value indicating whether the node ever reached the desired state.

    Example:
        >>> wait_for_node_state('node1', 'IDLE', timeout=60, poll_interval=5)
        True
        >>> wait_for_node_state('node2', 'DOWN', timeout=30, fatal=True)
        False
    """

    # Figure out if we're waiting for the desired_node_state to be present or to be gone
    def has_state(state):
        return desired_node_state in state

    def not_state(state):
        return desired_node_state not in state

    # Wrapper for the repeat_until command to do all our state checking for us
    repeat_until(
        lambda: get_node_parameter(nodename, "state"),
        not_state if reverse else has_state,
        timeout=timeout,
        poll_interval=poll_interval,
        fatal=fatal,
    )

    return (desired_node_state in get_node_parameter(nodename, "state")) != reverse


def wait_for_step(job_id, step_id, **repeat_until_kwargs):
    """Wait for the specified step of a job to be running.

    Continuously polls the step state until it becomes running or until a
    timeout occurs.

    Args:
        job_id (integer): The id of the job.
        step_id (integer): The id of the step within the job.

    Returns:
        A boolean value indicating whether the specified step is running or not.

    Example:
        >>> wait_for_step(1234, 0, timeout=60, poll_interval=5, fatal=True)
        True
        >>> wait_for_step(5678, 1, timeout=30)
        False
    """

    step_str = f"{job_id}.{step_id}"
    return repeat_until(
        lambda: run_command_output(f"scontrol -o show step {step_str}"),
        lambda out: re.search(rf"StepId={step_str}", out) is not None,
        **repeat_until_kwargs,
    )


def wait_for_step_accounted(job_id, step_id, **repeat_until_kwargs):
    """Wait for specified job step to appear in accounting database (`sacct`).

    Continuously polls the database until the step is accounted for or until a
    timeout occurs.

    Args:
        job_id (integer): The id of the job.
        step_id (integer): The id of the step within the job.

    Returns:
        A boolean value indicating whether the specified step is accounted for
        in the database or not.

    Example:
        >>> wait_for_step_accounted(1234, 0, timeout=60, poll_interval=5, fatal=True)
        True
        >>> wait_for_step_accounted(5678, 1, timeout=30)
        False
    """

    step_str = f"{job_id}.{step_id}"
    return repeat_until(
        lambda: run_command_output(f"sacct -j {job_id} -o jobid"),
        lambda out: re.search(rf"{step_str}", out) is not None,
        **repeat_until_kwargs,
    )


def wait_for_job_accounted(job_id, field="Start", value=None, **repeat_until_kwargs):
    """Wait for specified job field to appear in accounting database (`sacct`).

    Continuously polls the database until the job field is accounted for or until a
    timeout occurs. If a value is specified, it waits until the field has that value.

    Args:
        job_id (integer): The id of the job.
        field (string): The field to query from the job.
        value (string): The desired value to wait for.
                        None means "any value but empty"

    Returns:
        A boolean value indicating whether the specified job is accounted for
        in the database or not.

    Example:
        >>> wait_for_job_accounted(1234, timeout=60, poll_interval=5, fatal=True)
        True
        >>> wait_for_job_accounted(5678, "End", timeout=30)
        False
    """

    if not value:
        value = r"."
    return repeat_until(
        lambda: run_command_output(f"sacct -Xnj {job_id} -o {field}"),
        lambda out: re.search(value, out.strip()) is not None,
        **repeat_until_kwargs,
    )


def wait_for_job_state(
    job_id,
    desired_job_state,
    desired_reason=None,
    timeout=default_polling_timeout,
    poll_interval=None,
    fatal=False,
    quiet=False,
    xfail=False,
):
    """Wait for the specified job to reach the desired state.

    Continuously polls the job state until the desired state is reached or until
    a timeout occurs.

    Some supported job states are aggregate states, which may include multiple
    discrete states. Some logic is built-in to fail if a job reaches a state
    that makes the desired job state impossible to reach.

    Current supported aggregate states:
    - DONE

    Args:
        job_id (integer): The id of the job.
        desired_job_state (string): The desired state of the job.
        desired_reason (string): Optional reason to also match.
        timeout (integer): The number of seconds to poll before timing out.
        poll_interval (float): Time (in seconds) between job state polls.
        fatal (boolean): If True, a timeout will cause the test to fail.
        quiet (boolean): If True, logging is performed at the TRACE log level.
        xfail (boolean): If True, state (or reason) are not expected to be reached.

    Returns:
        Boolean value indicating whether the job reached the desired state.

    Example:
        >>> wait_for_job_state(1234, 'COMPLETED', timeout=300, poll_interval=10, fatal=True)
        True
        >>> wait_for_job_state(5678, 'RUNNING', timeout=60)
        False
    """

    if poll_interval is None:
        if timeout <= 5:
            poll_interval = 0.1
        elif timeout <= 10:
            poll_interval = 0.2
        else:
            poll_interval = 1

    if quiet:
        log_level = logging.TRACE
    else:
        log_level = logging.DEBUG

    # We don't use repeat_until here because we support pseudo-job states and
    # we want to allow early return (e.g. for a DONE state if we want RUNNING)

    xfail_str = ""
    if xfail:
        xfail_str = "not "
    message = (
        f"Waiting for job ({job_id}) to {xfail_str}reach state {desired_job_state}"
    )
    if desired_reason is not None:
        message += f" and reason {desired_reason}"
    logging.log(log_level, message)

    begin_time = time.time()
    while time.time() < begin_time + timeout:
        job_state = get_job_parameter(
            job_id, "JobState", default="NOT_FOUND", quiet=True
        )

        message = f"Job ({job_id}) is in state {job_state}, but we are waiting for {desired_job_state}"
        if job_state in [
            "NOT_FOUND",
            "BOOT_FAIL",
            "CANCELLED",
            "COMPLETED",
            "DEADLINE",
            "FAILED",
            "NODE_FAIL",
            "OUT_OF_MEMORY",
            "TIMEOUT",
            "PREEMPTED",
        ]:
            if desired_job_state == "DONE" or job_state == desired_job_state:
                message = f"Job ({job_id}) is in the {xfail_str}desired state {desired_job_state}"
                reason = get_job_parameter(
                    job_id, "Reason", default="NOT_FOUND", quiet=True
                )
                if desired_reason is None or reason == desired_reason:
                    if desired_reason is not None:
                        message += (
                            f" with the {xfail_str}desired reason {desired_reason}"
                        )
                    if not xfail:
                        logging.log(log_level, message)
                    else:
                        logging.warning(message)
                    return True
                else:
                    message += (
                        f", but with reason {reason} and we waited for {desired_reason}"
                    )

            if not xfail:
                if fatal:
                    pytest.fail(message)
                else:
                    logging.warning(message)
            else:
                logging.log(log_level, message)
            return False
        elif job_state == desired_job_state:
            message = (
                f"Job ({job_id}) is in the {xfail_str}desired state {desired_job_state}"
            )
            reason = get_job_parameter(
                job_id, "Reason", default="NOT_FOUND", quiet=True
            )
            if desired_reason is None or reason == desired_reason:
                if desired_reason is not None:
                    message += f" with the {xfail_str}desired reason {desired_reason}"
                if not xfail:
                    logging.log(log_level, message)
                else:
                    logging.warning(message)
                return True
            else:
                message += (
                    f", but with reason {reason} and we waited for {desired_reason}"
                )

        logging.log(log_level, message)
        time.sleep(poll_interval)

    message = f"Job ({job_id}) did not reach the {desired_job_state} state"
    if desired_reason is not None:
        message += f" or the reason {desired_reason}"
    message += f" within the {timeout} second(s) timeout"
    if not xfail:
        if fatal:
            pytest.fail(message)
        else:
            logging.warning(message)
    else:
        logging.log(log_level, message)

    return False


def check_steps_delayed(job_id, job_output, expected_delayed):
    """Check the output file of a job for expected delayed steps.

    This function checks that the output file for a job contains the expected
    pattern of delayed job steps. Note that at the time of writing, this
    requires srun steps to have at least a verbosity level of "-vv" to log
    their f"srun: Step completed in JobId={job_id}, retrying" notification.

    Args:
        job_id (integer): The job ID that we're interested in.
        job_output (string): The content of the output file of the job.
        expected_delayed (integer): The initial number of delayed job steps. It
            is verified that this initial number of job steps are delayed and
            then this number of delayed job steps decrements one by one as
            running job steps finish.

    Returns:
        True if steps were delayed in the correct amounts and order, else False.

    Example:
        >>> check_steps_delayed(123, "srun: Received task exit notification for 1 task of StepId=1.0 (status=0x0000).\nsrun: node1: task 0: Completed\nsrun: debug:  task 0 done\nsrun: Step completed in JobId=1, retrying\nsrun: Step completed in JobId=1, retrying\nsrun: Step created for StepId=1.1\nsrun: Received task exit notification for 1 task of StepId=1.1 (status=0x0000).\nsrun: node2: task 1: Completed\nsrun: debug:  task 1 done\nsrun: debug:  IO thread exiting\nsrun: Step completed in JobId=1, retrying\nsrun: Step created for StepId=1.2", 2)
        True
        >>> check_steps_delayed(456, "srun: Received task exit notification for 1 task of StepId=1.0 (status=0x0000).\nsrun: node1: task 0: Completed\nsrun: debug:  task 0 done\nsrun: Step completed in JobId=1, retrying\nsrun: Step completed in JobId=1, retrying\nsrun: Step created for StepId=1.1\nsrun: Received task exit notification for 1 task of StepId=1.1 (status=0x0000).\nsrun: node2: task 1: Completed\nsrun: debug:  task 1 done\nsrun: debug:  IO thread exiting", 2)
        False
    """

    # Iterate through each group of expected delayed steps. For example,
    # if there was a job that had 5 steps that could run in parallel but, due to
    # resource constrains, only allowed 3 steps to run at a time, we would
    # expect a group of 2 delayed job steps followed by a group of 1 delayed job
    # steps. For this example job, expected_delayed=2.
    #
    # The idea of the for loop below is to iterate through each group of delayed
    # job steps and replace the expected output as we go with re.sub. This
    # ensures that the delayed job step groups occur in the correct order.
    #
    # Each regex pattern matches part of the pattern we'd expect to see in the
    # output, replaces the matched text (see previous paragraph), and then makes
    # sure there is still text left to match for the rest of the pattern. If the
    # regex pattern doesn't match anything, then re.sub will match and replace
    # all the rest of the output and leave job_output empty.
    for delayed_grp_size in range(expected_delayed, 0, -1):
        # Match all lines before receiving an exit notification. This regex
        # pattern will match any line that doesn't contain "srun: Received task
        # exit notification".
        before_start_pattern = r"(^((?!srun: Received task exit notification).)*$\n)*"
        job_output = re.sub(before_start_pattern, "", job_output, 1, re.MULTILINE)
        if not job_output:
            logging.error(f"Pattern not found: {before_start_pattern}")
            return False

        # Match receiving the next exit notification. This regex pattern will
        # match the line where the exit notification is received when a step
        # that was already running finishes.
        exit_pattern = rf"srun: Received task exit notification for \[0-9]+ task of StepId={job_id}\.[0-9]+ \(status=0x[0-9A-Fa-f]+\)\.\n"
        job_output = re.sub(exit_pattern, "", job_output, 1, re.MULTILINE)
        if not job_output:
            logging.error(f"Pattern not found: {exit_pattern}")
            return False

        # Match lines we don't want before a step completion. After the exit
        # notification, we now match all lines that don't contain "srun: Step
        # completed". Sometimes an exit notification can be received multiple
        # times and any redundant exit notifications are also matched by this
        # pattern.
        before_completed_pattern = r"(^((?!srun: Step completed).)*$\n)*"
        job_output = re.sub(before_completed_pattern, "", job_output, 1, re.MULTILINE)
        if not job_output:
            logging.error(f"Pattern not found: {before_completed_pattern}")
            return False

        # Match number of lines retrying to start a delayed job step. Note that
        # this pattern searched for steps retrying "delayed_grp_size" number of
        # times. This is because every step that is delayed retries every time a
        # previously running step finishes.
        completed_pattern = rf"(srun: Step completed in JobId={job_id}, retrying\n){{{delayed_grp_size}}}"
        job_output = re.sub(completed_pattern, "", job_output, 1, re.MULTILINE)
        if not job_output:
            logging.error(f"Pattern not found: {completed_pattern}")
            return False

        # Match lines we don't want before a step creation. Due to steps running
        # in parallel, other lines of text can be output from already running
        # steps before we're told a new step has been created. This regex
        # pattern matches all lines that don't contain "srun: Step created".
        before_created_pattern = r"(^((?!srun: Step created).)*$\n)*"
        job_output = re.sub(before_created_pattern, "", job_output, 1, re.MULTILINE)
        if not job_output:
            logging.error(f"Pattern not found: {before_created_pattern}")
            return False

        # Match the step creation line for the delayed step
        created_pattern = rf"srun: Step created for StepId={job_id}\.[0-9]+"
        job_output = re.sub(created_pattern, "", job_output, 1, re.MULTILINE)
        if not job_output:
            logging.error(f"Pattern not found: {created_pattern}")
            return False

    return True


def create_node(node_dict):
    """Creates a node with the properties described by the supplied dictionary.

    This function is currently only used as a helper function within other
    library functions (e.g. require_nodes). It modifies the Slurm configuration
    file and restarts the relevant Slurm daemons. A backup is automatically
    created, and the original configuration is restored after the test
    completes. This function may only be used in auto-config mode.

    Args:
        node_dict (dictionary): A dictionary containing the desired node
            properties.

    Returns:
        None

    Example:
        >>> create_node({'NodeName': 'new_node', 'CPUs': 4, 'RealMemory': 16384})
    """

    if not properties["auto-config"]:
        require_auto_config("wants to add a node")

    config_file = f"{properties['slurm-config-dir']}/slurm.conf"

    # Read the original slurm.conf into a list of lines
    output = run_command_output(
        f"cat {config_file}", user=properties["slurm-user"], quiet=True
    )
    original_config_lines = output.splitlines()
    new_config_lines = original_config_lines.copy()

    # Locate the last NodeName definition
    last_node_line_index = 0
    for line_index in range(len(original_config_lines)):
        line = original_config_lines[line_index]

        if re.search(r"(?i)^ *NodeName", line) is not None:
            last_node_line_index = line_index
    if last_node_line_index == 0:
        last_node_line_index = line_index

    # Build up the new node line
    node_line = ""
    if "NodeName" in node_dict:
        node_line = f"NodeName={node_dict['NodeName']}"
        node_range = node_dict.pop("NodeName")
    if "Port" in node_dict:
        node_line += f" Port={node_dict['Port']}"
        node_dict.pop("Port")
    for parameter_name, parameter_value in sorted(node_dict.items()):
        node_line += f" {parameter_name}={parameter_value}"

    # Add the new node line
    new_config_lines.insert(last_node_line_index + 1, node_line)

    # Write the config file back out with the modifications
    new_config_string = "\n".join(new_config_lines)
    run_command(
        f"echo '{new_config_string}' > {config_file}",
        user=properties["slurm-user"],
        fatal=True,
        quiet=True,
    )

    # Create the required node directories
    for node_name in node_range_to_list(node_range):
        spool_dir = properties["slurm-spool-dir"].replace("%n", node_name)
        tmpfs_dir = properties["slurm-tmpfs"].replace("%n", node_name)
        properties["nodes"].append(node_name)

        run_command(f"sudo mkdir -p {spool_dir}", fatal=True, quiet=True)
        run_command(f"sudo mkdir -p {tmpfs_dir}", fatal=True, quiet=True)

    # Restart slurm if it is already running
    if is_slurmctld_running(quiet=True):
        restart_slurm(quiet=True)


# requirements_list is a list of (parameter_name, parameter_value) tuples.
# Uses non-live node info because must copy from existing node config line
# We implemented requirements_list as a list of tuples so that this could
# later be extended to include a comparator, etc.
# atf.require_nodes(1, [('CPUs', 4), ('RealMemory', 40)])
# atf.require_nodes(2, [('Gres', 'gpu:1,mps:100')])
def require_nodes(requested_node_count, requirements_list=[]):
    """Ensure that a requested number of nodes have the required properties.

    In local-config mode, the test is skipped if an insufficient number of
    nodes possess the required properties.
    In auto-config mode, nodes are created as needed with the required
    properties.

    Args:
        requested_node_count (integer): Number of required nodes.
        requirements_list (list of tuples): List of (parameter_name,
            parameter_value) tuples.

    Currently supported node requirement types include:
        CPUs
        Cores
        RealMemory
        Gres
        Features

    Other node requirement types will still be appended to the requirements,
    but this could stop slurm from starting.

    Returns:
        None

    Example:
        >>> require_nodes(2, [('CPUs', 4), ('RealMemory', 40)])
        >>> require_nodes(2, [('CPUs', 2), ('RealMemory', 30), ('Features', 'gpu,mpi')])
        >>> require_nodes(2, [('CPUs', 4), ('Sockets', 1)])
    """

    # If using local-config and slurm is running, use live node information
    # so that a test is not incorrectly skipped when slurm derives a non-single
    # CPUTot while the slurm.conf does not contain a CPUs property.
    if not properties["auto-config"] and is_slurmctld_running(quiet=True):
        live = True
    else:
        live = False

    # This should return separate nodes and a DEFAULT (unless live)
    nodes_dict = get_nodes(live=live, quiet=True)
    original_nodes = {}
    default_node = {}

    # Instantiate original node names from nodes_dict
    for node_name in nodes_dict:
        # Split out the default node
        if node_name == "DEFAULT":
            default_node = nodes_dict[node_name]
        else:
            original_nodes[node_name] = {}

    # Populate with any default parameters
    if default_node:
        for node_name in original_nodes:
            for parameter_name, parameter_value in default_node.items():
                if parameter_name.lower() != "nodename":
                    original_nodes[node_name][parameter_name] = parameter_value

    # Merge in parameters from nodes_dict
    for node_name in original_nodes:
        for parameter_name, parameter_value in nodes_dict[node_name].items():
            if parameter_name.lower() != "nodename":
                # Translate CPUTot to CPUs for screening qualifying nodes
                if parameter_name.lower() == "cputot":
                    parameter_name = "CPUs"
                original_nodes[node_name][parameter_name] = parameter_value

    # Check to see how many qualifying nodes we have
    qualifying_node_count = 0
    node_count = 0
    nonqualifying_node_count = 0
    first_node_name = ""
    first_qualifying_node_name = ""
    node_indices = {}
    augmentation_dict = {}
    for node_name in sorted(original_nodes):
        lower_node_dict = dict(
            (key.lower(), value) for key, value in original_nodes[node_name].items()
        )
        node_count += 1

        if node_count == 1:
            first_node_name = node_name

        # Build up node indices for use when having to create new nodes
        match = re.search(r"^(.*?)(\d*)$", node_name)
        node_prefix, node_index = match.group(1), match.group(2)
        if node_index == "":
            node_indices[node_prefix] = node_indices.get(node_prefix, [])
        else:
            node_indices[node_prefix] = node_indices.get(node_prefix, []) + [
                int(node_index)
            ]

        node_qualifies = True
        for requirement_tuple in requirements_list:
            parameter_name, parameter_value = requirement_tuple[0:2]
            if parameter_name in ["CPUs", "RealMemory"]:
                if parameter_name.lower() in lower_node_dict:
                    if lower_node_dict[parameter_name.lower()] < parameter_value:
                        if node_qualifies:
                            node_qualifies = False
                            nonqualifying_node_count += 1
                        if nonqualifying_node_count == 1:
                            augmentation_dict[parameter_name] = parameter_value
                else:
                    if node_qualifies:
                        node_qualifies = False
                        nonqualifying_node_count += 1
                    if nonqualifying_node_count == 1:
                        augmentation_dict[parameter_name] = parameter_value
            elif parameter_name == "Cores":
                boards = lower_node_dict.get("boards", 1)
                sockets_per_board = lower_node_dict.get("socketsperboard", 1)
                cores_per_socket = lower_node_dict.get("corespersocket", 1)
                sockets = boards * sockets_per_board
                cores = sockets * cores_per_socket
                if cores < parameter_value:
                    if node_qualifies:
                        node_qualifies = False
                        nonqualifying_node_count += 1
                    if nonqualifying_node_count == 1:
                        augmentation_dict["CoresPerSocket"] = math.ceil(
                            parameter_value / sockets
                        )
            elif parameter_name == "Gres":
                if parameter_name.lower() in lower_node_dict:
                    gres_list = parameter_value.split(",")
                    for gres_value in gres_list:
                        if match := re.search(r"^(\w+):(\d+)$", gres_value):
                            (required_gres_name, required_gres_value) = (
                                match.group(1),
                                match.group(2),
                            )
                        else:
                            pytest.fail(
                                "Gres requirement must be of the form <name>:<count>"
                            )
                        if match := re.search(
                            rf"{required_gres_name}:(\d+)",
                            lower_node_dict[parameter_name.lower()],
                        ):
                            if match.group(1) < required_gres_value:
                                if node_qualifies:
                                    node_qualifies = False
                                    nonqualifying_node_count += 1
                                if nonqualifying_node_count == 1:
                                    augmentation_dict[parameter_name] = gres_value
                        else:
                            if node_qualifies:
                                node_qualifies = False
                                nonqualifying_node_count += 1
                            if nonqualifying_node_count == 1:
                                augmentation_dict[parameter_name] = gres_value
                else:
                    if node_qualifies:
                        node_qualifies = False
                        nonqualifying_node_count += 1
                    if nonqualifying_node_count == 1:
                        augmentation_dict[parameter_name] = parameter_value
            elif parameter_name == "Features":
                required_features = set(parameter_value.split(","))
                features = lower_node_dict.get("features", [])
                features = features[0] if features else ""
                node_features = set(features.split(","))
                if not required_features.issubset(node_features):
                    if node_qualifies:
                        node_qualifies = False
                        nonqualifying_node_count += 1
                    if nonqualifying_node_count == 1:
                        augmentation_dict[parameter_name] = parameter_value
            else:
                logging.debug(
                    f"{parameter_name} is not a supported node requirement type."
                )
                logging.debug(
                    f"{parameter_name}={parameter_value} will be added anyways!"
                )
                augmentation_dict[parameter_name] = parameter_value
                if node_qualifies:
                    node_qualifies = False
                    nonqualifying_node_count += 1
        if node_qualifies:
            qualifying_node_count += 1
            if first_qualifying_node_name == "":
                first_qualifying_node_name = node_name

    # Not enough qualifying nodes
    if qualifying_node_count < requested_node_count:
        # If auto-config, configure what is required
        if properties["auto-config"]:
            # Create new nodes to meet requirements ignoring default node0
            new_node_count = requested_node_count

            # If we already have a qualifying node, we will use it as the template
            if qualifying_node_count > 0:
                template_node_name = first_qualifying_node_name
                template_node = nodes_dict[template_node_name].copy()
            # Otherwise we will use the first node as a template and augment it
            else:
                template_node_name = first_node_name
                template_node = nodes_dict[template_node_name].copy()
                for parameter_name, parameter_value in augmentation_dict.items():
                    template_node[parameter_name] = parameter_value

            base_port = int(nodes_dict[template_node_name]["Port"])

            # Build up a list of available new indices starting after the template
            match = re.search(r"^(.*?)(\d*)$", template_node_name)
            template_node_prefix, template_node_index = match.group(1), int(
                match.group(2)
            )
            used_indices = sorted(node_indices[template_node_prefix])
            new_indices = []
            new_index = template_node_index
            for i in range(new_node_count):
                new_index += 1
                while new_index in used_indices:
                    new_index += 1
                new_indices.append(new_index)

            # Create a new aggregate node
            # Later, we could consider collapsing the node into the template node if unmodified
            new_node_dict = template_node.copy()
            if new_node_count == 1:
                new_node_dict["NodeName"] = template_node_prefix + str(new_indices[0])
                new_node_dict["Port"] = base_port - template_node_index + new_indices[0]
            else:
                new_node_dict["NodeName"] = (
                    f"{template_node_prefix}[{list_to_range(new_indices)}]"
                )
                new_node_dict["Port"] = list_to_range(
                    list(
                        map(lambda x: base_port - template_node_index + x, new_indices)
                    )
                )
            create_node(new_node_dict)

        # If local-config, skip
        else:
            message = f"This test requires {requested_node_count} nodes"
            if requirements_list:
                message += f" with {requirements_list}"
            pytest.skip(message, allow_module_level=True)


def make_bash_script(script_name, script_contents):
    """Creates an executable Bash script with the specified contents.

    Args:
        script_name (string): Name of the script to create.
        script_contents (string): Contents of the script.

    Returns:
        None

    Example:
        >>> make_bash_script("my_script.sh", "echo 'Hello, World!'") # This creates an executable Bash script named "my_script.sh" with the contents "echo 'Hello, World!'"
    """

    with open(script_name, "w") as f:
        f.write("#!/bin/bash\n")
        f.write(script_contents)

    run_command(f"chmod 777 {script_name}", user="root", fatal=True, quiet=True)

    if properties["test-user-set"]:
        run_command(
            f"chown {properties['test-user']} {script_name}",
            user="root",
            fatal=True,
            quiet=True,
        )


def wait_for_file(file_name, **repeat_until_kwargs):
    """Waits for the specified file to be present.

    This function waits up to the specified timeout seconds for the file to be
    present, polling every poll_interval seconds. The default timeout and
    poll_interval are inherited from repeat_until.

    Args:
        file_name (string): The file name.

    Returns:
        True if the file was found, False otherwise.

    Example:
        >>> wait_for_file("my_file.txt", timeout=30, poll_interval=0.5)
        True
    """

    logging.debug(f"Waiting for file ({file_name}) to be present")
    return repeat_until(
        lambda: os.path.isfile(file_name), lambda exists: exists, **repeat_until_kwargs
    )


# Assuming this will only be called internally after validating accounting is configured and auto-config is set
def backup_accounting_database():
    """Backs up the accounting database.

    This function may only be used in auto-config mode. The database dump is
    automatically restored when the test ends.

    Args:
        None

    Returns:
        None

    Example:
        >>> backup_accounting_database() # Backs up Slurm accounting database to file in the test's temporary directory.
    """

    if not properties["auto-config"]:
        return

    mysqldump_path = shutil.which("mysqldump")
    if mysqldump_path is None:
        pytest.fail(
            "Unable to backup the accounting database. mysqldump was not found in your path"
        )
    mysql_path = shutil.which("mysql")
    if mysql_path is None:
        pytest.fail(
            "Unable to backup the accounting database. mysql was not found in your path"
        )

    sql_dump_file = f"{str(module_tmp_path / '../../slurm_acct_db.sql')}"

    # If a dump already exists, issue a warning and return (honor existing dump)
    if os.path.isfile(sql_dump_file):
        logging.warning(f"Dump file already exists ({sql_dump_file})")
        return

    slurmdbd_dict = get_config(live=False, source="slurmdbd", quiet=True)
    database_host, database_port, database_name, database_user, database_password = (
        slurmdbd_dict.get(key)
        for key in [
            "StorageHost",
            "StoragePort",
            "StorageLoc",
            "StorageUser",
            "StoragePass",
        ]
    )

    mysql_options = ""
    if database_host:
        mysql_options += f" -h {database_host}"
    if database_port:
        mysql_options += f" -P {database_port}"
    if database_user:
        mysql_options += f" -u {database_user}"
    else:
        mysql_options += f" -u {properties['slurm-user']}"
    if database_password:
        mysql_options += f" -p {database_password}"

    if not database_name:
        database_name = "slurm_acct_db"

    mysql_command = f"{mysql_path} {mysql_options} -e \"USE '{database_name}'\""
    if run_command_exit(mysql_command, quiet=True) != 0:
        logging.debug(f"Slurm accounting database ({database_name}) is not present")
    else:
        mysqldump_command = (
            f"{mysqldump_path} {mysql_options} {database_name} > {sql_dump_file}"
        )
        run_command(
            mysqldump_command, fatal=True, quiet=False, timeout=default_sql_cmd_timeout
        )


def restore_accounting_database():
    """Restores the accounting database from the backup.

    This function may only be used in auto-config mode.

    Args:
        None

    Returns:
        None

    Example:
        >>> restore_accounting_database() # Restores Slurm accounting database from previously created backup.
    """

    if not properties["auto-config"]:
        return

    mysql_path = shutil.which("mysql")
    if mysql_path is None:
        pytest.fail(
            "Unable to restore the accounting database. mysql was not found in your path"
        )

    slurmdbd_dict = get_config(live=False, source="slurmdbd", quiet=True)
    database_host, database_port, database_name, database_user, database_password = (
        slurmdbd_dict.get(key)
        for key in [
            "StorageHost",
            "StoragePort",
            "StorageLoc",
            "StorageUser",
            "StoragePass",
        ]
    )
    if not database_name:
        database_name = "slurm_acct_db"

    base_command = mysql_path
    if database_host:
        base_command += f" -h {database_host}"
    if database_port:
        base_command += f" -P {database_port}"
    if database_user:
        base_command += f" -u {database_user}"
    else:
        base_command += f" -u {properties['slurm-user']}"
    if database_password:
        base_command += f" -p {database_password}"

    # If DB exists, drop it and try to restore the dump file
    mysql_command = f"{base_command} -e \"USE '{database_name}'\""
    if run_command_exit(mysql_command, quiet=True) == 0:
        run_command(
            f'{base_command} -e "drop database {database_name}"',
            fatal=True,
            quiet=False,
            timeout=default_sql_cmd_timeout,
        )

    sql_dump_file = f"{str(module_tmp_path / '../../slurm_acct_db.sql')}"

    # If the dump file doesn't exist, it has probably already been
    # restored by a previous call to restore_accounting_database
    if not os.path.isfile(sql_dump_file):
        logging.debug(
            f"Slurm accounting database backup ({sql_dump_file}) is not present. It has probably already been restored."
        )
        return

    dump_stat = os.stat(sql_dump_file)
    if not (dump_stat.st_size == 0 and dump_stat.st_mode & stat.S_ISVTX):
        run_command(
            f'{base_command} -e "create database {database_name}"',
            fatal=True,
            quiet=False,
        )
        run_command(
            f"{base_command} {database_name} < {sql_dump_file}",
            fatal=True,
            quiet=False,
            timeout=default_sql_cmd_timeout,
        )

    # In either case, remove the dump file
    run_command(f"rm -f {sql_dump_file}", fatal=True, quiet=False)


def run_check_test(source_file, build_args=""):
    """Compiles and runs a libcheck test
    Args:
        source_file (string): The name of the source file, relative to testsuite_check_dir.
        build_args (string): Additional string to be appended to the build command.
    """

    import xmltodict

    check_test = (
        f"{module_tmp_path}/{os.path.splitext(os.path.basename(source_file))[0]}"
    )
    xml_test = check_test + ".xml"

    compile_against_libslurm(
        f"{properties['testsuite_check_dir']}/{source_file}",
        check_test,
        full=True,
        build_args=build_args + "-lcheck -lm -lsubunit",
        fatal=True,
        quiet=True,
    )

    # Run the libcheck test setting an xml output
    result = run_command(
        check_test, quiet=True, env_vars=f"CK_XML_LOG_FILE_NAME={xml_test}"
    )

    # Parse the xml output
    if not os.path.exists(xml_test):
        pytest.fail(f"Test results not found: {xml_test}")

    with open(xml_test) as f:
        xml_data = xmltodict.parse(f.read())

    logging.info(f"{result['stdout']}")
    if result["exit_code"]:
        logging.error(f"\n{result['stderr']}")

    return xml_data["testsuites"]["suite"]


def compile_against_libslurm(
    source_file,
    dest_file,
    build_args="",
    full=False,
    shared=False,
    new_prefixes=False,
    **run_command_kwargs,
):
    """Compiles a test program against either libslurm.so or libslurmfull.so.

    This function compiles the specified source file against the Slurm library,
    either libslurm.so or libslurmfull.so, and creates the target binary file.

    Args:
        source_file (string): The name of the source file.
        dest_file (string): The name of the target binary file.
        build_args (string): Additional string to be appended to the build command.
        full (boolean): Use libslurmfull.so instead of libslurm.so.
        shared (boolean): Produces a shared library (adds the -shared compiler option
            and adds a .so suffix to the output file name).
        **run_command_kwargs: Auxiliary arguments to be passed to the
            run_command function (e.g., quiet, fatal, timeout, etc.).

    Returns:
        None

    Example:
        >>> compile_against_libslurm("my_test.c", "my_test", build_args="-Wall -Werror")
    """

    slurm_prefix = properties["slurm-prefix"]
    slurm_source = properties["slurm-source-dir"]
    slurm_build = properties["slurm-build-dir"]
    if new_prefixes:
        slurm_prefix = properties["new-slurm-prefix"]
        slurm_source = properties["new-source-prefix"]
        slurm_build = properties["new-build-prefix"]

    if full:
        slurm_library = "slurmfull"
    else:
        slurm_library = "slurm"
    if os.path.isfile(f"{slurm_prefix}/lib64/slurm/lib{slurm_library}.so"):
        lib_dir = "lib64"
    else:
        lib_dir = "lib"
    if full:
        lib_path = f"{slurm_prefix}/{lib_dir}/slurm"
    else:
        lib_path = f"{slurm_prefix}/{lib_dir}"

    command = f"gcc {source_file} -g -pthread"
    if shared:
        command += " -fPIC -shared"
    command += f" -o {dest_file}"
    command += f" -I{slurm_source} -I{slurm_build} -I{slurm_prefix}/include -Wl,-rpath={lib_path} -L{lib_path} -l{slurm_library} -lresolv"
    if build_args != "":
        command += f" {build_args}"
    run_command(command, **run_command_kwargs)


def get_partitions(**run_command_kwargs):
    """Returns the Slurm partition configuration as a dictionary of dictionaries.

    Args:
        **run_command_kwargs: Auxiliary arguments to be passed to the
            run_command function (e.g., quiet, fatal, timeout, etc.).

    Returns:
        A dictionary of dictionaries, where the first-level keys are the
        partition names, and the values are dictionaries containing the
        configuration parameters for the respective partitions.

    Example:
        >>> get_partitions(quiet=True)
        {'partition1': {'PartitionName': 'partition1', 'AllowGroups': 'ALL', 'Defaults': 'YES', ...},
         'partition2': {'PartitionName': 'partition2', 'AllowGroups': 'group1,group2', 'Defaults': 'YES', ...}}
    """

    partitions_dict = {}

    output = run_command_output(
        "scontrol show partition -o", fatal=True, **run_command_kwargs
    )

    partition_dict = {}
    for line in output.splitlines():
        if line == "":
            continue

        while match := re.search(r"^ *([^ =]+)=(.*?)(?= +[^ =]+=| *$)", line):
            param_name, param_value = match.group(1), match.group(2)

            # Remove the consumed parameter from the line
            line = re.sub(r"^ *([^ =]+)=(.*?)(?= +[^ =]+=| *$)", "", line)

            # Reformat the value if necessary
            if is_integer(param_value):
                param_value = int(param_value)
            elif is_float(param_value):
                param_value = float(param_value)
            elif param_value == "(null)":
                param_value = None

            # Add it to the temporary partition dictionary
            partition_dict[param_name] = param_value

        # Add the partition dictionary to the partitions dictionary
        partitions_dict[partition_dict["PartitionName"]] = partition_dict

        # Clear the partition dictionary for use by the next partition
        partition_dict = {}

    return partitions_dict


def get_partition_parameter(partition_name, parameter_name, default=None):
    """Obtains the value for a Slurm partition configuration parameter.

    This function retrieves the value of the specified parameter for the given
    partition. If the parameter is not present, the default value is returned.

    Args:
        partition_name (string): The name of the partition.
        parameter_name (string): The name of the parameter to retrieve.
        default (string or None): The default value to return if the parameter is
            not found.

    Returns:
        The value of the specified partition parameter, or the default if not
        found.

    Example:
        >>> get_partition_parameter('my_partition', 'AllowAccounts')
        'ALL'
        >>> get_partition_parameter('second_partition', 'DefaultTime', '00:30:00')
        '00:30:00'
    """

    partitions_dict = get_partitions()

    if partition_name in partitions_dict:
        partition_dict = partitions_dict[partition_name]
    else:
        pytest.fail(
            f"Partition ({partition_name}) was not found in the partition configuration"
        )

    if parameter_name in partition_dict:
        return partition_dict[parameter_name]
    else:
        return default


def set_partition_parameter(partition_name, new_parameter_name, new_parameter_value):
    """Sets the value of the specified partition configuration parameter.

    This function the specified partition property and reconfigures
    the slurm daemons. A backup is automatically created and the
    original configuration is restored after the test completes.

    This function may only be used in auto-config mode.

    Args:
        partition_name (string): The partition name.
        new_parameter_name (string): The parameter name.
        new_parameter_value (string): The parameter value.
            Use a value of None to unset a partition parameter.

    Example:
        >>> set_partition_parameter('partition1', 'MaxTime', 'INFINITE')
    """

    if not properties["auto-config"]:
        require_auto_config("wants to modify partition parameters")

    config_file = f"{properties['slurm-config-dir']}/slurm.conf"

    # Read the original slurm.conf into a list of lines
    output = run_command_output(
        f"cat {config_file}", user=properties["slurm-user"], quiet=True
    )
    original_config_lines = output.splitlines()
    new_config_lines = original_config_lines.copy()

    # Locate the partition among the various Partition definitions
    found_partition_name = False
    for line_index in range(len(original_config_lines)):
        line = original_config_lines[line_index]

        words = re.split(r" +", line.strip())
        if len(words) < 1:
            continue
        if words[0][0] == "#":
            continue
        parameter_name, parameter_value = words[0].split("=", 1)
        if parameter_name.lower() != "partitionname":
            continue

        if parameter_value == partition_name:
            # We found a matching PartitionName line
            found_partition_name = True

            # Read in the partition parameters
            original_partition_parameters = collections.OrderedDict()
            for word in words[1:]:
                parameter_name, parameter_value = word.split("=", 1)
                original_partition_parameters[parameter_name] = parameter_value

            # Delete the original partition definition
            new_config_lines.pop(line_index)

            # Add the modified definition for the specified partition
            modified_partition_parameters = original_partition_parameters.copy()
            if new_parameter_value is None:
                if new_parameter_name in modified_partition_parameters:
                    del modified_partition_parameters[new_parameter_name]
            else:
                modified_partition_parameters[new_parameter_name] = new_parameter_value
            modified_partition_line = f"PartitionName={partition_name}"
            for (
                parameter_name,
                parameter_value,
            ) in modified_partition_parameters.items():
                modified_partition_line += f" {parameter_name}={parameter_value}"
            new_config_lines.insert(line_index, modified_partition_line)

            break

    if not found_partition_name:
        pytest.fail(
            f"Invalid partition name specified in set_partition_parameter(). Partition {partition_name} does not exist"
        )

    # Write the config file back out with the modifications
    new_config_string = "\n".join(new_config_lines)
    run_command(
        f"echo '{new_config_string}' > {config_file}",
        user=properties["slurm-user"],
        fatal=True,
        quiet=True,
    )

    # Reconfigure slurm controller if it is already running
    if is_slurmctld_running(quiet=True):
        run_command("scontrol reconfigure", user=properties["slurm-user"], quiet=True)


def default_partition():
    """Returns the name of the default Slurm partition.

    This function retrieves the Slurm partition configuration and returns the
    name of the partition that is marked as the default.

    Args:
        None

    Returns:
        The name of the default Slurm partition.

    Example:
        >>> default_partition()
        'my_default_partition'
    """

    partitions_dict = get_partitions()

    for partition_name in partitions_dict:
        if partitions_dict[partition_name]["Default"] == "YES":
            return partition_name


# This is supplied for ease-of-use in test development only.
# Tests should not use this permanently. Use logging.debug() instead.
def log_debug(msg):
    logging.debug(msg)


##############################################################################
# ATF module initialization
##############################################################################


# This is a logging filter that adds a new LogRecord traceback attribute
class TraceBackFilter(logging.Filter):
    def filter(self, record):
        call_stack = []
        within_atf_context = False

        for frame_summary in (traceback.extract_stack())[-5::-1]:
            if within_atf_context:
                if "testsuite/python" not in frame_summary.filename:
                    break
            else:
                if "testsuite/python" in frame_summary.filename:
                    within_atf_context = True
                else:
                    continue

            function = frame_summary.name
            short_filename = frame_summary.filename.rpartition("testsuite/python/")[2]
            lineno = frame_summary.lineno

            call_stack.append(f"{function}@{short_filename}:{lineno}")

        record.traceback = ",".join(call_stack)

        return True


# Add a new traceback LogRecord attribute
logging.getLogger().addFilter(TraceBackFilter())

# Add a custom TRACE logging level
# This has to be done early enough to allow pytest --log-level=TRACE to be used
logging.TRACE = logging.NOTSET + 5
logging.addLevelName(logging.TRACE, "TRACE")


def _trace(message, *args, **kwargs):
    logging.log(logging.TRACE, message, *args, **kwargs)


logging.trace = _trace
logging.getLogger().trace = _trace

# Add a custom NOTE logging level in between INFO and DEBUG
logging.NOTE = logging.DEBUG + 5
logging.addLevelName(logging.NOTE, "NOTE")


def _note(message, *args, **kwargs):
    logging.log(logging.NOTE, message, *args, **kwargs)


logging.note = _note
logging.getLogger().note = _note

# The module-level temporary directory is initialized in conftest.py
module_tmp_path = None

# Instantiate and populate testrun-level properties
properties = {}

# Initialize directory properties
properties["testsuite_base_dir"] = str(pathlib.Path(__file__).resolve().parents[2])
properties["testsuite_python_lib"] = properties["testsuite_base_dir"] + "/python/lib"
properties["slurm-source-dir"] = str(pathlib.Path(__file__).resolve().parents[3])
properties["slurm-build-dir"] = properties["slurm-source-dir"]
properties["slurm-prefix"] = "/usr/local"
properties["testsuite_scripts_dir"] = (
    properties["testsuite_base_dir"] + "/python/scripts"
)
properties["testsuite_data_dir"] = properties["testsuite_base_dir"] + "/python/data"
properties["testsuite_check_dir"] = properties["testsuite_base_dir"] + "/python/check"
properties["influxdb_host"] = "localhost"
properties["influxdb_port"] = 8086
properties["influxdb_db"] = "slurm"

# Override directory properties with values from testsuite.conf file
testsuite_config = {}
# The default location for the testsuite.conf file (in SRCDIR/testsuite)
# can be overridden with the SLURM_TESTSUITE_CONF environment variable.
testsuite_config_file = os.getenv(
    "SLURM_TESTSUITE_CONF", f"{properties['testsuite_base_dir']}/testsuite.conf"
)
if not os.path.isfile(testsuite_config_file):
    pytest.fail(
        f"The unified testsuite configuration file (testsuite.conf) was not found. This file can be created from a copy of the autogenerated sample found in BUILDDIR/testsuite/testsuite.conf.sample. By default, this file is expected to be found in SRCDIR/testsuite ({properties['testsuite_base_dir']}). If placed elsewhere, set the SLURM_TESTSUITE_CONF environment variable to the full path of your testsuite.conf file."
    )
with open(testsuite_config_file, "r") as f:
    for line in f.readlines():
        if match := re.search(r"^\s*(\w+)\s*=\s*(.*)$", line):
            testsuite_config[match.group(1).lower()] = match.group(2)
if "slurmsourcedir" in testsuite_config:
    properties["slurm-source-dir"] = testsuite_config["slurmsourcedir"]
if "slurmbuilddir" in testsuite_config:
    properties["slurm-build-dir"] = testsuite_config["slurmbuilddir"]
if "slurminstalldir" in testsuite_config:
    properties["slurm-prefix"] = testsuite_config["slurminstalldir"]
if "slurmconfigdir" in testsuite_config:
    properties["slurm-config-dir"] = testsuite_config["slurmconfigdir"]

if "influxdb_host" in testsuite_config:
    properties["influxdb_host"] = testsuite_config["influxdb_host"]
if "influxdb_port" in testsuite_config:
    properties["influxdb_port"] = testsuite_config["influxdb_port"]
if "influxdb_db" in testsuite_config:
    properties["influxdb_db"] = testsuite_config["influxdb_db"]

# TODO: The SlurmConfigDirBase should be passed from the testsuite.conf too,
#       instead of the hadcoded prefix+etc.orig form here.
properties["slurm-config-orig-dir"] = properties["slurm-prefix"] + "/etc.orig"

# Set derived directory properties
# The environment (e.g. PATH, SLURM_CONF) overrides the configuration.
# If the Slurm clients and daemons are not in the current PATH
# but can be found using the configured SlurmInstallDir, add the
# derived bin and sbin dir to the current PATH.
properties["slurm-bin-dir"] = f"{properties['slurm-prefix']}/bin"
if squeue_path := shutil.which("squeue"):
    properties["slurm-bin-dir"] = os.path.dirname(squeue_path)
elif os.access(f"{properties['slurm-bin-dir']}/squeue", os.X_OK):
    os.environ["PATH"] += ":" + properties["slurm-bin-dir"]
properties["slurm-sbin-dir"] = f"{properties['slurm-prefix']}/sbin"
if slurmctld_path := shutil.which("slurmctld"):
    properties["slurm-sbin-dir"] = os.path.dirname(slurmctld_path)
elif os.access(f"{properties['slurm-sbin-dir']}/slurmctld", os.X_OK):
    os.environ["PATH"] += ":" + properties["slurm-sbin-dir"]
properties["slurm-config-dir"] = re.sub(
    r"\${prefix}", properties["slurm-prefix"], properties["slurm-config-dir"]
)
if slurm_conf_path := os.getenv("SLURM_CONF"):
    properties["slurm-config-dir"] = os.path.dirname(slurm_conf_path)

# Derive the slurm-user value
properties["slurm-user"] = "root"
slurm_config_file = f"{properties['slurm-config-dir']}/slurm.conf"
if not os.path.isfile(slurm_config_file):
    pytest.fail(
        f"The python testsuite was expecting your slurm.conf to be found in {properties['slurm-config-dir']}. Please create it or use the SLURM_CONF environment variable to indicate its location."
    )
if os.access(slurm_config_file, os.R_OK):
    with open(slurm_config_file, "r") as f:
        for line in f.readlines():
            if match := re.search(r"^\s*(?i:SlurmUser)\s*=\s*(.*)$", line):
                properties["slurm-user"] = match.group(1)
else:
    # slurm.conf is not readable, we will try reading it as root
    results = run_command(
        f"grep -i SlurmUser {slurm_config_file}", user="root", quiet=True
    )
    if results["exit_code"] == 0:
        pytest.fail(f"Unable to read {slurm_config_file}")
    for line in results["stdout"].splitlines():
        if match := re.search(r"^\s*(?i:SlurmUser)\s*=\s*(.*)$", line):
            properties["slurm-user"] = match.group(1)

properties["submitted-jobs"] = []
if "slurmtestuser" in testsuite_config:
    properties["test-user"] = testsuite_config["slurmtestuser"]
    properties["test-user-set"] = True
else:
    properties["test-user"] = pwd.getpwuid(os.getuid()).pw_name
    properties["test-user-set"] = False

properties["test-user-uid"] = pwd.getpwnam(properties["test-user"]).pw_uid
properties["test-user-gid"] = pwd.getpwnam(properties["test-user"]).pw_gid

properties["auto-config"] = False
properties["allow-slurmdbd-modify"] = False
properties["slurmrestd-started"] = False

# Instantiate a nodes dictionary. These are populated in require_slurm_running.
nodes = {}

# Check if user has sudo privileges
results = subprocess.run(
    "sudo -ln | grep -q '(ALL.*) NOPASSWD: ALL'",
    shell=True,
    capture_output=True,
    text=True,
)
if results.returncode == 0:
    properties["sudo-rights"] = True
else:
    properties["sudo-rights"] = False