From 1d40adab5474c06e61bba3cdbefdf2dbe968b3da Mon Sep 17 00:00:00 2001
From: Sriram Madapusi Vasudevan <srirammv@amazon.com>
Date: Thu, 27 Oct 2022 17:28:11 -0700
Subject: [PATCH] fix: `Makefile` workflow to stream output

---
 .../workflows/custom_make/make.py             | 39 +++++++++++++------
 tests/unit/workflows/custom_make/test_make.py | 10 ++++-
 2 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/aws_lambda_builders/workflows/custom_make/make.py b/aws_lambda_builders/workflows/custom_make/make.py
index 964912d44..5f097aa04 100644
--- a/aws_lambda_builders/workflows/custom_make/make.py
+++ b/aws_lambda_builders/workflows/custom_make/make.py
@@ -1,8 +1,11 @@
 """
 Wrapper around calling make through a subprocess.
 """
-
+import io
 import logging
+import subprocess
+import shutil
+import threading
 
 LOG = logging.getLogger(__name__)
 
@@ -82,14 +85,26 @@ def run(self, args, env=None, cwd=None):
 
         p = self.osutils.popen(invoke_make, stdout=self.osutils.pipe, stderr=self.osutils.pipe, cwd=cwd, env=env)
 
-        out, err = p.communicate()
-
-        # Typically this type of information is logged to DEBUG, however, since the Make builder
-        # can be different per customer's use case, it is helpful to always log the output so
-        # developers can diagnose any issues.
-        LOG.info(out.decode("utf8").strip())
-
-        if p.returncode != 0:
-            raise MakeExecutionError(message=err.decode("utf8").strip())
-
-        return out.decode("utf8").strip()
+        # Create a stdout variable that will contain the final stitched stdout result
+        stdout = ""
+        # Create a buffer and use a thread to gather the stderr stream into the buffer
+        stderr_buf = io.BytesIO()
+        stderr_thread = threading.Thread(target=shutil.copyfileobj, args=(p.stderr, stderr_buf), daemon=True)
+        stderr_thread.start()
+
+        # Log every stdout line by iterating
+        for line in p.stdout:
+            decoded_line = line.decode("utf-8").strip()
+            LOG.info(decoded_line)
+            # Gather total stdout
+            stdout += decoded_line
+
+        # Wait for the process to exit and stderr thread to end.
+        return_code = p.wait()
+        stderr_thread.join()
+
+        if return_code != 0:
+            # Raise an Error with the appropriate value from the stderr buffer.
+            raise MakeExecutionError(message=stderr_buf.getvalue().decode("utf8").strip())
+
+        return stdout
diff --git a/tests/unit/workflows/custom_make/test_make.py b/tests/unit/workflows/custom_make/test_make.py
index 99d566fd2..0696725d9 100644
--- a/tests/unit/workflows/custom_make/test_make.py
+++ b/tests/unit/workflows/custom_make/test_make.py
@@ -1,3 +1,4 @@
+import io
 from unittest import TestCase
 from mock import patch
 
@@ -8,11 +9,16 @@ class FakePopen:
     def __init__(self, out=b"out", err=b"err", retcode=0):
         self.out = out
         self.err = err
+        self.stderr = io.BytesIO(err)
+        self.stdout = [out]
         self.returncode = retcode
 
     def communicate(self):
         return self.out, self.err
 
+    def wait(self):
+        return self.returncode
+
 
 class TestSubprocessMake(TestCase):
     @patch("aws_lambda_builders.workflows.custom_make.utils.OSUtils")
@@ -67,7 +73,7 @@ def test_uses_env_and_cwd_if_supplied(self):
         )
 
     def test_returns_popen_out_decoded_if_retcode_is_0(self):
-        self.popen.out = b"some encoded text\n\n"
+        self.popen.stdout = [b"some encoded text\n\n"]
 
         result = self.under_test.run(["build_logical_id"])
 
@@ -75,7 +81,7 @@ def test_returns_popen_out_decoded_if_retcode_is_0(self):
 
     def test_raises_MakeExecutionError_with_err_text_if_retcode_is_not_0(self):
         self.popen.returncode = 1
-        self.popen.err = b"some error text\n\n"
+        self.popen.stderr = io.BytesIO(b"some error text\n\n")
 
         with self.assertRaises(MakeExecutionError) as raised:
             self.under_test.run(["build-logical_id"])