fix handling of specified command environment in create_cmd_scripts + add test for env option in run_shell_cmd

easybuild/tools/run.py

@@ -205,17 +205,26 @@ def create_cmd_scripts(cmd_str, work_dir, env, tmpdir):
       and with the command in shell history;
     """
     # Save environment variables in env.sh which can be sourced to restore environment
-    full_env = os.environ.copy()
-    if env is not None:
-        full_env.update(env)
+    if env is None:
+        env = os.environ.copy()
+
     env_fp = os.path.join(tmpdir, 'env.sh')
     with open(env_fp, 'w') as fid:
+        # unset all environment variables in current environment first to start from a clean slate;
+        # we need to be careful to filter out functions definitions, so first undefine those
+        fid.write("unset -f $(env | grep '%=' | cut -f1 -d'%' | sed 's/BASH_FUNC_//g')\n")
+        fid.write("unset $(env | cut -f1 -d=)\n")
+
         # excludes bash functions (environment variables ending with %)
-        fid.write('\n'.join(f'export {key}={shlex.quote(value)}' for key, value in sorted(full_env.items())
-                            if not key.endswith('%')))
+        fid.write('\n'.join(f'export {key}={shlex.quote(value)}' for key, value in sorted(env.items())
+                            if not key.endswith('%')) + '\n')
+
         fid.write('\n\nPS1="eb-shell> "')
+
         # also change to working directory (to ensure that working directory is correct for interactive bash shell)
         fid.write(f'\ncd "{work_dir}"')
+
+        # reset shell history to only include executed command
         fid.write(f'\nhistory -s {shlex.quote(cmd_str)}')
 
     # Make script that sets up bash shell with specified environment and working directory

test/framework/run.py

@@ -209,9 +209,10 @@ def test_run_shell_cmd_basic(self):
         self.assertIn("history -s 'echo hello'", env_script_txt)
 
         with self.mocked_stdout_stderr():
-            res = run_shell_cmd(f"source {env_script}; echo $FOOBAR; history")
+            res = run_shell_cmd(f"source {env_script}; echo $USER; echo $FOOBAR; history")
         self.assertEqual(res.exit_code, 0)
-        self.assertTrue(res.output.startswith('foobar\n'))
+        user = os.getenv('USER')
+        self.assertTrue(res.output.startswith(f'{user}\nfoobar\n'))
         self.assertTrue(res.output.endswith("echo hello\n"))
 
         # check on cmd.sh script that can be used to create interactive shell environment for command
@@ -247,6 +248,45 @@ def test_run_shell_cmd_basic(self):
             self.assertTrue(isinstance(res.output, str))
             self.assertTrue(res.work_dir and isinstance(res.work_dir, str))
 
+    def test_run_shell_cmd_env(self):
+        """Test env option in run_shell_cmd."""
+
+        # use 'env' to define environment in which command should be run;
+        # with a few exceptions (like $_, $PWD) no other environment variables will be defined,
+        # so $HOME and $USER will not be set
+        cmd = "env | sort"
+        with self.mocked_stdout_stderr():
+            res = run_shell_cmd(cmd, env={'FOOBAR': 'foobar', 'PATH': os.getenv('PATH')})
+        self.assertEqual(res.cmd, cmd)
+        self.assertEqual(res.exit_code, 0)
+        self.assertIn("FOOBAR=foobar\n", res.output)
+        self.assertTrue(re.search("^_=.*/env$", res.output, re.M))
+        for var in ('HOME', 'USER'):
+            self.assertFalse(re.search('^' + var + '=.*', res.output, re.M))
+
+        # check on helper scripts that were generated for this command
+        paths = glob.glob(os.path.join(self.test_prefix, 'eb-*', 'run-shell-cmd-output', 'env-*'))
+        self.assertEqual(len(paths), 1)
+        cmd_tmpdir = paths[0]
+
+        # set environment variable in current environment,
+        # this should not be set in shell environment produced by scripts
+        os.environ['TEST123'] = 'test123'
+
+        env_script = os.path.join(cmd_tmpdir, 'env.sh')
+        self.assertExists(env_script)
+        env_script_txt = read_file(env_script)
+        self.assertTrue(env_script_txt.startswith('unset -f $('))
+        self.assertIn('\nexport FOOBAR=foobar\nexport PATH', env_script_txt)
+
+        cmd_script = os.path.join(cmd_tmpdir, 'cmd.sh')
+        self.assertExists(cmd_script)
+
+        with self.mocked_stdout_stderr():
+            res = run_shell_cmd(f"{cmd_script} -c 'echo $FOOBAR; echo TEST123:$TEST123'", fail_on_error=False)
+        self.assertEqual(res.exit_code, 0)
+        self.assertTrue(res.output.endswith('\nfoobar\nTEST123:\n'))
+
     def test_fileprefix_from_cmd(self):
         """test simplifications from fileprefix_from_cmd."""
         cmds = {