############################################################################
# Copyright (C) SchedMD LLC.
############################################################################
import atf
import pytest
import pathlib
import re
import signal
import time
import logging


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


def test_signal_forwarding():
    """Test of srun signal forwarding"""

    sig1_count = 2
    sig2_count = 2

    file_in = atf.module_tmp_path / "file_in"
    file_out = atf.module_tmp_path / "file_out"
    script_file = (
        pathlib.Path(atf.properties["testsuite_python_lib"])
        / "signal_forwarding_script.py"
    )
    atf.make_bash_script(file_in, f"""srun --output={file_out} python3 {script_file}""")
    job_id = atf.submit_job_sbatch(f"{file_in}")

    for i in range(sig1_count):
        atf.run_command(f"scancel --signal {signal.SIGUSR1.value} {job_id}")
        time.sleep(1)
    for i in range(sig2_count):
        atf.run_command(f"scancel --signal {signal.SIGUSR2.value} {job_id}")
        time.sleep(1)
    atf.wait_for_job_state(job_id, "DONE", timeout=60)

    received_all_sigs = False
    user1_count = 0
    user2_count = 0
    with open(file_out, "r") as fp:
        for line in fp:
            if re.search(str(signal.SIGUSR1.value), line) is not None:
                user1_count += 1
            if re.search(str(signal.SIGUSR2.value), line) is not None:
                user2_count += 1
            if user1_count == sig1_count and user2_count == sig2_count:
                received_all_sigs = True
                break
    assert (
        received_all_sigs == True
    ), f"Did not receive all signals. Signal {signal.SIGUSR1.value} count: {user1_count}, Signal {signal.SIGUSR2.value} count: {user2_count}"
