Use method with regexp in class var

dask · Aug 24, 2018 · f72a7f1 · f72a7f1
1 parent 2602192
commit f72a7f1
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 15 deletions.
diff --git a/dask_jobqueue/core.py b/dask_jobqueue/core.py
@@ -139,6 +139,7 @@ class JobQueueCluster(Cluster):
     cancel_command = None
     scheduler_name = ''
     _adaptive_options = {'worker_key': lambda ws: _job_id_from_worker_name(ws.name)}
+    job_id_regexp = r'\d+'
 
     def __init__(self,
                  name=None,
@@ -407,6 +408,12 @@ def _del_pending_jobs(self):
             del self.pending_jobs[job_id]
         return jobs
 
-    @staticmethod
-    def _job_id_from_submit_output(out):
-        return re.findall(r'\d+', out)[0]
+    def _job_id_from_submit_output(self, out):
+        regexp = r'(?P<job_id>{}'.format(self.job_id_regexp)
+        match = re.search(regexp, out)
+        job_id = match.group('job_id')
+        if job_id is None:
+            raise ValueError('Could not parse job id from submission command'
+                             'output. Job id regexp is {}, submission command'
+                             'output is: {}'.format(self.job_id_regexp, out))
+        return job_id
diff --git a/dask_jobqueue/tests/test_jobqueue_core.py b/dask_jobqueue/tests/test_jobqueue_core.py
@@ -46,17 +46,19 @@ def test_forward_ip():
         assert cluster.local_cluster.scheduler.ip == default_ip
 
 
+@pytest.mark.parametrize('Cluster', [PBSCluster, MoabCluster, SLURMCluster,
+                                     SGECluster, LSFCluster, JobQueueCluster])
 @pytest.mark.parametrize(
     'qsub_return_string',
-    ['{jobid}.admin01',
-     'Request {jobid}.asdf was sumbitted to queue: standard.',
-     'sbatch: Submitted batch job {jobid}',
-     '{jobid};cluster',
-     'Job <{jobid}> is submitted to default queue <normal>.',
-     '{jobid}'])
-def test_jobid_from_qsub(qsub_return_string):
-    original_jobid = '654321'
-    qsub_return_string = qsub_return_string.format(jobid=original_jobid)
-    parsed_jobid = JobQueueCluster._job_id_from_submit_output(
-        qsub_return_string.format(jobid=original_jobid))
-    assert parsed_jobid == original_jobid
+    ['{job_id}.admin01',
+     'Request {job_id}.asdf was sumbitted to queue: standard.',
+     'sbatch: Submitted batch job {job_id}',
+     '{job_id};cluster',
+     'Job <{job_id}> is submitted to default queue <normal>.',
+     '{job_id}'])
+def test_job_id_from_qsub(Cluster, qsub_return_string):
+    original_job_id = '654321'
+    qsub_return_string.format(job_id=original_job_id)
+    with Cluster() as cluster:
+        assert (original_job_id
+                == cluster._job_id_from_submit_output(sub_return_string))