# Path: pyjob/slurm.py
# class SlurmTask(ClusterTask):
#     """Slurm executable :obj:`~pyjob.task.Task`"""
# 
#     JOB_ARRAY_INDEX = "$SLURM_ARRAY_TASK_ID"
#     SCRIPT_DIRECTIVE = "#SBATCH"
# 
#     @property
#     def info(self):
#         """:obj:`~pyjob.slurm.SlurmTask` information"""
#         if self.pid is None:
#             return {}
#         try:
#             cexec(["squeue", "-j", str(self.pid)])
#         except (PyJobExecutableNotFoundError, Exception):
#             return {}
#         else:
#             return {"job_number": self.pid, "status": "Running"}
# 
#     def kill(self):
#         """Immediately terminate the :obj:`~pyjob.slurm.SlurmTask`"""
#         if self.pid is None:
#             return
#         cexec(["scancel", str(self.pid)])
#         logger.debug("Terminated task: %d", self.pid)
# 
#     def _check_requirements(self):
#         """Check if the requirements for task execution are met"""
#         self._ensure_exec_available("squeue")
# 
#     def _run(self):
#         """Method to initialise :obj:`~pyjob.slurm.SlurmTask` execution"""
#         self.runscript = self._create_runscript()
#         self.runscript.write()
#         stdout = cexec(["sbatch", self.runscript.path], cwd=self.directory)
#         self.pid = int(stdout.strip().split()[-1])
#         logger.debug(
#             "%s [%d] submission script is %s",
#             self.__class__.__qualname__,
#             self.pid,
#             self.runscript.path,
#         )
# 
#     def _create_runscript(self):
#         """Utility method to create runscript"""
#         runscript = Script(
#             directory=self.directory,
#             prefix="slurm_",
#             suffix=".script",
#             stem=str(uuid.uuid1().int),
#         )
#         runscript.append(self.__class__.SCRIPT_DIRECTIVE + " --export=ALL")
#         runscript.append(self.__class__.SCRIPT_DIRECTIVE + f" --job-name={self.name}")
#         if self.dependency:
#             cmd = f'--depend=afterok:{":".join(map(str, self.dependency))}'
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + " " + cmd)
#         if self.queue:
#             cmd = f"-p {self.queue}"
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + " " + cmd)
#         if self.nprocesses:
#             cmd = f"-n {self.nprocesses}"
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + " " + cmd)
#         if self.directory:
#             cmd = f"--workdir={self.directory}"
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + " " + cmd)
#         if self.runtime:
#             cmd = f"-t {self.runtime}"
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + " " + cmd)
#         if self.extra:
#             cmd = " ".join(map(str, self.extra))
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + " " + cmd)
#         if len(self.script) > 1:
#             logf = runscript.path.replace(".script", ".log")
#             jobsf = runscript.path.replace(".script", ".jobs")
#             with open(jobsf, "w") as f_out:
#                 f_out.write("\n".join(self.script))
#             cmd = f"--array=1-{len(self.script)}%{self.max_array_size}"
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + " " + cmd)
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + f" -o {logf}")
#             runscript.extend(self.get_array_bash_extension(jobsf, 0))
#         else:
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + f" -o {self.log[0]}")
#             runscript.append(self.script[0])
#         return runscript


import os
import pytest
from unittest import mock
from pyjob.slurm import SlurmTask



@pytest.mark.skipif(pytest.on_windows, reason="Unavailable on Windows")
@mock.patch("pyjob.slurm.SlurmTask._check_requirements")
class TestCreateRunscript(object):
    def test_1(self, check_requirements_mock):
        check_requirements_mock.return_value = None
        scripts = [pytest.helpers.get_py_script(i, 1) for i in range(1)]
        [s.write() for s in scripts]
        paths = [s.path for s in scripts]
        task = SlurmTask(paths)
        runscript = task._create_runscript()
        pytest.helpers.unlink(paths)
        assert runscript.shebang == "#!/bin/bash"
        assert runscript.content == [
            "#SBATCH --export=ALL",
            "#SBATCH --job-name=pyjob",
            "#SBATCH -n 1",
            "#SBATCH --workdir=" + os.getcwd(),
            "#SBATCH -o " + paths[0].replace(".py", ".log"),
            paths[0],
        ]

    def test_2(self, check_requirements_mock):
        check_requirements_mock.return_value = None
        scripts = [pytest.helpers.get_py_script(i, 1) for i in range(3)]
        [s.write() for s in scripts]
        paths = [s.path for s in scripts]

# Path: pyjob/slurm.py
# class SlurmTask(ClusterTask):
#     """Slurm executable :obj:`~pyjob.task.Task`"""
# 
#     JOB_ARRAY_INDEX = "$SLURM_ARRAY_TASK_ID"
#     SCRIPT_DIRECTIVE = "#SBATCH"
# 
#     @property
#     def info(self):
#         """:obj:`~pyjob.slurm.SlurmTask` information"""
#         if self.pid is None:
#             return {}
#         try:
#             cexec(["squeue", "-j", str(self.pid)])
#         except (PyJobExecutableNotFoundError, Exception):
#             return {}
#         else:
#             return {"job_number": self.pid, "status": "Running"}
# 
#     def kill(self):
#         """Immediately terminate the :obj:`~pyjob.slurm.SlurmTask`"""
#         if self.pid is None:
#             return
#         cexec(["scancel", str(self.pid)])
#         logger.debug("Terminated task: %d", self.pid)
# 
#     def _check_requirements(self):
#         """Check if the requirements for task execution are met"""
#         self._ensure_exec_available("squeue")
# 
#     def _run(self):
#         """Method to initialise :obj:`~pyjob.slurm.SlurmTask` execution"""
#         self.runscript = self._create_runscript()
#         self.runscript.write()
#         stdout = cexec(["sbatch", self.runscript.path], cwd=self.directory)
#         self.pid = int(stdout.strip().split()[-1])
#         logger.debug(
#             "%s [%d] submission script is %s",
#             self.__class__.__qualname__,
#             self.pid,
#             self.runscript.path,
#         )
# 
#     def _create_runscript(self):
#         """Utility method to create runscript"""
#         runscript = Script(
#             directory=self.directory,
#             prefix="slurm_",
#             suffix=".script",
#             stem=str(uuid.uuid1().int),
#         )
#         runscript.append(self.__class__.SCRIPT_DIRECTIVE + " --export=ALL")
#         runscript.append(self.__class__.SCRIPT_DIRECTIVE + f" --job-name={self.name}")
#         if self.dependency:
#             cmd = f'--depend=afterok:{":".join(map(str, self.dependency))}'
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + " " + cmd)
#         if self.queue:
#             cmd = f"-p {self.queue}"
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + " " + cmd)
#         if self.nprocesses:
#             cmd = f"-n {self.nprocesses}"
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + " " + cmd)
#         if self.directory:
#             cmd = f"--workdir={self.directory}"
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + " " + cmd)
#         if self.runtime:
#             cmd = f"-t {self.runtime}"
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + " " + cmd)
#         if self.extra:
#             cmd = " ".join(map(str, self.extra))
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + " " + cmd)
#         if len(self.script) > 1:
#             logf = runscript.path.replace(".script", ".log")
#             jobsf = runscript.path.replace(".script", ".jobs")
#             with open(jobsf, "w") as f_out:
#                 f_out.write("\n".join(self.script))
#             cmd = f"--array=1-{len(self.script)}%{self.max_array_size}"
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + " " + cmd)
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + f" -o {logf}")
#             runscript.extend(self.get_array_bash_extension(jobsf, 0))
#         else:
#             runscript.append(self.__class__.SCRIPT_DIRECTIVE + f" -o {self.log[0]}")
#             runscript.append(self.script[0])
#         return runscript


# Path: pyjob/tests/test_slurm.py
import os
import pytest
from unittest import mock
from pyjob.slurm import SlurmTask




@pytest.mark.skipif(pytest.on_windows, reason="Unavailable on Windows")
@mock.patch("pyjob.slurm.SlurmTask._check_requirements")
class TestCreateRunscript(object):
    def test_1(self, check_requirements_mock):
        check_requirements_mock.return_value = None
        scripts = [pytest.helpers.get_py_script(i, 1) for i in range(1)]
        [s.write() for s in scripts]
        paths = [s.path for s in scripts]
        task = SlurmTask(paths)
        runscript = task._create_runscript()
        pytest.helpers.unlink(paths)
        assert runscript.shebang == "#!/bin/bash"
        assert runscript.content == [
            "#SBATCH --export=ALL",
            "#SBATCH --job-name=pyjob",
            "#SBATCH -n 1",
            "#SBATCH --workdir=" + os.getcwd(),
            "#SBATCH -o " + paths[0].replace(".py", ".log"),
            paths[0],
        ]

    def test_2(self, check_requirements_mock):
        check_requirements_mock.return_value = None
        scripts = [pytest.helpers.get_py_script(i, 1) for i in range(3)]
        [s.write() for s in scripts]
        paths = [s.path for s in scripts]

        task = SlurmTask(paths)