From a6677280ada3fa389f421d25e830db8f02ec5191 Mon Sep 17 00:00:00 2001
From: pryce-turner <pryce.turner@gmail.com>
Date: Wed, 27 Dec 2023 09:51:04 -0800
Subject: [PATCH] Generalized subproc_execute and added tests

Signed-off-by: pryce-turner <pryce.turner@gmail.com>
---
 flytekit/extras/tasks/shell.py                 | 13 +++++++++++--
 tests/flytekit/unit/extras/tasks/test_shell.py | 18 +++++++++++++++++-
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index 319a927743..57d7eb0109 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -34,7 +34,7 @@ class OutputLocation:
     location: typing.Union[os.PathLike, str]
 
 
-def subproc_execute(command: List[str]) -> Tuple[str, str]:
+def subproc_execute(command: typing.Union[List[str], str], **kwargs) -> Tuple[str, str]:
     """
     Execute a command and capture its stdout and stderr. Useful for executing
     shell commands from within a python task.
@@ -52,9 +52,18 @@ def subproc_execute(command: List[str]) -> Tuple[str, str]:
             guidance on specifying a container image in the task definition when
             using custom dependencies.
     """
+    defaults = {
+        "stdout": subprocess.PIPE,
+        "stderr": subprocess.PIPE,
+        "text": True,
+        "check": True,
+    }
+
+    kwargs = {**defaults, **kwargs}
+
     try:
         # Execute the command and capture stdout and stderr
-        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)
+        result = subprocess.run(command, **kwargs)
 
         # Access the stdout and stderr output
         return result.stdout, result.stderr
diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index 08e5bb92af..0ca7ca66bf 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -11,7 +11,7 @@
 import flytekit
 from flytekit import kwtypes
 from flytekit.exceptions.user import FlyteRecoverableException
-from flytekit.extras.tasks.shell import OutputLocation, RawShellTask, ShellTask, get_raw_shell_task
+from flytekit.extras.tasks.shell import OutputLocation, RawShellTask, ShellTask, get_raw_shell_task, subproc_execute
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import CSVFile, FlyteFile
 
@@ -323,3 +323,19 @@ def test_long_run_script():
         name="long-running",
         script=script,
     )()
+
+
+def test_subproc_execute():
+    cmd = ["echo", "hello"]
+    o, e = subproc_execute(cmd)
+    assert o == "hello\n"
+    assert e == ""
+
+
+def test_subproc_execute_with_shell():
+    with tempfile.TemporaryDirectory() as tmp:
+        opth = os.path.join(tmp, "test.txt")
+        cmd = f"echo hello > {opth}"
+        subproc_execute(cmd, shell=True)
+        cont = open(opth).read()
+        assert "hello" in cont