From 1a044de4106afb3f34f076118d91124c3263c948 Mon Sep 17 00:00:00 2001
From: Tom Parker-Shemilt <palfrey@tevp.net>
Date: Sat, 16 Nov 2024 20:54:29 +0000
Subject: [PATCH] Redo log streaming with _stream_lambda_image_build_logs

---
 samcli/lib/build/app_builder.py                 | 12 +++++-------
 samcli/lib/docker/log_streamer.py               | 10 ++++++++--
 tests/unit/lib/build_module/test_app_builder.py |  7 +++++--
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/samcli/lib/build/app_builder.py b/samcli/lib/build/app_builder.py
index e483df72a6..f99a9d3fa8 100644
--- a/samcli/lib/build/app_builder.py
+++ b/samcli/lib/build/app_builder.py
@@ -430,11 +430,7 @@ def _build_lambda_image(self, function_name: str, metadata: Dict, architecture:
             LOG.debug("%s image is built for %s function", build_image, function_name)
         except docker.errors.BuildError as ex:
             LOG.error("Failed building function %s", function_name)
-            for log in ex.build_log:
-                if "stream" in log:
-                    self._stream_writer.write_str(log["stream"])
-                elif "error" in log:
-                    self._stream_writer.write_str(log["error"])
+            self._stream_lambda_image_build_logs(ex.build_log, function_name, False)
             raise DockerBuildFailed(str(ex)) from ex
 
         # The Docker-py low level api will stream logs back but if an exception is raised by the api
@@ -450,7 +446,9 @@ def _build_lambda_image(self, function_name: str, metadata: Dict, architecture:
 
         return docker_tag
 
-    def _stream_lambda_image_build_logs(self, build_logs: List[Dict[str, str]], function_name: str) -> None:
+    def _stream_lambda_image_build_logs(
+        self, build_logs: List[Dict[str, str]], function_name: str, throw_on_error: bool = True
+    ) -> None:
         """
         Stream logs to the console from an Lambda image build.
 
@@ -461,7 +459,7 @@ def _stream_lambda_image_build_logs(self, build_logs: List[Dict[str, str]], func
         function_name str
             Name of the function that is being built
         """
-        build_log_streamer = LogStreamer(self._stream_writer)
+        build_log_streamer = LogStreamer(self._stream_writer, throw_on_error)
         try:
             build_log_streamer.stream_progress(build_logs)
         except LogStreamError as ex:
diff --git a/samcli/lib/docker/log_streamer.py b/samcli/lib/docker/log_streamer.py
index c4aacfd09d..c82998354b 100644
--- a/samcli/lib/docker/log_streamer.py
+++ b/samcli/lib/docker/log_streamer.py
@@ -22,12 +22,13 @@ def __init__(self, msg: str) -> None:
 
 
 class LogStreamer:
-    def __init__(self, stream: StreamWriter):
+    def __init__(self, stream: StreamWriter, throw_on_error: bool = True):
         self._stream = stream
         self._cursor_up_formatter = CursorUpFormatter()
         self._cursor_down_formatter = CursorDownFormatter()
         self._cursor_left_formatter = CursorLeftFormatter()
         self._cursor_clear_formatter = ClearLineFormatter()
+        self._throw_on_error = throw_on_error
 
     def stream_progress(self, logs: docker.APIClient.logs):
         """
@@ -73,7 +74,12 @@ def _stream_write(self, _id: str, status: str, stream: str, progress: str, error
         :param error: docker log error
         """
         if error:
-            raise LogStreamError(msg=error)
+            if self._throw_on_error:
+                raise LogStreamError(msg=error)
+            else:
+                self._stream.write_str(error)
+                return
+
         if not status and not stream:
             return
 
diff --git a/tests/unit/lib/build_module/test_app_builder.py b/tests/unit/lib/build_module/test_app_builder.py
index d4dacc1c8a..f1bbff70cf 100644
--- a/tests/unit/lib/build_module/test_app_builder.py
+++ b/tests/unit/lib/build_module/test_app_builder.py
@@ -1743,8 +1743,11 @@ def test_can_raise_build_error(self):
             self.builder._build_lambda_image("Name", {"Dockerfile": "Dockerfile", "DockerContext": "context"}, X86_64)
 
         self.assertEqual(ex.exception.args, ("Build failure",))
-        self.assertEqual(self.stream_mock.write_str.call_count, 2)
-        self.assertEqual(self.stream_mock.write_str.call_args_list, [call("Some earlier log"), call("Build failed")])
+        self.assertEqual(self.stream_mock.write_str.call_count, 4, self.stream_mock.write_str.call_args_list)
+        self.assertEqual(
+            self.stream_mock.write_str.call_args_list,
+            [call("Some earlier log"), call(""), call("Build failed"), call("\n")],
+        )
 
 
 class TestApplicationBuilder_load_lambda_image_function(TestCase):