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

# import re


# Setup
@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_auto_config("wants to add custom generic resources")
    atf.require_config_parameter("GresTypes", "r1")
    atf.require_nodes(1, [("Gres", "r1:1")])
    atf.require_slurm_running()


def test_job_array_with_gres():
    """Test creating job array with gres requested"""

    output_pattern = f"{atf.module_tmp_path}/%A-%a.out"
    job_id = atf.submit_job_sbatch(
        f"--array=1-2 --gres=r1:1 --wrap='echo DONE' \
                    --output={output_pattern}"
    )
    output_file_1 = f"{atf.module_tmp_path}/{job_id}-1.out"
    output_file_2 = f"{atf.module_tmp_path}/{job_id}-2.out"
    atf.wait_for_job_state(job_id, "DONE", timeout=15, fatal=True)
    with open(output_file_1, "r") as f:
        output = f.read()
        assert "DONE" in output, "Expect job to finish"
    with open(output_file_2, "r") as f:
        output = f.read()
        assert "DONE" in output, "Expect job to finish"
    assert atf.is_slurmctld_running()