############################################################################
# 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 gres')
    atf.require_config_parameter('SelectType', 'select/cons_tres')
    atf.require_config_parameter('GresTypes', 'r1,r2')
    atf.require_nodes(1,
            [('Gres', 'r1:no_consume:1,r2:1'), ('CPUs', 2), ('RealMemory', 2)])
    atf.require_slurm_running()


def test_no_consume():
    """Test gres with no_consume"""

    # Get the last node only
    no_consume_output = atf.run_command_output(
        'srun --gres=r1 scontrol show node -d').split('NodeName')[-1]
    assert re.search('GresUsed=.*r1:0', no_consume_output) is not None, \
        'Expect no_consume resources to not be consumed'
    consume_output = atf.run_command_output(
        'srun --gres=r2 scontrol show node -d').split('NodeName')[-1]
    assert re.search('GresUsed=.*r2:1', consume_output) is not None, \
        'Expect consumable resource to be consumed'


def test_no_consume_parallel():
    """Test no_consume gres with parallel jobs"""

    job_id_1 = atf.submit_job_sbatch('--gres=r1 --mem=1 --wrap="sleep 20"')
    job_id_2 = atf.submit_job_sbatch('--gres=r1 --mem=1 --wrap="sleep 20"')
    atf.wait_for_job_state(job_id_1, 'RUNNING')
    atf.wait_for_job_state(job_id_2, 'RUNNING')
    squeue = atf.run_command_output('squeue')
    assert re.search(
        f"{job_id_1}( +[^ ]+)( +[^ ]+)( +[^ ]+) +R", squeue) is not None, \
        'Expect first job to be running'
    assert re.search(
        f"{job_id_2}( +[^ ]+)( +[^ ]+)( +[^ ]+) +R", squeue) is not None, \
        'Expect second job to be running'
