From 88f6523fc005fdd1c85f43cf8383ff8b4d597828 Mon Sep 17 00:00:00 2001 From: Richard Darst Date: Wed, 5 Dec 2018 11:21:51 +0200 Subject: [PATCH] Fix all the environment variable handling, make tests work - This fixes up all the previous commits on this branch --- batchspawner/batchspawner.py | 26 +++++++-- batchspawner/tests/test_spawners.py | 81 +++++++++++++++++++++-------- 2 files changed, 79 insertions(+), 28 deletions(-) diff --git a/batchspawner/batchspawner.py b/batchspawner/batchspawner.py index a54ab288..1cf087fe 100644 --- a/batchspawner/batchspawner.py +++ b/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. diff --git a/batchspawner/tests/test_spawners.py b/batchspawner/tests/test_spawners.py index adbd1ff8..b71b3536 100644 --- a/batchspawner/tests/test_spawners.py +++ b/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)