Skip to content

Commit

Permalink
fix(invoke): Write in UTF-8 string instead of bytes. (aws#5232)
Browse files Browse the repository at this point in the history
* fix(invoke): Write in UTF-8 string instead of bytes.

It appears that we were using sys.stdout.buffer to support python2
and python3 at the same time. Switching to just write to sys.stdout
allows us to write a utf-8 encoding string. When using sys.stdout.buffer,
we can only write bytes and I couldn't get the correct UTF8 encoded
string to print correctly.

* Fix ruff errors

* Update log_streamer.py to remove encoding

* More updates to make everything work better in general

* Fix with ruff again

* Explictingly write to stream for building images

* More patching writes

* More patching

* Fix long line

* Use mock over io.string

* More fixing of tests

* Assert mock instead of data directly

* More small edits in test

* Verify through calls instead of value

* run make black

* Fix when we flush to match pervious behavior and output

* add integration tests

* run make black

---------

Co-authored-by: Jacob Fuss <jfuss@users.noreply.github.com>
Co-authored-by: Mehmet Nuri Deveci <5735811+mndeveci@users.noreply.github.com>
  • Loading branch information
3 people authored and lucashuy committed Jun 22, 2023
1 parent 431fc8b commit c38d085
Show file tree
Hide file tree
Showing 10 changed files with 64 additions and 34 deletions.
4 changes: 2 additions & 2 deletions samcli/lib/utils/osutils.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ def stdout():
io.BytesIO
Byte stream of Stdout
"""
return sys.stdout.buffer
return sys.stdout


def stderr():
Expand All @@ -99,7 +99,7 @@ def stderr():
io.BytesIO
Byte stream of stderr
"""
return sys.stderr.buffer
return sys.stderr


def remove(path):
Expand Down
7 changes: 5 additions & 2 deletions samcli/lib/utils/stream_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ def __init__(self, stream, auto_flush=False):
def stream(self):
return self._stream

def write(self, output, encode=False):
def write(self, output, encode=False, write_to_buffer=True):
"""
Writes specified text to the underlying stream
Expand All @@ -31,7 +31,10 @@ def write(self, output, encode=False):
output bytes-like object
Bytes to write
"""
self._stream.write(output.encode() if encode else output)
if write_to_buffer:
self._stream.buffer.write(output.encode() if encode else output)
else:
self._stream.write(output)

if self._auto_flush:
self._stream.flush()
Expand Down
5 changes: 4 additions & 1 deletion samcli/local/docker/container.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""
Representation of a generic Docker container
"""
import json
import logging
import os
import pathlib
Expand Down Expand Up @@ -324,7 +325,8 @@ def wait_for_http_response(self, name, event, stdout):
data=event.encode("utf-8"),
timeout=(self.RAPID_CONNECTION_TIMEOUT, None),
)
stdout.write(resp.content)
stdout.write(json.dumps(json.loads(resp.content), ensure_ascii=False), write_to_buffer=False)
stdout.flush()

def wait_for_result(self, full_path, event, stdout, stderr, start_timer=None):
# NOTE(sriram-mv): Let logging happen in its own thread, so that a http request can be sent.
Expand Down Expand Up @@ -434,6 +436,7 @@ def _write_container_output(output_itr, stdout=None, stderr=None):

if stderr_data and stderr:
stderr.write(stderr_data)

except Exception as ex:
LOG.debug("Failed to get the logs from the container", exc_info=ex)

Expand Down
10 changes: 5 additions & 5 deletions samcli/local/docker/lambda_image.py
Original file line number Diff line number Diff line change
Expand Up @@ -226,7 +226,7 @@ def build(self, runtime, packagetype, image, layers, architecture, stream=None,
or not runtime
):
stream_writer = stream or StreamWriter(sys.stderr)
stream_writer.write("Building image...")
stream_writer.write("Building image...", write_to_buffer=False)
stream_writer.flush()
self._build_image(
image if image else base_image, rapid_image, downloaded_layers, architecture, stream=stream_writer
Expand Down Expand Up @@ -337,15 +337,15 @@ def set_item_permission(tar_info):
platform=get_docker_platform(architecture),
)
for log in resp_stream:
stream_writer.write(".")
stream_writer.write(".", write_to_buffer=False)
stream_writer.flush()
if "error" in log:
stream_writer.write("\n")
stream_writer.write("\n", write_to_buffer=False)
LOG.exception("Failed to build Docker Image")
raise ImageBuildException("Error building docker image: {}".format(log["error"]))
stream_writer.write("\n")
stream_writer.write("\n", write_to_buffer=False)
except (docker.errors.BuildError, docker.errors.APIError) as ex:
stream_writer.write("\n")
stream_writer.write("\n", write_to_buffer=False)
LOG.exception("Failed to build Docker Image")
raise ImageBuildException("Building Image failed.") from ex
finally:
Expand Down
8 changes: 5 additions & 3 deletions samcli/local/docker/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -168,16 +168,18 @@ def pull_image(self, image_name, tag=None, stream=None):
raise DockerImagePullFailedException(str(ex)) from ex

# io streams, especially StringIO, work only with unicode strings
stream_writer.write("\nFetching {}:{} Docker container image...".format(image_name, tag))
stream_writer.write(
"\nFetching {}:{} Docker container image...".format(image_name, tag), write_to_buffer=False
)

# 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_writer.write(".")
stream_writer.write(".", write_to_buffer=False)
stream_writer.flush()

# We are done. Go to the next line
stream_writer.write("\n")
stream_writer.write("\n", write_to_buffer=False)

def has_image(self, image_name):
"""
Expand Down
21 changes: 21 additions & 0 deletions tests/integration/local/invoke/test_integrations_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -291,6 +291,27 @@ def test_invoke_returns_expected_result_when_no_event_given(self):
self.assertEqual(process.returncode, 0)
self.assertEqual("{}", process_stdout.decode("utf-8"))

# @pytest.mark.flaky(reruns=3)
def test_invoke_returns_utf8(self):
command_list = InvokeIntegBase.get_command_list(
"EchoEventFunction", template_path=self.template_path, event_path=self.event_utf8_path
)

process = Popen(command_list, stdout=PIPE)
try:
stdout, _ = process.communicate(timeout=TIMEOUT)
except TimeoutExpired:
process.kill()
raise

process_stdout = stdout.strip()

with open(self.event_utf8_path) as f:
expected_output = json.dumps(json.load(f), ensure_ascii=False)

self.assertEqual(process.returncode, 0)
self.assertEqual(expected_output, process_stdout.decode("utf-8"))

@pytest.mark.flaky(reruns=3)
def test_invoke_with_env_using_parameters(self):
command_list = InvokeIntegBase.get_command_list(
Expand Down
8 changes: 0 additions & 8 deletions tests/unit/lib/utils/test_osutils.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,29 +34,21 @@ def test_raises_on_cleanup_failure(self, rmdir_mock):
@patch("os.rmdir")
def test_handles_ignore_error_case(self, rmdir_mock):
rmdir_mock.side_effect = OSError("fail")
dir_name = None
with osutils.mkdir_temp(ignore_errors=True) as tempdir:
dir_name = tempdir
self.assertTrue(os.path.exists(tempdir))


class Test_stderr(TestCase):
def test_must_return_sys_stderr(self):
expected_stderr = sys.stderr

if sys.version_info.major > 2:
expected_stderr = sys.stderr.buffer

self.assertEqual(expected_stderr, osutils.stderr())


class Test_stdout(TestCase):
def test_must_return_sys_stdout(self):
expected_stdout = sys.stdout

if sys.version_info.major > 2:
expected_stdout = sys.stdout.buffer

self.assertEqual(expected_stdout, osutils.stdout())


Expand Down
2 changes: 1 addition & 1 deletion tests/unit/lib/utils/test_stream_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ def test_must_write_to_stream(self):
writer = StreamWriter(stream_mock)
writer.write(buffer)

stream_mock.write.assert_called_once_with(buffer)
stream_mock.buffer.write.assert_called_once_with(buffer)

def test_must_flush_underlying_stream(self):
stream_mock = Mock()
Expand Down
9 changes: 4 additions & 5 deletions tests/unit/local/docker/test_lambda_image.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import io
import tempfile

from unittest import TestCase
Expand Down Expand Up @@ -271,7 +270,7 @@ def test_force_building_image_that_doesnt_already_exists(
docker_client_mock.images.get.side_effect = ImageNotFound("image not found")
docker_client_mock.images.list.return_value = []

stream = io.StringIO()
stream = Mock()

lambda_image = LambdaImage(layer_downloader_mock, False, True, docker_client=docker_client_mock)
actual_image_id = lambda_image.build(
Expand Down Expand Up @@ -311,7 +310,7 @@ def test_force_building_image_on_daemon_404(
docker_client_mock.images.get.side_effect = NotFound("image not found")
docker_client_mock.images.list.return_value = []

stream = io.StringIO()
stream = Mock()

lambda_image = LambdaImage(layer_downloader_mock, False, True, docker_client=docker_client_mock)
actual_image_id = lambda_image.build(
Expand Down Expand Up @@ -351,7 +350,7 @@ def test_docker_distribution_api_error_on_daemon_api_error(
docker_client_mock.images.get.side_effect = APIError("error from docker daemon")
docker_client_mock.images.list.return_value = []

stream = io.StringIO()
stream = Mock()

lambda_image = LambdaImage(layer_downloader_mock, False, True, docker_client=docker_client_mock)
with self.assertRaises(DockerDistributionAPIError):
Expand All @@ -377,7 +376,7 @@ def test_not_force_building_image_that_doesnt_already_exists(
docker_client_mock.images.get.side_effect = ImageNotFound("image not found")
docker_client_mock.images.list.return_value = []

stream = io.StringIO()
stream = Mock()

lambda_image = LambdaImage(layer_downloader_mock, False, False, docker_client=docker_client_mock)
actual_image_id = lambda_image.build(
Expand Down
24 changes: 17 additions & 7 deletions tests/unit/local/docker/test_manager.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
"""
Tests container manager
"""

import io
import importlib
from unittest import TestCase
from unittest.mock import Mock, patch, MagicMock, ANY, call
Expand Down Expand Up @@ -218,17 +216,29 @@ def setUp(self):
self.manager = ContainerManager(docker_client=self.mock_docker_client)

def test_must_pull_and_print_progress_dots(self):
stream = io.StringIO()
stream = Mock()
pull_result = [1, 2, 3, 4, 5, 6, 7, 8, 9, 0]
self.mock_docker_client.api.pull.return_value = pull_result
expected_stream_output = "\nFetching {}:latest Docker container image...{}\n".format(
self.image_name, "." * len(pull_result) # Progress bar will print one dot per response from pull API
)
expected_stream_calls = [
call(f"\nFetching {self.image_name}:latest Docker container image...", write_to_buffer=False),
call(".", write_to_buffer=False),
call(".", write_to_buffer=False),
call(".", write_to_buffer=False),
call(".", write_to_buffer=False),
call(".", write_to_buffer=False),
call(".", write_to_buffer=False),
call(".", write_to_buffer=False),
call(".", write_to_buffer=False),
call(".", write_to_buffer=False),
call(".", write_to_buffer=False),
call("\n", write_to_buffer=False),
]

self.manager.pull_image(self.image_name, stream=stream)

self.mock_docker_client.api.pull.assert_called_with(self.image_name, stream=True, decode=True, tag="latest")
self.assertEqual(stream.getvalue(), expected_stream_output)

stream.write.assert_has_calls(expected_stream_calls)

def test_must_raise_if_image_not_found(self):
msg = "some error"
Expand Down

0 comments on commit c38d085

Please sign in to comment.