admin_environment: test and make active

Author: rkdarst

batchspawner/batchspawner.py

@@ -284,7 +284,7 @@ async def submit_batch_script(self):
         script = await self._get_batch_script(**subvars)
         self.log.info('Spawner submitting job using ' + cmd)
         self.log.info('Spawner submitted script:\n' + script)
-        out = await self.run_command(cmd, input=script, env=self.get_env())
+        out = await self.run_command(cmd, input=script, env=self.get_admin_env())
         try:
             self.log.info('Job submitted. cmd: ' + cmd + ' output: ' + out)
             self.job_id = self.parse_job_id(out)
@@ -310,7 +310,7 @@ async def read_job_state(self):
                         format_template(self.batch_query_cmd, **subvars)))
         self.log.debug('Spawner querying job: ' + cmd)
         try:
-            out = await self.run_command(cmd)
+            out = await self.run_command(cmd, env=self.get_admin_env())
             self.job_status = out
         except Exception as e:
             self.log.error('Error querying job ' + self.job_id)
@@ -328,7 +328,7 @@ async def cancel_batch_job(self):
         cmd = ' '.join((format_template(self.exec_prefix, **subvars),
                         format_template(self.batch_cancel_cmd, **subvars)))
         self.log.info('Cancelling job ' + self.job_id + ': ' + cmd)
-        await self.run_command(cmd)
+        await self.run_command(cmd, env=self.get_admin_env())
 
     def load_state(self, state):
         """load job_id from state"""

batchspawner/tests/test_spawners.py

@@ -1,5 +1,6 @@
 """Test BatchSpawner and subclasses"""
 
+import itertools
 import re
 from unittest import mock
 from .. import BatchSpawnerRegexStates
@@ -29,11 +30,17 @@ class BatchDummy(BatchSpawnerRegexStates):
     state_running_re = Unicode('RUN')
     state_exechost_re = Unicode('RUN (.*)$')
 
-    cmd_expectlist = None
+    cmd_expectlist = None  #List of (re)
     out_expectlist = None
+    env_testlist = None    # List of functions to call on env dict, function should assert within.
     def run_command(self, *args, **kwargs):
         """Overwriten run command to test templating and outputs"""
         cmd = args[0]
+        # Test the environment
+        if self.env_testlist:   # if first item is None, also pop and advance
+            env_test = self.env_testlist.pop(0)
+            if env_test:
+                env_test(kwargs['env'])
         # Test that the command matches the expectations
         if self.cmd_expectlist:
             run_re = self.cmd_expectlist.pop(0)
@@ -221,13 +228,14 @@ def run_command(self, cmd, *args, **kwargs):
     assert status == 1
 
 def run_spawner_script(db, io_loop, spawner, script,
-                       batch_script_re_list=None, spawner_kwargs={}):
+                       batch_script_re_list=None, spawner_kwargs={},
+                       env_test=None):
     """Run a spawner script and test that the output and behavior is as expected.
 
     db: same as in this module
     io_loop: same as in this module
     spawner: the BatchSpawnerBase subclass to test
-    script: list of (input_re_to_match, output)
+    script: list of (input_re_to_match, output, env_testfunc)
     batch_script_re_list: if given, assert batch script matches all of these
     """
     # Create the expected scripts
@@ -238,6 +246,9 @@ def run_spawner_script(db, io_loop, spawner, script,
     class BatchDummyTestScript(spawner):
         @gen.coroutine
         def run_command(self, cmd, input=None, env=None):
+            # Test the environment
+            if env_test:
+                env_test(env)
             # Test the input
             run_re = cmd_expectlist.pop(0)
             if run_re:
@@ -402,15 +413,17 @@ def run_typical_slurm_spawner(db, io_loop,
         spawner=SlurmSpawner,
         script=normal_slurm_script,
         batch_script_re_list=None,
-        spawner_kwargs={}):
+        spawner_kwargs={},
+        env_test=None):
     """Run a full slurm job with default (overrideable) parameters.
 
     This is useful, for example, for changing options and testing effect
     of batch scripts.
     """
     return run_spawner_script(db, io_loop, spawner, script,
                        batch_script_re_list=batch_script_re_list,
-                       spawner_kwargs=spawner_kwargs)
+                       spawner_kwargs=spawner_kwargs,
+                       env_test=env_test)
 
 
 #def test_gridengine(db, io_loop):
@@ -489,16 +502,21 @@ def test_lfs(db, io_loop):
 
 
 def test_keepvars(db, io_loop):
+    """Test of environment handling
+    """
     # req_keepvars
     spawner_kwargs = {
         'req_keepvars_default': 'ABCDE',
         }
     batch_script_re_list = [
         re.compile(r'--export=ABCDE', re.X|re.M),
         ]
+    def env_test(env):
+        assert 'ABCDE' in env
     run_typical_slurm_spawner(db, io_loop,
                               spawner_kwargs=spawner_kwargs,
-                              batch_script_re_list=batch_script_re_list)
+                              batch_script_re_list=batch_script_re_list,
+                              env_test=env_test)
 
     # req_keepvars
     spawner_kwargs = {
@@ -516,11 +534,18 @@ def test_keepvars(db, io_loop):
         'admin_environment': 'ABCDE',
         }
     batch_script_re_list = [
-        re.compile(r'^((?!ABCDE).)*$', re.X|re.S),
+        re.compile(r'^((?!ABCDE).)*$', re.X|re.S),  # ABCDE not in the script
         ]
+    def env_test(env):
+        assert 'ABCDE' in env
+        assert 'VWXYZ' not in env
+    os.environ['ABCDE'] = 'TEST1'
+    os.environ['VWXYZ'] = 'TEST2'
     run_typical_slurm_spawner(db, io_loop,
                               spawner_kwargs=spawner_kwargs,
-                              batch_script_re_list=batch_script_re_list)
+                              batch_script_re_list=batch_script_re_list,
+                              env_test=env_test)
+    del os.environ['ABCDE'], os.environ['VWXYZ']
 
     # req_keepvars AND req_keepvars together
     spawner_kwargs = {