Fix all the environment variable handling, make tests work

- This fixes up all the previous commits on this branch

Author: rkdarst

batchspawner/batchspawner.py

@@ -148,16 +148,18 @@ def _req_homedir_default(self):
              "to a default by JupyterHub, and if you change this list "
              "the previous ones are _not_ included and your submissions will "
              "break unless you re-add necessary variables.  Consider "
-             "req_keepvars for most use cases."
-)
-    @default('req_keepvars_all')
-    def _req_keepvars_all_default(self):
-        return ','.join(self.get_env().keys())
+             "req_keepvars for most use cases, don't edit this under normal "
+             "use.").tag(config=True)
+
+    @default('req_keepvarsdefault')
+    def _req_keepvars_default_default(self):
+        return ','.join(super(BatchSpawnerBase, self).get_env().keys())
 
     req_keepvars = Unicode(
         help="Extra environment variables which should be configured, "
               "added to the defaults in keepvars, "
               "comma separated list.").tag(config=True)
+
     admin_environment = Unicode(
         help="Comma-separated list of environment variables to be passed to "
              "the batch submit/cancel commands, but _not_ to the batch script "
@@ -213,6 +215,20 @@ def cmd_formatted_for_batch(self):
         """The command which is substituted inside of the batch script"""
         return ' '.join([self.batchspawner_singleuser_cmd] + self.cmd + self.get_args())
 
+    def get_env(self):
+        """Get the env dict from JH, adding req_keepvars options
+
+        get_env() returns the variables given by JH, but not anything
+        specified by req_keepvars since that's our creation.  Add those
+        (from the JH environment) to the environment passed to the batch
+        start/stop/cancel/etc commands.
+        """
+        env = super(BatchSpawnerBase, self).get_env()
+        for k in self.req_keepvars.split(',') + self.req_keepvars_default.split(','):
+            if k in os.environ:
+                env[k] = os.environ[k]
+        return env
+
     def get_admin_env(self):
         """Get the environment passed to the batch submit/cancel/etc commands.
 

batchspawner/tests/test_spawners.py

@@ -1,6 +1,8 @@
 """Test BatchSpawner and subclasses"""
 
+import contextlib
 import itertools
+import os
 import re
 from unittest import mock
 from .. import BatchSpawnerRegexStates
@@ -20,6 +22,21 @@
 testjob  = "12345"
 testport = 54321
 
+@contextlib.contextmanager
+def setenv_context(**kwargs):
+    """Context manage which sets and restores environment variables."""
+    orig = { }
+    for k, v in kwargs.items():
+        orig[k] = os.environ.get(k, None)
+        os.environ[k] = v
+    yield
+    for k in kwargs:
+        if orig[k] is not None:
+            os.environ[k] = orig[k]
+        else:
+            del os.environ[k]
+
+
 class BatchDummy(BatchSpawnerRegexStates):
     exec_prefix = ''
     batch_submit_cmd = Unicode('cat > /dev/null; echo '+testjob)
@@ -504,32 +521,50 @@ def test_lfs(db, io_loop):
 def test_keepvars(db, io_loop):
     """Test of environment handling
     """
-    # req_keepvars
+    environment = {'ABCDE': 'TEST1', 'VWXYZ': 'TEST2', 'XYZ': 'TEST3',}
+
+
+    # req_keepvars_default - anything NOT here should not be propogated.
     spawner_kwargs = {
         'req_keepvars_default': 'ABCDE',
         }
     batch_script_re_list = [
         re.compile(r'--export=ABCDE', re.X|re.M),
+        re.compile(r'^((?!JUPYTERHUB_API_TOKEN).)*$', re.X|re.S),  # *not* in the script
         ]
     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,
-                              env_test=env_test)
-
-    # req_keepvars
+        # We can't test these - becasue removing these from the environment is
+        # a job of the batch system itself, which we do *not* run here.
+        #assert 'ABCDE' in env
+        #assert 'JUPYTERHUB_API_TOKEN' not in env
+        pass
+    with setenv_context(**environment):
+        run_typical_slurm_spawner(db, io_loop,
+                                  spawner_kwargs=spawner_kwargs,
+                                  batch_script_re_list=batch_script_re_list,
+                                  env_test=env_test)
+
+    # req_keepvars - this should be added to the environment
     spawner_kwargs = {
         'req_keepvars': 'ABCDE',
         }
     batch_script_re_list = [
         re.compile(r'--export=.*ABCDE', re.X|re.M),
+        re.compile(r'^((?!VWXYZ).)*$', re.X|re.M), # *not* in line
+        re.compile(r'--export=.*JUPYTERHUB_API_TOKEN', re.X|re.S),
         ]
-    run_typical_slurm_spawner(db, io_loop,
-                              spawner_kwargs=spawner_kwargs,
-                              batch_script_re_list=batch_script_re_list)
-
-    # req_keepvars
+    def env_test(env):
+        assert 'ABCDE' in env
+        assert 'VWXYZ' not in env
+    with setenv_context(**environment):
+        run_typical_slurm_spawner(db, io_loop,
+                                  spawner_kwargs=spawner_kwargs,
+                                  batch_script_re_list=batch_script_re_list,
+                                  env_test=env_test)
+
+    # admin_environment - this should be in the environment passed to
+    # run commands but not the --export command which is included in
+    # the batch scripts
     spawner_kwargs = {
         'admin_environment': 'ABCDE',
         }
@@ -539,13 +574,12 @@ def env_test(env):
     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,
-                              env_test=env_test)
-    del os.environ['ABCDE'], os.environ['VWXYZ']
+        assert 'JUPYTERHUB_API_TOKEN' in env
+    with setenv_context(**environment):
+        run_typical_slurm_spawner(db, io_loop,
+                                  spawner_kwargs=spawner_kwargs,
+                                  batch_script_re_list=batch_script_re_list,
+                                  env_test=env_test)
 
     # req_keepvars AND req_keepvars together
     spawner_kwargs = {
@@ -555,6 +589,7 @@ def env_test(env):
     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)
+    with setenv_context(**environment):
+        run_typical_slurm_spawner(db, io_loop,
+                                  spawner_kwargs=spawner_kwargs,
+                                  batch_script_re_list=batch_script_re_list)