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


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


def test_quit_on_interrupt():
    """Test srun handling of SIGINT to get task status or kill the job (--quit-on-interrupt option)."""

    file_in = atf.module_tmp_path / "file_in"
    atf.make_bash_script(
        file_in,
        """trap \"\" INT
echo WAITING
sleep 1000""",
    )

    child = pexpect.spawn(f"srun -v -N1 --unbuffered {file_in}")
    child.expect(r"launching StepId=(\d+)")
    step_id = re.search(r"\d+", str(child.after)).group(0)
    child.expect("WAITING")
    time.sleep(1)
    atf.run_command(f"kill -INT {child.pid}")
    assert (
        child.expect(["srun: interrupt", pexpect.EOF, pexpect.TIMEOUT]) == 0
    ), f"srun failed to process command (kill -INT {child.pid})"
    atf.run_command(f"scancel {step_id}")
    assert (
        child.expect(["Job step aborted", pexpect.EOF, pexpect.TIMEOUT]) == 0
    ), f"srun failed to process command (scancel {step_id})"
    child.close()

    child = pexpect.spawn(f"srun -v -N1 --unbuffered --quit-on-interrupt {file_in}")
    child.expect(r"launching StepId=(\d+)")
    step_id = re.search(r"\d+", str(child.after)).group(0)
    child.expect("WAITING")
    time.sleep(1)
    atf.run_command(f"kill -INT {child.pid}")
    assert (
        child.expect(["task 0: Killed", pexpect.EOF, pexpect.TIMEOUT]) == 0
    ), "Did not get message (task 0: Killed) from srun"