diff --git a/pysqa/cmd.py b/pysqa/cmd.py index 4a1c589f..b8f6bb10 100644 --- a/pysqa/cmd.py +++ b/pysqa/cmd.py @@ -3,14 +3,16 @@ import json import getopt from pysqa.queueadapter import QueueAdapter +from pysqa.utils.execute import execute_command -def command_line(argv): +def command_line(argv, execute_command=execute_command): """ Parse the command line arguments. Args: argv: Command line arguments + execute_command: function to comunicate with shell process """ directory = "~/.queues" @@ -94,7 +96,7 @@ def command_line(argv): print("cmd.py help ... coming soon.") sys.exit() if mode_submit or mode_delete or mode_reservation or mode_status: - qa = QueueAdapter(directory=directory) + qa = QueueAdapter(directory=directory, execute_command=execute_command) if mode_submit: print( qa.submit_job( diff --git a/pysqa/ext/modular.py b/pysqa/ext/modular.py index d1e703a6..02a861fe 100644 --- a/pysqa/ext/modular.py +++ b/pysqa/ext/modular.py @@ -3,11 +3,14 @@ import pandas from pysqa.utils.basic import BasisQueueAdapter +from pysqa.utils.execute import execute_command class ModularQueueAdapter(BasisQueueAdapter): - def __init__(self, config, directory="~/.queues"): - super(ModularQueueAdapter, self).__init__(config=config, directory=directory) + def __init__(self, config, directory="~/.queues", execute_command=execute_command): + super(ModularQueueAdapter, self).__init__( + config=config, directory=directory, execute_command=execute_command + ) self._queue_to_cluster_dict = { k: v["cluster"] for k, v in self._config["queues"].items() } diff --git a/pysqa/ext/remote.py b/pysqa/ext/remote.py index 6e3e731a..ef9e7b33 100644 --- a/pysqa/ext/remote.py +++ b/pysqa/ext/remote.py @@ -8,11 +8,14 @@ import warnings from tqdm import tqdm from pysqa.utils.basic import BasisQueueAdapter +from pysqa.utils.execute import execute_command class RemoteQueueAdapter(BasisQueueAdapter): - def __init__(self, config, directory="~/.queues"): - super(RemoteQueueAdapter, self).__init__(config=config, directory=directory) + def __init__(self, config, directory="~/.queues", execute_command=execute_command): + super(RemoteQueueAdapter, self).__init__( + config=config, directory=directory, execute_command=execute_command + ) self._ssh_host = config["ssh_host"] self._ssh_username = config["ssh_username"] self._ssh_known_hosts = os.path.abspath( diff --git a/pysqa/queueadapter.py b/pysqa/queueadapter.py index 114d7d36..95a02dfb 100644 --- a/pysqa/queueadapter.py +++ b/pysqa/queueadapter.py @@ -6,6 +6,7 @@ from pysqa.ext.modular import ModularQueueAdapter from pysqa.ext.remote import RemoteQueueAdapter from pysqa.utils.config import read_config +from pysqa.utils.execute import execute_command __author__ = "Jan Janssen" __copyright__ = "Copyright 2019, Jan Janssen" @@ -44,14 +45,16 @@ class QueueAdapter(object): Queues available for auto completion QueueAdapter().queues. returns the queue name. """ - def __init__(self, directory="~/.queues"): + def __init__(self, directory="~/.queues", execute_command=execute_command): queue_yaml = os.path.join(directory, "queue.yaml") clusters_yaml = os.path.join(directory, "clusters.yaml") self._adapter = None if os.path.exists(queue_yaml): self._queue_dict = { "default": set_queue_adapter( - config=read_config(file_name=queue_yaml), directory=directory + config=read_config(file_name=queue_yaml), + directory=directory, + execute_command=execute_command, ) } primary_queue = "default" @@ -61,6 +64,7 @@ def __init__(self, directory="~/.queues"): k: set_queue_adapter( config=read_config(file_name=os.path.join(directory, v)), directory=directory, + execute_command=execute_command, ) for k, v in config["cluster"].items() } @@ -296,7 +300,7 @@ def check_queue_parameters( ) -def set_queue_adapter(config, directory): +def set_queue_adapter(config, directory, execute_command=execute_command): """ Initialize the queue adapter @@ -305,10 +309,16 @@ def set_queue_adapter(config, directory): directory (str): directory which contains the queue configurations """ if config["queue_type"] in ["SGE", "TORQUE", "SLURM", "LSF", "MOAB"]: - return BasisQueueAdapter(config=config, directory=directory) + return BasisQueueAdapter( + config=config, directory=directory, execute_command=execute_command + ) elif config["queue_type"] in ["GENT"]: - return ModularQueueAdapter(config=config, directory=directory) + return ModularQueueAdapter( + config=config, directory=directory, execute_command=execute_command + ) elif config["queue_type"] in ["REMOTE"]: - return RemoteQueueAdapter(config=config, directory=directory) + return RemoteQueueAdapter( + config=config, directory=directory, execute_command=execute_command + ) else: raise ValueError diff --git a/tests/test_cmd.py b/tests/test_cmd.py new file mode 100644 index 00000000..5d53f530 --- /dev/null +++ b/tests/test_cmd.py @@ -0,0 +1,57 @@ +import os +import unittest +from pysqa.cmd import command_line + + +class TestCMD(unittest.TestCase): + def test_help(self): + with self.assertRaises(SystemExit): + command_line(["--help"]) + + def test_wrong_option(self): + with self.assertRaises(SystemExit): + command_line(["--error"]) + + def test_submit(self): + def execute_command( + commands, + working_directory=None, + split_output=True, + shell=False, + error_filename="pysqa.err", + ): + return "1\n" + + test_dir = os.path.abspath(os.path.dirname(__file__)) + with self.assertRaises(SystemExit): + command_line( + [ + "--config_directory", os.path.join(test_dir, "config", "slurm"), + "--submit", + "--queue", "slurm", + "--job_name", "test", + "--working_directory", ".", + "--cores", "2", + "--memory", "1GB", + "--run_time", "10", + "--command", "echo hello" + ], + execute_command=execute_command + ) + with open("run_queue.sh") as f: + output = f.readlines() + content = [ + '#!/bin/bash\n', + '#SBATCH --output=time.out\n', + '#SBATCH --job-name=test\n', + '#SBATCH --chdir=.\n', + '#SBATCH --get-user-env=L\n', + '#SBATCH --partition=slurm\n', + '#SBATCH --time=4320\n', + '#SBATCH --mem=1GBG\n', + '#SBATCH --cpus-per-task=10\n', + '\n', + 'echo hello' + ] + self.assertEqual(output, content) + os.remove("run_queue.sh")