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


@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_accounting()
    atf.require_nodes(3, [("CPUs", 8), ("RealMemory", 80)])

    # only to make the test faster
    atf.require_config_parameter("SchedulerParameters", "bf_interval=1")

    atf.require_slurm_running()


def test_hetjob():
    file_in = "hetjob.in"
    atf.make_bash_script(
        file_in,
        """
#SBATCH --cpus-per-task=4 --mem-per-cpu=10 --ntasks=1
#SBATCH hetjob
#SBATCH --cpus-per-task=2 --mem-per-cpu=2  --ntasks=1 -t1
#SBATCH hetjob
#SBATCH --cpus-per-task=1 --mem-per-cpu=6  --ntasks=1 -t1

srun sleep 60
        """,
    )
    leader_job_id = atf.submit_job_sbatch(f"-t1 {file_in}", fatal=True)
    atf.wait_for_job_state(leader_job_id, "RUNNING", fatal=True)

    jobs_dict = atf.get_jobs()

    # Verify details about leader job
    assert leader_job_id in jobs_dict
    leader_job_dict = jobs_dict[leader_job_id]
    assert leader_job_dict["JobId"] == leader_job_id
    assert leader_job_dict["HetJobId"] == leader_job_id
    assert leader_job_dict["HetJobOffset"] == 0
    het_job_expression = leader_job_dict["HetJobIdSet"]
    assert re.search(r"[\d,-]+", het_job_expression) is not None
    assert leader_job_dict["CPUs/Task"] == 4
    assert leader_job_dict["MinMemoryCPU"] == "10M"
    het_job_component_list = atf.range_to_list(het_job_expression)
    het_job_component_list.remove(leader_job_id)
    assert len(het_job_component_list) == 2

    # Verify details about first component job
    component1_job_id = het_job_component_list[0]
    assert component1_job_id in jobs_dict
    component1_job_dict = jobs_dict[component1_job_id]
    assert component1_job_dict["JobId"] == component1_job_id
    assert component1_job_dict["HetJobId"] == leader_job_id
    assert component1_job_dict["HetJobOffset"] == 1
    assert re.search(r"[\d,-]+", component1_job_dict["HetJobIdSet"]) is not None
    assert component1_job_dict["CPUs/Task"] == 2
    assert component1_job_dict["MinMemoryCPU"] == "2M"

    # Verify details about second component job
    component2_job_id = het_job_component_list[1]
    assert component2_job_id in jobs_dict
    component2_job_dict = jobs_dict[component2_job_id]
    assert component2_job_dict["JobId"] == component2_job_id
    assert component2_job_dict["HetJobId"] == leader_job_id
    assert component2_job_dict["HetJobOffset"] == 2
    assert re.search(r"[\d,-]+", component2_job_dict["HetJobIdSet"]) is not None
    assert component2_job_dict["CPUs/Task"] == 1
    assert component2_job_dict["MinMemoryCPU"] == "6M"