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


# Setup
@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_accounting()
    atf.require_nodes(1, [("CPUs", 4), ("RealMemory", 40)])
    atf.require_slurm_running()


def test_hetjob(tmp_path):
    file_in = str(tmp_path / "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

$bin_sleep 300""",
    )
    leader_job_id = atf.submit_job_sbatch(f"-t1 {file_in}", 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"
