fix: Zip source directory should read from sh_work_dir (terraform-aws-modules#560)

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

Signed-off-by: ANGkeith <[email protected]>

* 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 <[email protected]>

Author: ANGkeith

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)

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()}"