Chef - Use temporary folder for stateful_shell artifacts (#19362)

examples/chef/stateful_shell.py

@@ -16,6 +16,7 @@
 import shlex
 import subprocess
 import sys
+import tempfile
 import time
 from typing import Dict, Optional
 
@@ -49,11 +50,6 @@ def __init__(self) -> None:
         self.env: Dict[str, str] = os.environ.copy()
         self.cwd: str = self.env["PWD"]
 
-        # This file holds the env after running a command. This is a better approach
-        # than writing to stdout because commands could redirect the stdout.
-        self._envfile_path: str = os.path.join(_HERE, _ENV_FILENAME)
-        self._cmd_output_path: str = os.path.join(_HERE, _OUTPUT_FILENAME)
-
     def print_env(self) -> None:
         """Print environment variables in commandline friendly format for export.
 
@@ -87,72 +83,70 @@ def run_cmd(
         Returns:
           Output of command if return_cmd_output set to True.
         """
-        env_dict = {}
-        # Set OLDPWD at beginning because opening the shell clears this. This handles 'cd -'.
-        # env -0 prints the env variables separated by null characters for easy parsing.
-
-        if return_cmd_output:
-            # Piping won't work here because piping will affect how environment variables
-            # are propagated. This solution uses tee without piping to preserve env variables.
-            redirect = f" > >(tee \"{self._cmd_output_path}\") 2>&1 "  # include stderr
-
-            # Delete the file before running the command so we can later check if the file
-            # exists as a signal that tee has finished writing to the file.
-            if os.path.isfile(self._cmd_output_path):
-                os.remove(self._cmd_output_path)
-        else:
-            redirect = ""
-
-        # TODO: Use env -0 when `macos-latest` refers to macos-12 in github actions.
-        # env -0 is ideal because it will support cases where an env variable that has newline
-        # characters. The flag "-0" is requires MacOS 12 which is still in beta in Github Actions.
-        # The less ideal `env` command is used by itself, with the caveat that newline chars
-        # are unsupported in env variables.
-        save_env_cmd = f"env > {self._envfile_path}"
-
-        command_with_state = (
-            f"OLDPWD={self.env.get('OLDPWD', '')}; {cmd} {redirect}; RETCODE=$?; "
-            f"{save_env_cmd}; exit $RETCODE")
-        with subprocess.Popen(
-            [command_with_state],
-            env=self.env, cwd=self.cwd,
-            shell=True, executable=self.shell_app
-        ) as proc:
-            returncode = proc.wait()
-
-        # Load env state from envfile.
-        with open(self._envfile_path, encoding="latin1") as f:
-            # TODO: Split on null char after updating to env -0 - requires MacOS 12.
-            env_entries = f.read().split("\n")
-            for entry in env_entries:
-                parts = entry.split("=")
-                # Handle case where an env variable contains text with '='.
-                env_dict[parts[0]] = "=".join(parts[1:])
-            self.env = env_dict
-            self.cwd = self.env["PWD"]
-
-        if raise_on_returncode and returncode != 0:
-            raise RuntimeError(
-                "Error. Nonzero return code."
-                f"\nReturncode: {returncode}"
-                f"\nCmd: {cmd}")
-
-        if return_cmd_output:
-            # Poll for file due to give 'tee' time to close.
-            # This is necessary because 'tee' waits for all subshells to finish before writing.
-            start_time = time.time()
-            while time.time() - start_time < _TEE_WAIT_TIMEOUT:
-                try:
-                    with open(self._cmd_output_path, encoding="latin1") as f:
-                        output = f.read()
-                    break
-                except FileNotFoundError:
-                    pass
-                time.sleep(0.1)
+        with tempfile.TemporaryDirectory(dir=os.path.dirname(_HERE)) as temp_dir:
+            envfile_path: str = os.path.join(temp_dir, _ENV_FILENAME)
+            cmd_output_path: str = os.path.join(temp_dir, _OUTPUT_FILENAME)
+
+            env_dict = {}
+            # Set OLDPWD at beginning because opening the shell clears this. This handles 'cd -'.
+            # env -0 prints the env variables separated by null characters for easy parsing.
+
+            if return_cmd_output:
+                # Piping won't work here because piping will affect how environment variables
+                # are propagated. This solution uses tee without piping to preserve env variables.
+                redirect = f" > >(tee \"{cmd_output_path}\") 2>&1 "  # include stderr
             else:
-                raise TimeoutError(
-                    f"Error. Output file: {self._cmd_output_path} not created within "
-                    f"the alloted time of: {_TEE_WAIT_TIMEOUT}s"
-                )
-
-            return output
+                redirect = ""
+
+            # TODO: Use env -0 when `macos-latest` refers to macos-12 in github actions.
+            # env -0 is ideal because it will support cases where an env variable that has newline
+            # characters. The flag "-0" is requires MacOS 12 which is still in beta in Github Actions.
+            # The less ideal `env` command is used by itself, with the caveat that newline chars
+            # are unsupported in env variables.
+            save_env_cmd = f"env > {envfile_path}"
+
+            command_with_state = (
+                f"OLDPWD={self.env.get('OLDPWD', '')}; {cmd} {redirect}; RETCODE=$?; "
+                f"{save_env_cmd}; exit $RETCODE")
+            with subprocess.Popen(
+                [command_with_state],
+                env=self.env, cwd=self.cwd,
+                shell=True, executable=self.shell_app
+            ) as proc:
+                returncode = proc.wait()
+
+            # Load env state from envfile.
+            with open(envfile_path, encoding="latin1") as f:
+                # TODO: Split on null char after updating to env -0 - requires MacOS 12.
+                env_entries = f.read().split("\n")
+                for entry in env_entries:
+                    parts = entry.split("=")
+                    # Handle case where an env variable contains text with '='.
+                    env_dict[parts[0]] = "=".join(parts[1:])
+                self.env = env_dict
+                self.cwd = self.env["PWD"]
+
+            if raise_on_returncode and returncode != 0:
+                raise RuntimeError(
+                    "Error. Nonzero return code."
+                    f"\nReturncode: {returncode}"
+                    f"\nCmd: {cmd}")
+
+            if return_cmd_output:
+                # Poll for file due to give 'tee' time to close.
+                # This is necessary because 'tee' waits for all subshells to finish before writing.
+                start_time = time.time()
+                while time.time() - start_time < _TEE_WAIT_TIMEOUT:
+                    if os.path.isfile(cmd_output_path):
+                        with open(cmd_output_path, encoding="latin1") as f:
+                            output = f.read()
+                            if output:  # Ensure that file has been written to.
+                                break
+                    time.sleep(0.1)
+                else:
+                    raise TimeoutError(
+                        f"Error. Output file: {cmd_output_path} not created within "
+                        f"the alloted time of: {_TEE_WAIT_TIMEOUT}s"
+                    )
+
+                return output