fix: Flush stdout and stderr on each write when debugging (aws#843)

* Introduce StreamWriter to wrap output streams. Fix aws#835

* Fix typo

samcli/commands/local/cli_common/invoke_context.py

@@ -7,6 +7,7 @@
 import os
 
 import samcli.lib.utils.osutils as osutils
+from samcli.lib.utils.stream_writer import StreamWriter
 from samcli.commands.local.lib.local_lambda import LocalLambdaRunner
 from samcli.commands.local.lib.debug_context import DebugContext
 from samcli.local.lambdafn.runtime import LambdaRuntime
@@ -202,26 +203,28 @@ def local_lambda_runner(self):
     @property
     def stdout(self):
         """
-        Returns a stdout stream to output Lambda function logs to
+        Returns stream writer for stdout to output Lambda function logs to
 
-        :return File like object: Stream where the output of the function is sent to
+        Returns
+        -------
+        samcli.lib.utils.stream_writer.StreamWriter
+            Stream writer for stdout
         """
-        if self._log_file_handle:
-            return self._log_file_handle
-
-        return osutils.stdout()
+        stream = self._log_file_handle if self._log_file_handle else osutils.stdout()
+        return StreamWriter(stream, self._is_debugging)
 
     @property
     def stderr(self):
         """
-        Returns stderr stream to output Lambda function errors to
+        Returns stream writer for stderr to output Lambda function errors to
 
-        :return File like object: Stream where the stderr of the function is sent to
+        Returns
+        -------
+        samcli.lib.utils.stream_writer.StreamWriter
+            Stream writer for stderr
         """
-        if self._log_file_handle:
-            return self._log_file_handle
-
-        return osutils.stderr()
+        stream = self._log_file_handle if self._log_file_handle else osutils.stderr()
+        return StreamWriter(stream, self._is_debugging)
 
     @property
     def template(self):
@@ -256,6 +259,10 @@ def parameter_overrides(self):
 
         return self._parameter_overrides
 
+    @property
+    def _is_debugging(self):
+        return bool(self._debug_context)
+
     @staticmethod
     def _get_template_data(template_file):
         """

samcli/commands/local/lib/debug_context.py

@@ -1,7 +1,6 @@
 """
 Information and debug options for a specific runtime.
 """
-import os
 
 
 class DebugContext(object):
@@ -14,8 +13,6 @@ def __init__(self,
         self.debug_port = debug_port
         self.debugger_path = debugger_path
         self.debug_args = debug_args
-        if self.debug_port:
-            os.environ["PYTHONUNBUFFERED"] = "1"
 
     def __bool__(self):
         return bool(self.debug_port)

samcli/commands/local/lib/local_lambda.py

@@ -53,11 +53,21 @@ def invoke(self, function_name, event, stdout=None, stderr=None):
 
         This function will block until either the function completes or times out.
 
-        :param string function_name: Name of the Lambda function to invoke
-        :param string event: Event data passed to the function. Must be a valid JSON String.
-        :param io.BaseIO stdout: Stream to write the output of the Lambda function to.
-        :param io.BaseIO stderr: Stream to write the Lambda runtime logs to.
-        :raises FunctionNotfound: When we cannot find a function with the given name
+        Parameters
+        ----------
+        function_name str
+            Name of the Lambda function to invoke
+        event str
+            Event data passed to the function. Must be a valid JSON String.
+        stdout samcli.lib.utils.stream_writer.StreamWriter
+            Stream writer to write the output of the Lambda function to.
+        stderr samcli.lib.utils.stream_writer.StreamWriter
+            Stream writer to write the Lambda runtime logs to.
+
+        Raises
+        ------
+        FunctionNotfound
+            When we cannot find a function with the given name
         """
 
         # Generate the correct configuration based on given inputs

samcli/lib/utils/stream_writer.py

@@ -0,0 +1,37 @@
+"""
+This class acts like a wrapper around output streams to provide any flexibility with output we need
+"""
+
+
+class StreamWriter(object):
+
+    def __init__(self, stream, auto_flush=False):
+        """
+        Instatiates new StreamWriter to the specified stream
+
+        Parameters
+        ----------
+        stream io.RawIOBase
+            Stream to wrap
+        auto_flush bool
+            Whether to autoflush the stream upon writing
+        """
+        self._stream = stream
+        self._auto_flush = auto_flush
+
+    def write(self, output):
+        """
+        Writes specified text to the underlying stream
+
+        Parameters
+        ----------
+        output bytes-like object
+            Bytes to write
+        """
+        self._stream.write(output)
+
+        if self._auto_flush:
+            self._stream.flush()
+
+    def flush(self):
+        self._stream.flush()

samcli/local/apigw/local_apigw_service.py

@@ -7,6 +7,7 @@
 from flask import Flask, request
 
 from samcli.local.services.base_local_service import BaseLocalService, LambdaOutputParser, CaseInsensitiveDict
+from samcli.lib.utils.stream_writer import StreamWriter
 from samcli.local.lambdafn.exceptions import FunctionNotFound
 from samcli.local.events.api_event import ContextIdentity, RequestContext, ApiGatewayLambdaEvent
 from .service_error_responses import ServiceErrorResponses
@@ -40,16 +41,22 @@ def __init__(self, routing_list, lambda_runner, static_dir=None, port=None, host
         """
         Creates an ApiGatewayService
 
-        :param list(ApiGatewayCallModel) routing_list: A list of the Model that represent
-          the service paths to create.
-        :param samcli.commands.local.lib.local_lambda.LocalLambdaRunner lambda_runner: The Lambda runner class capable
-            of invoking the function
-        :param str static_dir: Directory from which to serve static files
-        :param int port: Optional. port for the service to start listening on
-          Defaults to 3000
-        :param str host: Optional. host to start the service on
-          Defaults to '127.0.0.1
-        :param io.BaseIO stderr: Optional stream where the stderr from Docker container should be written to
+        Parameters
+        ----------
+        routing_list list(ApiGatewayCallModel)
+            A list of the Model that represent the service paths to create.
+        lambda_runner samcli.commands.local.lib.local_lambda.LocalLambdaRunner
+            The Lambda runner class capable of invoking the function
+        static_dir str
+            Directory from which to serve static files
+        port int
+            Optional. port for the service to start listening on
+            Defaults to 3000
+        host str
+            Optional. host to start the service on
+            Defaults to '127.0.0.1
+        stderr samcli.lib.utils.stream_writer.StreamWriter
+            Optional stream writer where the stderr from Docker container should be written to
         """
         super(LocalApigwService, self).__init__(lambda_runner.is_debugging(), port=port, host=host)
         self.routing_list = routing_list
@@ -123,9 +130,14 @@ def _request_handler(self, **kwargs):
         * We then transform the response or errors we get from the Invoke and return the data back to
           the caller
 
-        :param kwargs dict: Keyword Args that are passed to the function from Flask. This happens when we have
-            Path Parameters.
-        :return: Response object
+        Parameters
+        ----------
+        kwargs dict
+            Keyword Args that are passed to the function from Flask. This happens when we have path parameters
+
+        Returns
+        -------
+        Response object
         """
         route = self._get_current_route(request)
 
@@ -135,9 +147,10 @@ def _request_handler(self, **kwargs):
             return ServiceErrorResponses.lambda_failure_response()
 
         stdout_stream = io.BytesIO()
+        stdout_stream_writer = StreamWriter(stdout_stream, self.is_debugging)
 
         try:
-            self.lambda_runner.invoke(route.function_name, event, stdout=stdout_stream, stderr=self.stderr)
+            self.lambda_runner.invoke(route.function_name, event, stdout=stdout_stream_writer, stderr=self.stderr)
         except FunctionNotFound:
             return ServiceErrorResponses.lambda_not_found_response()
 

samcli/local/docker/container.py

@@ -168,9 +168,10 @@ def start(self, input_data=None):
         It waits for the container to complete, fetches both stdout and stderr logs and returns through the
         given streams.
 
-        :param input_data: Optional. Input data sent to the container through container's stdin.
-        :param io.StringIO stdout: Optional. IO Stream to that receives stdout text from container.
-        :param io.StringIO stderr: Optional. IO Stream that receives stderr text from container
+        Parameters
+        ----------
+        input_data
+            Optional. Input data sent to the container through container's stdin.
         """
 
         if input_data:
@@ -233,10 +234,10 @@ def _write_container_output(output_itr, stdout=None, stderr=None):
         ----------
         output_itr: Iterator
             Iterator returned by the Docker Attach command
-        stdout: io.BaseIO, optional
-            Stream to write stdout data from Container into
-        stderr: io.BaseIO, optional
-            Stream to write stderr data from the Container into
+        stdout: samcli.lib.utils.stream_writer.StreamWriter, optional
+            Stream writer to write stdout data from Container into
+        stderr: samcli.lib.utils.stream_writer.StreamWriter, optional
+            Stream writer to write stderr data from the Container into
         """
 
         # Iterator returns a tuple of (frame_type, data) where the frame type determines which stream we write output

samcli/local/docker/manager.py

@@ -3,11 +3,13 @@
 """
 
 import logging
-import sys
 
+import sys
 import docker
 import requests
 
+from samcli.lib.utils.stream_writer import StreamWriter
+
 LOG = logging.getLogger(__name__)
 
 
@@ -107,28 +109,37 @@ def pull_image(self, image_name, stream=None):
         """
         Ask Docker to pull the container image with given name.
 
-        :param string image_name: Name of the image
-        :param stream: Optional stream to write output to. Defaults to stderr
-        :raises DockerImagePullFailedException: If the Docker image was not available in the server
+        Parameters
+        ----------
+        image_name str
+            Name of the image
+        stream samcli.lib.utils.stream_writer.StreamWriter
+            Optional stream writer to output to. Defaults to stderr
+
+        Raises
+        ------
+        DockerImagePullFailedException
+            If the Docker image was not available in the server
         """
-        stream = stream or sys.stderr
+        stream_writer = stream or StreamWriter(sys.stderr)
+
         try:
             result_itr = self.docker_client.api.pull(image_name, stream=True, decode=True)
         except docker.errors.APIError as ex:
             LOG.debug("Failed to download image with name %s", image_name)
             raise DockerImagePullFailedException(str(ex))
 
         # io streams, especially StringIO, work only with unicode strings
-        stream.write(u"\nFetching {} Docker container image...".format(image_name))
+        stream_writer.write(u"\nFetching {} Docker container image...".format(image_name))
 
         # Each line contains information on progress of the pull. Each line is a JSON string
         for _ in result_itr:
             # For every line, print a dot to show progress
-            stream.write(u'.')
-            stream.flush()
+            stream_writer.write(u'.')
+            stream_writer.flush()
 
         # We are done. Go to the next line
-        stream.write(u"\n")
+        stream_writer.write(u"\n")
 
     def has_image(self, image_name):
         """

samcli/local/lambda_service/local_lambda_invoke_service.py

@@ -6,7 +6,7 @@
 
 from flask import Flask, request
 
-
+from samcli.lib.utils.stream_writer import StreamWriter
 from samcli.local.services.base_local_service import BaseLocalService, LambdaOutputParser, CaseInsensitiveDict
 from samcli.local.lambdafn.exceptions import FunctionNotFound
 from .lambda_error_responses import LambdaErrorResponses
@@ -139,9 +139,10 @@ def _invoke_request_handler(self, function_name):
         request_data = request_data.decode('utf-8')
 
         stdout_stream = io.BytesIO()
+        stdout_stream_writer = StreamWriter(stdout_stream, self.is_debugging)
 
         try:
-            self.lambda_runner.invoke(function_name, request_data, stdout=stdout_stream, stderr=self.stderr)
+            self.lambda_runner.invoke(function_name, request_data, stdout=stdout_stream_writer, stderr=self.stderr)
         except FunctionNotFound:
             LOG.debug('%s was not found to invoke.', function_name)
             return LambdaErrorResponses.resource_not_found(function_name)

tests/functional/commands/local/lib/test_local_lambda.py

@@ -8,6 +8,7 @@
 import shutil
 import logging
 
+from samcli.lib.utils.stream_writer import StreamWriter
 from samcli.commands.local.lib import provider
 from samcli.commands.local.lib.local_lambda import LocalLambdaRunner
 from samcli.local.lambdafn.runtime import LambdaRuntime
@@ -82,7 +83,10 @@ def test_must_invoke(self):
 
         stdout_stream = io.BytesIO()
         stderr_stream = io.BytesIO()
-        runner.invoke(self.function_name, input_event, stdout=stdout_stream, stderr=stderr_stream)
+
+        stdout_stream_writer = StreamWriter(stdout_stream)
+        stderr_stream_writer = StreamWriter(stderr_stream)
+        runner.invoke(self.function_name, input_event, stdout=stdout_stream_writer, stderr=stderr_stream_writer)
 
         # stderr is where the Lambda container runtime logs are available. It usually contains requestId, start time
         # etc. So it is non-zero in size
@@ -93,4 +97,4 @@ def test_must_invoke(self):
 
         for key, value in expected_env_vars.items():
             self.assertTrue(key in actual_output, "Key '{}' must be in function output".format(key))
-            self.assertEquals(actual_output.get(key), value)
+            self.assertEqual(actual_output.get(key), value)

tests/functional/local/docker/test_lambda_container.py

@@ -10,6 +10,7 @@
 from contextlib import contextmanager
 from unittest import TestCase
 
+from samcli.lib.utils.stream_writer import StreamWriter
 from samcli.commands.local.lib.debug_context import DebugContext
 from tests.functional.function_code import nodejs_lambda
 from samcli.local.docker.lambda_container import LambdaContainer
@@ -130,13 +131,17 @@ def test_function_result_is_available_in_stdout_and_logs_in_stderr(self):
         layer_downloader = LayerDownloader("./", "./")
         image_builder = LambdaImage(layer_downloader, False, False)
         container = LambdaContainer(self.runtime, self.handler, self.code_dir, self.layers, image_builder)
+
         stdout_stream = io.BytesIO()
         stderr_stream = io.BytesIO()
 
+        stdout_stream_writer = StreamWriter(stdout_stream)
+        stderr_stream_writer = StreamWriter(stderr_stream)
+
         with self._create(container):
 
             container.start()
-            container.wait_for_logs(stdout=stdout_stream, stderr=stderr_stream)
+            container.wait_for_logs(stdout=stdout_stream_writer, stderr=stderr_stream_writer)
 
             function_output = stdout_stream.getvalue()
             function_stderr = stderr_stream.getvalue()