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

num_nodes = 3


# Setup
@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_nodes(num_nodes * 2)
    atf.require_slurm_running()


def test_preserve_env():
    """Test that a job correctly uses the -E or --preserve-env flag."""

    file_in = atf.module_tmp_path / 'file_in'
    file_out = atf.module_tmp_path / 'file_out'
    num_tasks = num_nodes * 2
    srun_nodes = 1
    srun_tasks = 1
    atf.make_bash_script(file_in, f"""printenv SLURM_NNODES
srun -E -n{srun_tasks} -N{srun_nodes} printenv SLURM_NNODES
printenv SLURM_NTASKS
srun --preserve-env -n{srun_tasks} -N{srun_nodes} printenv SLURM_NTASKS
srun -n{srun_tasks} -N{srun_nodes} printenv SLURM_NNODES
srun -n{srun_tasks} -N{srun_nodes} printenv SLURM_NTASKS""")

    atf.run_command("srun printenv SLURM_NNODES")
    job_id = atf.submit_job_sbatch(f"-O --output={file_out} -N{num_nodes} -n{num_tasks} {file_in}")
    atf.wait_for_job_state(job_id, 'DONE')
    atf.wait_for_file(file_out)
    f = open(file_out, 'r')
    error_msg = "Incorrect output from:"
    assert num_nodes == int(f.readline()), f"{error_msg} printenv SLURM_NNODES"
    assert num_nodes == int(f.readline()), f"{error_msg} srun -E -n{srun_tasks} -N{srun_nodes} printenv SLURM_NNODES"
    assert num_tasks == int(f.readline()), f"{error_msg} printenv SLURM_NTASKS"
    assert num_tasks == int(f.readline()), f"{error_msg} srun --preserve-env -n{srun_tasks} -N{srun_nodes} printenv SLURM_NTASKS"
    assert srun_nodes == int(f.readline()), f"{error_msg} srun -n{srun_tasks} -N{srun_nodes} printenv SLURM_NNODES"
    assert srun_tasks == int(f.readline()), f"{error_msg} srun -n{srun_tasks} -N{srun_nodes} printenv SLURM_NTASKS"
