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

# TODO: Bug 17619
#       From the docs:
#       "A single srun opens 3 listening ports plus 2 more for every 48 hosts."
#       It seems that docs are not right, but (4 + 2 * ((hosts-1)//48))

port_range = 8
srun_port_lower = 60000
srun_port_upper = srun_port_lower + port_range


# Setup
@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_auto_config("wants to edit SrunPortRange and create a node")
    atf.require_config_parameter(
        "SrunPortRange", f"{srun_port_lower}-{srun_port_upper}"
    )
    atf.require_nodes(145, [("CPUs", 2), ("RealMemory", 2)])
    atf.require_slurm_running()


@pytest.mark.parametrize("nodes", [1, 10, 48, 49, 96, 100, 144])
def test_srun_ports_in_range(nodes):
    """Test srun uses the right SrunPortRange"""

    command = '''bash -c "
        task_id=\\$(scontrol show step \\$SLURM_JOBID.\\$SLURM_STEPID \
        | grep SrunHost | awk -F: '{print \\$3}')
        lsof -P -p \\$task_id 2>/dev/null | grep LISTEN | awk '{print \\$9}' \
        | awk -F: '{print \\$2}'"'''
    output = atf.run_command_output(f"srun -N{nodes} {command}").split("\n")
    count = 0
    for port_string in output:
        # Ignore blank lines
        if len(port_string) < 4:
            continue
        count += 1
        port_int = int(port_string)
        assert (
            port_int >= srun_port_lower and port_int <= srun_port_upper
        ), f"Port {port_int} is not in range {srun_port_lower}-{srun_port_upper}"

    ports = nodes * (4 + 2 * ((nodes - 1) // 48))
    assert count == ports, f"srun with -N{nodes} should use {ports} ports, not {count}"


@pytest.mark.parametrize("nodes", [145])
def test_srun_ports_out_of_range(nodes):
    """Test sruns with too many nodes, so with not enough SrunPortRange"""

    result = atf.run_command(f"srun -t1 -N{nodes} sleep 1", xfail=True)
    assert (
        result["exit_code"] != 0
    ), f"srun with -N{nodes} should fail because it needs more than {port_range} ports"

    regex = rf"all ports in range .{srun_port_lower}, {srun_port_upper}. exhausted"
    assert (
        re.search(regex, result["stderr"]) is not None
    ), "srun's stderr should contain the 'all ports in range exhausted' message"


def test_out_of_srun_ports():
    """Test exhausted ports"""

    job_id1 = atf.submit_job_sbatch('-N1 -o/dev/null --wrap="srun sleep 30"')
    job_id2 = atf.submit_job_sbatch('-N1 -o/dev/null --wrap="srun sleep 30"')

    atf.wait_for_step(job_id1, 0, fatal=True)
    atf.wait_for_step(job_id2, 0, fatal=True)

    result = atf.run_command("srun -t1 -N1 sleep 1", xfail=True)
    assert (
        result["exit_code"] != 0
    ), "srun should fail because running job shoulb be using the whole SrunPortRange"

    regex = rf"all ports in range .{srun_port_lower}, {srun_port_upper}. exhausted"
    assert (
        re.search(regex, result["stderr"]) is not None
    ), "srun's stderr should contain the 'all ports in range exhausted' message"