Refactored tests

Author: aalitaiga

smartdispatch/tests/pbs_slurm_test.py

@@ -12,7 +12,6 @@
 #PBS -o arrayJob_%A_%a.out
 #PBS -t 1-2
 #PBS -l walltime=01:00:00
-#PBS -l naccelerators=1
 {}
 
 ######################
@@ -43,90 +42,61 @@
 nvidia-smi
 """
 
+def test_param(self, param_array, command, flag, string=pbs_string):
+    for param in param_array:
+        command = pbs_string.format(
+            string.format(command.format(param))
+        )
+        with open("test.pbs", "w") as text_file:
+            text_file.write(command)
+        process = Popen("sbatch test.pbs", stdout=PIPE, stderr=PIPE, shell=True)
+        stdout, stderr = process.communicate()
+        assert_true("Submitted batch job" in stdout)
+        job_id = stdout.split(" ")[-1].strip()
+
+        time.sleep(0.25)
+        process = Popen("squeue -u $USER -j {} -O {}".format(job_id, flag), stdout=PIPE, stderr=PIPE, shell=True)
+        stdout, stderr = process.communicate()
+        job_params = [c.strip() for c in stdout.split("\n")[1:] if c != '']
+        assert_true(all(p == param for p in job_params))
 
 class TestSlurm(unittest.TestCase):
 
+    def tearDown(self):
+        process = Popen("rm *.out", stdout=PIPE, stderr=PIPE, shell=True)
+        stdout, stderr = process.communicate()
+
     def test_priority(self):
-        priorities = ['high', 'low']
-        for priority in priorities:
-            string = pbs_string.format(
-                "#SBATCH --qos={priority}".format(priority=priority)
-            )
-            with open("test.pbs", "w") as text_file:
-                text_file.write(string)
-            process = Popen("sbatch test.pbs", stdout=PIPE, stderr=PIPE, shell=True)
-            stdout, stderr = process.communicate()
-            assert_true("Submitted batch job" in stdout)
-            job_id = stdout.split(" ")[-1]
-
-            time.sleep(0.25)
-            process = Popen("squeue -u $USER -j {} -O qos".format(job_id), stdout=PIPE, stderr=PIPE, shell=True)
-            stdout, stderr = process.communicate()
-            job_priorities = [prio.strip() for prio in stdout.split("\n")[1:] if prio != '']
-            assert_true(all(pri == priority for pri in job_priorities))
+        test_param(
+            ['high', 'low'],
+            "#SBATCH --qos={}",
+            "qos",
+            pbs_string
+        )
 
     def test_gres(self):
-        gress = ['gpu', 'gpu:titanblack']
-        for gres in gress:
-            string = pbs_string.format(
-                "#PBS -l gpu={gres}".format(gres=gres)
-            )
-            with open("test.pbs", "w") as text_file:
-                text_file.write(string)
-            process = Popen("sbatch test.pbs", stdout=PIPE, stderr=PIPE, shell=True)
-            stdout, stderr = process.communicate()
-            assert_true("Submitted batch job" in stdout)
-            job_id = stdout.split(" ")[-1]
-
-            time.sleep(0.25)
-            process = Popen("squeue -u $USER -j {} -O gres".format(job_id), stdout=PIPE, stderr=PIPE, shell=True)
-            stdout, stderr = process.communicate()
-            job_gres = [gre.strip() for gre in stdout.split("\n")[1:] if gre != '']
-            assert_true(all(gpu == gres for gpu in job_gres))
+        test_param(
+            ['titanblack],
+            "#PBS -l naccelerators={}",
+            "gres",
+            pbs_string
+        )
 
     def test_memory(self):
-        mems = ['2G', '4G']
-        for mem in mems:
-            string = pbs_string.format(
-                "#PBS -l mem_free={memory}".format(memory=mem)
-            )
-            with open("test.pbs", "w") as text_file:
-                text_file.write(string)
-            process = Popen("sbatch test.pbs", stdout=PIPE, stderr=PIPE, shell=True)
-            stdout, stderr = process.communicate()
-            assert_true("Submitted batch job" in stdout)
-            job_id = stdout.split(" ")[-1]
-
-            time.sleep(0.25)
-            process = Popen("squeue -u $USER -j {} -O minmemory".format(job_id), stdout=PIPE, stderr=PIPE, shell=True)
-            stdout, stderr = process.communicate()
-            job_mems = [m.strip() for m in stdout.split("\n")[1:] if m != '']
-            assert_true(all(me == mem for me in job_mems))
-
-    # def test_pbs_slurm(self):
-    #     priorities = ['unkillable', 'high', 'low']
-    #     # doesn't test 12 and 24 to avoid killing jobs just for testing
-    #     constraints = ['gpulowmem', 'gpu6gb', 'gpu8gb']
-    #     mem = ['2G', '4G', '6G']
-    #     mail = 'adrien.alitaiga@gmail.com'
-    #     # gpus = ['titanblack']
-    #
-    #     for priority in priorities:
-    #         for cons in constraints:
-    #             for m in mem:
-    #                 string = pbs_string.format(**{
-    #                     'cons': cons,
-    #                     'mail': mail,
-    #                     'priority': priority
-    #                 })
-    #                 with open("test.pbs", "w") as text_file:
-    #                     text_file.write(string)
-    #                 process = Popen("sbatch test.pbs", stdout=PIPE, stderr=PIPE, shell=True)
-    #                 time.sleep(1)
-    #
-    #                 stdout, stderr = process.communicate()
-    #                 print stdout
-    #                 assert_true("Submitted batch job" in stdout)
+        test_param(
+            ['2G', '4G'],
+            "#PBS -l mem={}",
+            "minmemory",
+            pbs_string
+        )
+
+    def test_nb_cpus(self):
+        test_param(
+            ["1", "2"],
+            "PBS -l ncpus={}",
+            "cpuspertask", # or numcpus
+            pbs_string
+        )
 
 
 pbs_string2 = """\