############################################################################
# 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 create custom gpu files and custom gres')
    atf.require_config_parameter('SelectType', 'select/cons_tres')
    atf.require_config_parameter('GresTypes', 'gpu')
    atf.require_nodes(1, [('Gres','gpu:2 Sockets=2 CoresPerSocket=1')])
    # GPU's need to point to existing files
    gpu_file = f"{str(atf.module_tmp_path)}/gpu"
    atf.run_command(f"touch {gpu_file + '1'}")
    atf.run_command(f"touch {gpu_file + '2'}")
    atf.require_config_parameter('NodeName',
            f"node2 Name=gpu Cores=0-1 File={gpu_file}[1-2]", source='gres')
    atf.require_slurm_running()


def test_gpu_socket_sharing():
    """Test allocating multiple gpus on the same core group with enforce-binding"""

    output = atf.run_command_output(
            'srun --gres-flags=enforce-binding --ntasks-per-socket=1 \
                    --cpus-per-task=1 --ntasks-per-node=2 -N1 \
                    --gpus-per-task=1 scontrol show nodes node2 -d',
            timeout=2, fatal=True)
    assert re.search('GresUsed=gpu.*:2', output) is not None, \
            'Expect 2 gpus used for job'


def test_gpu_socket_sharing_no_alloc():
    """Test allocating multiple gpus on the same core group with enforce-binding without enough resources"""

    output = atf.run_command(
            'srun --gres-flags=enforce-binding --ntasks-per-socket=1 \
                    --cpus-per-task=2 --ntasks-per-node=2 -N1 \
                    --gpus-per-task=1 scontrol show nodes node2 -d',
            timeout=1, fatal=False)
    assert output['exit_code'] != 0, 'Expect command to timeout'
    assert re.search('srun: job [0-9]+ queued and waiting for resources',
            str(output['stderr'])) is not None, 'Expect command to wait on resources'
