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,24 @@ 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:
         # 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')
+
+        # unset environment variables in current environment if they're not defined in environment used to run command
+        fid.write('\n'.join(f'unset {key}' for key in os.environ if key not in env and not key.endswith('%')))
+
         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

@@ -247,6 +247,40 @@ 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]
+
+        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('export FOOBAR=foobar\nexport PATH'))
+
+        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'", fail_on_error=False)
+        self.assertEqual(res.exit_code, 0)
+        self.assertTrue(res.output.endswith('\nfoobar\n'))
+
     def test_fileprefix_from_cmd(self):
         """test simplifications from fileprefix_from_cmd."""
         cmds = {