From f7866811bc1429ce224bf6a35448cb44aa5155e7 Mon Sep 17 00:00:00 2001
From: ANGkeith <angkeith@hotmail.sg>
Date: Fri, 12 Apr 2024 17:04:07 +0800
Subject: [PATCH] fix: Zip source directory should read from sh_work_dir (#560)

* fix: :zip:embedded source directory should read from sh_work_dir

Signed-off-by: ANGkeith <ANGkeith@hotmail.sg>

* refactor: write to a tempfile instead of using fd

* refactor: move tests out to own file

* docs: add comments to explain code

---------

Signed-off-by: ANGkeith <ANGkeith@hotmail.sg>
---
 package.py               | 48 ++++++++++++++++++++++++----------------
 tests/test_zip_source.py | 46 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 19 deletions(-)
 create mode 100644 tests/test_zip_source.py

diff --git a/package.py b/package.py
index d5ae6703..e1c5c910 100644
--- a/package.py
+++ b/package.py
@@ -916,27 +916,37 @@ def execute(self, build_plan, zip_stream, query):
                             # XXX: timestamp=0 - what actually do with it?
                             zs.write_dirs(rd, prefix=prefix, timestamp=0)
             elif cmd == "sh":
-                path, script = action[1:]
-                p = subprocess.Popen(
-                    script,
-                    shell=True,
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.PIPE,
-                    cwd=path,
-                )
+                with tempfile.NamedTemporaryFile(mode="w+t", delete=True) as temp_file:
+                    path, script = action[1:]
+                    # NOTE: Execute `pwd` to determine the subprocess shell's working directory after having executed all other commands.
+                    script = f"{script} && pwd >{temp_file.name}"
+                    p = subprocess.Popen(
+                        script,
+                        shell=True,
+                        stdout=subprocess.PIPE,
+                        stderr=subprocess.PIPE,
+                        cwd=path,
+                    )
 
-                p.wait()
-                call_stdout, call_stderr = p.communicate()
-                exit_code = p.returncode
-                log.info("exit_code: %s", exit_code)
-                if exit_code != 0:
-                    raise RuntimeError(
-                        "Script did not run successfully, exit code {}: {} - {}".format(
-                            exit_code,
-                            call_stdout.decode("utf-8").strip(),
-                            call_stderr.decode("utf-8").strip(),
+                    p.wait()
+                    temp_file.seek(0)
+
+                    # NOTE: This var `sh_work_dir` is consumed in cmd == "zip" loop
+                    sh_work_dir = temp_file.read().strip()
+
+                    log.info("WD: %s", sh_work_dir)
+
+                    call_stdout, call_stderr = p.communicate()
+                    exit_code = p.returncode
+                    log.info("exit_code: %s", exit_code)
+                    if exit_code != 0:
+                        raise RuntimeError(
+                            "Script did not run successfully, exit code {}: {} - {}".format(
+                                exit_code,
+                                call_stdout.decode("utf-8").strip(),
+                                call_stderr.decode("utf-8").strip(),
+                            )
                         )
-                    )
             elif cmd == "set:filter":
                 patterns = action[1]
                 pf = ZipContentFilter(args=self._args)
diff --git a/tests/test_zip_source.py b/tests/test_zip_source.py
new file mode 100644
index 00000000..f46466e1
--- /dev/null
+++ b/tests/test_zip_source.py
@@ -0,0 +1,46 @@
+import os
+from unittest.mock import MagicMock, Mock
+
+from package import BuildPlanManager
+
+
+def test_zip_source_path_sh_work_dir():
+    zs = Mock()
+    zs.write_dirs = MagicMock()
+
+    bpm = BuildPlanManager(args=Mock())
+
+    bpm.execute(
+        build_plan=[
+            ["sh", ".", "cd $(mktemp -d)\n echo pip install"],
+            ["zip:embedded", ".", "./python"],
+        ],
+        zip_stream=zs,
+        query=None,
+    )
+
+    zs.write_dirs.assert_called_once()
+
+    zip_source_path = zs.write_dirs.call_args_list[0][0][0]
+    assert zip_source_path != f"{os.getcwd()}"
+
+
+def test_zip_source_path():
+    zs = Mock()
+    zs.write_dirs = MagicMock()
+
+    bpm = BuildPlanManager(args=Mock())
+
+    bpm.execute(
+        build_plan=[
+            ["sh", ".", "echo pip install"],
+            ["zip:embedded", ".", "./python"],
+        ],
+        zip_stream=zs,
+        query=None,
+    )
+
+    zs.write_dirs.assert_called_once()
+
+    zip_source_path = zs.write_dirs.call_args_list[0][0][0]
+    assert zip_source_path == f"{os.getcwd()}"