Merge pull request #3 from boegel/cmdlog

minor tweaks to dumping of `env.sh` + `run.sh` helper scripts in `run_shell_cmd` + enhance test to verify they're working as intended

easybuild/tools/run.py

@@ -197,30 +197,38 @@ def fileprefix_from_cmd(cmd, allowed_chars=False):
     return ''.join([c for c in cmd if c in allowed_chars])
 
 
-def save_cmd(cmd_str, work_dir, env, tmpdir):
-    # Save environment variables in it's own environment file
+def create_cmd_scripts(cmd_str, work_dir, env, tmpdir):
+    """
+    Create helper scripts for specified command in specified directory:
+    - env.sh which can be sourced to define environment in which command was run;
+    - cmd.sh to create interactive (bash) shell session with working directory and environment,
+      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)
     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 full_env.items()
+        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\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}"')
         fid.write(f'\nhistory -s {shlex.quote(cmd_str)}')
 
-    # Make script that sets up bash shell with given environments set.
+    # Make script that sets up bash shell with specified environment and working directory
     cmd_fp = os.path.join(tmpdir, 'cmd.sh')
     with open(cmd_fp, 'w') as fid:
         fid.write('#!/usr/bin/env bash\n')
-        fid.write('# Run this script to replicate the environment that EB used to run the shell command\n')
+        fid.write('# Run this script to set up a shell environment that EasyBuild used to run the shell command\n')
         fid.write('\n'.join([
-            f'\ncd "{work_dir}"',
             'EB_SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )',
-            f'echo Shell for the command: {shlex.quote(cmd_str)}',
-            'echo Use command history, exit to stop',
-            'bash --rcfile $EB_SCRIPT_DIR/env.sh',
+            f'echo "# Shell for the command: {shlex.quote(cmd_str)}"',
+            'echo "# Use command history, exit to stop"',
+            # using -i to force interactive shell, so env.sh is also sourced when -c is used to run commands
+            'bash --rcfile $EB_SCRIPT_DIR/env.sh -i "$@"',
             ]))
     os.chmod(cmd_fp, 0o775)
 
@@ -358,14 +366,17 @@ def to_cmd_str(cmd):
         _log.info(f"Auto-enabling streaming output of '{cmd_str}' command because logging to stdout is enabled")
         stream_output = True
 
-    # temporary output file(s) for command output
+    # temporary output file(s) for command output, along with helper scripts
     if output_file:
         toptmpdir = os.path.join(tempfile.gettempdir(), 'run-shell-cmd-output')
         os.makedirs(toptmpdir, exist_ok=True)
         cmd_name = fileprefix_from_cmd(os.path.basename(cmd_str.split(' ')[0]))
         tmpdir = tempfile.mkdtemp(dir=toptmpdir, prefix=f'{cmd_name}-')
+
         _log.info(f'run_shell_cmd: command environment of "{cmd_str}" will be saved to {tmpdir}')
-        save_cmd(cmd_str, work_dir, env, tmpdir)
+
+        create_cmd_scripts(cmd_str, work_dir, env, tmpdir)
+
         cmd_out_fp = os.path.join(tmpdir, 'out.txt')
         _log.info(f'run_shell_cmd: Output of "{cmd_str}" will be logged to {cmd_out_fp}')
         if split_stderr:

test/framework/run.py

@@ -179,6 +179,10 @@ def test_run_cmd(self):
     def test_run_shell_cmd_basic(self):
         """Basic test for run_shell_cmd function."""
 
+        os.environ['FOOBAR'] = 'foobar'
+
+        cwd = change_dir(self.test_prefix)
+
         with self.mocked_stdout_stderr():
             res = run_shell_cmd("echo hello")
         self.assertEqual(res.output, "hello\n")
@@ -189,6 +193,42 @@ def test_run_shell_cmd_basic(self):
         self.assertEqual(res.stderr, None)
         self.assertTrue(res.work_dir and isinstance(res.work_dir, str))
 
+        change_dir(cwd)
+        del os.environ['FOOBAR']
+
+        # check on helper scripts that were generated for this command
+        paths = glob.glob(os.path.join(self.test_prefix, 'eb-*', 'run-shell-cmd-output', 'echo-*'))
+        self.assertEqual(len(paths), 1)
+        cmd_tmpdir = paths[0]
+
+        # check on env.sh script that can be used to set up environment in which command was run
+        env_script = os.path.join(cmd_tmpdir, 'env.sh')
+        self.assertExists(env_script)
+        env_script_txt = read_file(env_script)
+        self.assertIn("export FOOBAR=foobar", env_script_txt)
+        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")
+        self.assertEqual(res.exit_code, 0)
+        self.assertTrue(res.output.startswith('foobar\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
+        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 pwd: $PWD; echo $FOOBAR'", fail_on_error=False)
+        self.assertEqual(res.exit_code, 0)
+        self.assertTrue(res.output.endswith('foobar\n'))
+        # check whether working directory is what's expected
+        regex = re.compile('^pwd: .*', re.M)
+        res = regex.findall(res.output)
+        self.assertEqual(len(res), 1)
+        pwd = res[0].strip()[5:]
+        self.assertTrue(os.path.samefile(pwd, self.test_prefix))
+
         # test running command that emits non-UTF-8 characters
         # this is constructed to reproduce errors like:
         # UnicodeDecodeError: 'utf-8' codec can't decode byte 0xe2