############################################################################
# Copyright (C) SchedMD LLC.
############################################################################
import atf
import pytest
import re

node_num = 4


# Setup
@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_nodes(node_num)
    atf.require_slurm_running()


def test_threads():
    """Verify a job executes with various launch thread fanouts with --threads option"""

    node_count = "1-64"
    threads = 0
    output = atf.run_command_output(
        f"srun -N{node_count} -l --threads=1 printenv SLURMD_NODENAME"
    )
    assert re.search(r"0: (\S+)", output) is not None, "Did not get hostname of task 0"

    task_num = int(atf.get_config_parameter("MaxTasksPerNode"))
    if task_num <= 32:
        threads = task_num
    else:
        threads = 32
    task_num = threads * 2
    output = atf.run_command(
        f"srun -N{node_count} -n{task_num} -O -l --threads={threads} printenv SLURMD_NODENAME"
    )
    assert (
        output["exit_code"] == 0
    ), f"srun failed to run with {threads} threads and {task_num} tasks"
    assert (
        re.search(r"0: (\S+)", output["stdout"]) is not None
    ), f"Did not get hostname of task 0 when running {threads} threads"