Add tests for req_keepvars and req_keepvars_extra

jupyterhub · Jun 21, 2018 · 8eb2fd6 · 8eb2fd6
1 parent 72ac758
commit 8eb2fd6
Showing 1 changed file with 37 additions and 3 deletions.
diff --git a/batchspawner/tests/test_spawners.py b/batchspawner/tests/test_spawners.py
@@ -339,16 +339,25 @@ def test_slurm(db, io_loop):
         re.compile(r'^#SBATCH \s+ --time=3-05:10:10', re.X|re.M),
         re.compile(r'^#SBATCH \s+ some_option_asdf', re.X|re.M),
         ]
-    script = [
+    from .. import SlurmSpawner
+    run_spawner_script(db, io_loop, SlurmSpawner, normal_slurm_script,
+                       batch_script_re_list=batch_script_re_list,
+                       spawner_kwargs=spawner_kwargs)
+normal_slurm_script = [
         (re.compile(r'sudo.*sbatch'),   str(testjob)),
         (re.compile(r'sudo.*squeue'),   'PENDING '),          # pending
         (re.compile(r'sudo.*squeue'),   'RUNNING '+testhost), # running
         (re.compile(r'sudo.*squeue'),   'RUNNING '+testhost),
         (re.compile(r'sudo.*scancel'),  'STOP'),
         (re.compile(r'sudo.*squeue'),   ''),
         ]
-    from .. import SlurmSpawner
-    run_spawner_script(db, io_loop, SlurmSpawner, script,
+from .. import SlurmSpawner
+def run_typical_slurm_spawner(db, io_loop,
+        spawner=SlurmSpawner,
+        script=normal_slurm_script,
+        batch_script_re_list=None,
+        spawner_kwargs={}):
+    return run_spawner_script(db, io_loop, spawner, script,
                        batch_script_re_list=batch_script_re_list,
                        spawner_kwargs=spawner_kwargs)
 
@@ -424,3 +433,28 @@ def test_lfs(db, io_loop):
     run_spawner_script(db, io_loop, LsfSpawner, script,
                        batch_script_re_list=batch_script_re_list,
                        spawner_kwargs=spawner_kwargs)
+
+
+def test_keepvars(db, io_loop):
+    # req_keepvars
+    spawner_kwargs = {
+        'req_keepvars': 'ABCDE',
+        }
+    batch_script_re_list = [
+        re.compile(r'--export=ABCDE', re.X|re.M),
+        ]
+    run_typical_slurm_spawner(db, io_loop,
+                              spawner_kwargs=spawner_kwargs,
+                              batch_script_re_list=batch_script_re_list)
+
+    # req_keepvars AND req_keepvars together
+    spawner_kwargs = {
+        'req_keepvars': 'ABCDE',
+        'req_keepvars_extra': 'XYZ',
+        }
+    batch_script_re_list = [
+        re.compile(r'--export=ABCDE,XYZ', re.X|re.M),
+        ]
+    run_typical_slurm_spawner(db, io_loop,
+                              spawner_kwargs=spawner_kwargs,
+                              batch_script_re_list=batch_script_re_list)