"""
Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
"""

import json
import os
import re
import tempfile
import traceback
import unittest
from imp import C_BUILTIN
from io import StringIO
from tempfile import NamedTemporaryFile
from unittest.mock import patch, Mock, MagicMock

import awslambdaric.bootstrap as bootstrap
from awslambdaric.lambda_runtime_exception import FaultException
from awslambdaric.lambda_runtime_marshaller import LambdaMarshaller


class TestUpdateXrayEnv(unittest.TestCase):
    def setUp(self):
        self.org_os_environ = os.environ

    def tearDown(self):
        os.environ = self.org_os_environ

    def test_update_xray_env_variable_empty(self):
        os.environ = {}
        bootstrap.update_xray_env_variable(None)
        self.assertEqual(os.environ.get("_X_AMZN_TRACE_ID"), None)

    def test_update_xray_env_variable_remove_old_value(self):
        os.environ = {"_X_AMZN_TRACE_ID": "old-id"}
        bootstrap.update_xray_env_variable(None)
        self.assertEqual(os.environ.get("_X_AMZN_TRACE_ID"), None)

    def test_update_xray_env_variable_new_value(self):
        os.environ = {}
        bootstrap.update_xray_env_variable("new-id")
        self.assertEqual(os.environ.get("_X_AMZN_TRACE_ID"), "new-id")

    def test_update_xray_env_variable_overwrite(self):
        os.environ = {"_X_AMZN_TRACE_ID": "old-id"}
        bootstrap.update_xray_env_variable("new-id")
        self.assertEqual(os.environ.get("_X_AMZN_TRACE_ID"), "new-id")


class TestHandleEventRequest(unittest.TestCase):
    def setUp(self):
        self.lambda_runtime = Mock()
        self.lambda_runtime.marshaller = LambdaMarshaller()
        self.event_body = '"event_body"'
        self.working_directory = os.getcwd()

    @staticmethod
    def dummy_handler(json_input, lambda_context):
        return {"input": json_input, "aws_request_id": lambda_context.aws_request_id}

    def test_handle_event_request_happy_case(self):
        bootstrap.handle_event_request(
            self.lambda_runtime,
            self.dummy_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            bootstrap.StandardLogSink(),
        )
        self.lambda_runtime.post_invocation_result.assert_called_once_with(
            "invoke_id",
            '{"input": "event_body", "aws_request_id": "invoke_id"}',
            "application/json",
        )

    def test_handle_event_request_invalid_client_context(self):
        expected_response = {
            "errorType": "Runtime.LambdaContextUnmarshalError",
            "errorMessage": "Unable to parse Client Context JSON: Expecting value: line 1 column 1 (char 0)",
        }
        bootstrap.handle_event_request(
            self.lambda_runtime,
            self.dummy_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            "invalid_client_context_not_json",
            {},
            "invoked_function_arn",
            0,
            bootstrap.StandardLogSink(),
        )
        args, _ = self.lambda_runtime.post_invocation_error.call_args
        error_response = json.loads(args[1])
        self.assertEqual(args[0], "invoke_id")
        self.assertTrue(
            expected_response.items() <= error_response.items(),
            "Response doesn't contain all the necessary fields\nExpected: {}\nActual: {}".format(
                expected_response, error_response
            ),
        )
        self.assertEqual(
            json.loads(args[2]),
            {
                "working_directory": self.working_directory,
                "exceptions": [
                    {
                        "message": expected_response["errorMessage"],
                        "type": "LambdaValidationError",
                        "stack": [],
                    }
                ],
                "paths": [],
            },
        )

    def test_handle_event_request_invalid_cognito_idenity(self):
        expected_response = {
            "errorType": "Runtime.LambdaContextUnmarshalError",
            "errorMessage": "Unable to parse Cognito Identity JSON: Expecting value: line 1 column 1 (char 0)",
        }
        bootstrap.handle_event_request(
            self.lambda_runtime,
            self.dummy_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            "invalid_cognito_identity",
            "invoked_function_arn",
            0,
            bootstrap.StandardLogSink(),
        )
        args, _ = self.lambda_runtime.post_invocation_error.call_args
        error_response = json.loads(args[1])
        self.assertEqual(args[0], "invoke_id")
        self.assertTrue(
            expected_response.items() <= error_response.items(),
            "Response doesn't contain all the necessary fields\nExpected: {}\nActual: {}".format(
                expected_response, error_response
            ),
        )
        self.assertEqual(
            json.loads(args[2]),
            {
                "working_directory": self.working_directory,
                "exceptions": [
                    {
                        "message": expected_response["errorMessage"],
                        "type": "LambdaValidationError",
                        "stack": [],
                    }
                ],
                "paths": [],
            },
        )

    def test_handle_event_request_invalid_event_body(self):
        expected_response = {
            "errorType": "Runtime.UnmarshalError",
            "errorMessage": "Unable to unmarshal input: Expecting value: line 1 column 1 (char 0)",
        }
        invalid_event_body = "not_valid_json"
        bootstrap.handle_event_request(
            self.lambda_runtime,
            self.dummy_handler,
            "invoke_id",
            invalid_event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            bootstrap.StandardLogSink(),
        )
        args, _ = self.lambda_runtime.post_invocation_error.call_args
        error_response = json.loads(args[1])
        self.assertEqual(args[0], "invoke_id")
        self.assertTrue(
            expected_response.items() <= error_response.items(),
            "Response doesn't contain all the necessary fields\nExpected: {}\nActual: {}".format(
                expected_response, error_response
            ),
        )
        self.assertEqual(
            json.loads(args[2]),
            {
                "working_directory": self.working_directory,
                "exceptions": [
                    {
                        "message": expected_response["errorMessage"],
                        "type": "LambdaValidationError",
                        "stack": [],
                    }
                ],
                "paths": [],
            },
        )

    def test_handle_event_request_invalid_response(self):
        def invalid_json_response(json_input, lambda_context):
            return type("obj", (object,), {"propertyName": "propertyValue"})

        expected_response = {
            "errorType": "Runtime.MarshalError",
            "errorMessage": "Unable to marshal response: Object of type type is not JSON serializable",
        }
        bootstrap.handle_event_request(
            self.lambda_runtime,
            invalid_json_response,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            bootstrap.StandardLogSink(),
        )
        args, _ = self.lambda_runtime.post_invocation_error.call_args
        error_response = json.loads(args[1])
        self.assertEqual(args[0], "invoke_id")
        self.assertTrue(
            expected_response.items() <= error_response.items(),
            "Expected response is not a subset of the actual response\nExpected: {}\nActual: {}".format(
                expected_response, error_response
            ),
        )
        self.assertEqual(
            json.loads(args[2]),
            {
                "working_directory": self.working_directory,
                "exceptions": [
                    {
                        "message": expected_response["errorMessage"],
                        "type": "LambdaValidationError",
                        "stack": [],
                    }
                ],
                "paths": [],
            },
        )

    def test_handle_event_request_custom_exception(self):
        def raise_exception_handler(json_input, lambda_context):
            class MyError(Exception):
                def __init__(self, message):
                    self.message = message

            raise MyError("My error")

        expected_response = {"errorType": "MyError", "errorMessage": "My error"}
        bootstrap.handle_event_request(
            self.lambda_runtime,
            raise_exception_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            bootstrap.StandardLogSink(),
        )
        args, _ = self.lambda_runtime.post_invocation_error.call_args
        error_response = json.loads(args[1])
        self.assertEqual(args[0], "invoke_id")
        self.assertTrue(
            expected_response.items() <= error_response.items(),
            "Expected response is not a subset of the actual response\nExpected: {}\nActual: {}".format(
                expected_response, error_response
            ),
        )
        xray_fault = json.loads(args[2])
        self.assertEqual(xray_fault["working_directory"], self.working_directory)
        self.assertEqual(len(xray_fault["exceptions"]), 1)
        self.assertEqual(
            xray_fault["exceptions"][0]["message"], expected_response["errorMessage"]
        )
        self.assertEqual(
            xray_fault["exceptions"][0]["type"], expected_response["errorType"]
        )
        self.assertEqual(len(xray_fault["exceptions"][0]["stack"]), 1)
        self.assertEqual(
            xray_fault["exceptions"][0]["stack"][0]["label"], "raise_exception_handler"
        )
        self.assertIsInstance(xray_fault["exceptions"][0]["stack"][0]["line"], int)
        self.assertEqual(
            xray_fault["exceptions"][0]["stack"][0]["path"], os.path.realpath(__file__)
        )
        self.assertEqual(len(xray_fault["paths"]), 1)
        self.assertEqual(xray_fault["paths"][0], os.path.realpath(__file__))

    def test_handle_event_request_no_module(self):
        def unable_to_import_module(json_input, lambda_context):
            import invalid_module

        expected_response = {
            "errorType": "ModuleNotFoundError",
            "errorMessage": "No module named 'invalid_module'",
        }
        bootstrap.handle_event_request(
            self.lambda_runtime,
            unable_to_import_module,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            bootstrap.StandardLogSink(),
        )
        args, _ = self.lambda_runtime.post_invocation_error.call_args
        error_response = json.loads(args[1])
        self.assertEqual(args[0], "invoke_id")
        self.assertTrue(
            expected_response.items() <= error_response.items(),
            "Expected response is not a subset of the actual response\nExpected: {}\nActual: {}".format(
                expected_response, error_response
            ),
        )

    def test_handle_event_request_fault_exception(self):
        def raise_exception_handler(json_input, lambda_context):
            try:
                import invalid_module
            except ImportError as e:
                raise FaultException(
                    "FaultExceptionType",
                    "Fault exception msg",
                    ["trace_line1\ntrace_line2", "trace_line3\ntrace_line4"],
                )

        expected_response = {
            "errorType": "FaultExceptionType",
            "errorMessage": "Fault exception msg",
            "stackTrace": ["trace_line1\ntrace_line2", "trace_line3\ntrace_line4"],
        }
        bootstrap.handle_event_request(
            self.lambda_runtime,
            raise_exception_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            bootstrap.StandardLogSink(),
        )
        args, _ = self.lambda_runtime.post_invocation_error.call_args
        error_response = json.loads(args[1])
        self.assertEqual(args[0], "invoke_id")
        self.assertEqual(error_response.items(), expected_response.items())
        self.assertEqual(
            json.loads(args[2]),
            {
                "working_directory": self.working_directory,
                "exceptions": [
                    {
                        "message": expected_response["errorMessage"],
                        "type": "LambdaValidationError",
                        "stack": [],
                    }
                ],
                "paths": [],
            },
        )

    @patch("sys.stdout", new_callable=StringIO)
    def test_handle_event_request_fault_exception_logging(self, mock_stdout):
        def raise_exception_handler(json_input, lambda_context):
            try:
                import invalid_module
            except ImportError as e:
                raise bootstrap.FaultException(
                    "FaultExceptionType",
                    "Fault exception msg",
                    traceback.format_list(
                        [
                            ("spam.py", 3, "<module>", "spam.eggs()"),
                            ("eggs.py", 42, "eggs", 'return "bacon"'),
                        ]
                    ),
                )

        bootstrap.handle_event_request(
            self.lambda_runtime,
            raise_exception_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            bootstrap.StandardLogSink(),
        )

        # NOTE: Indentation characters are NO-BREAK SPACE (U+00A0) not SPACE (U+0020)
        error_logs = "[ERROR] FaultExceptionType: Fault exception msg\r"
        error_logs += "Traceback (most recent call last):\r"
        error_logs += '  File "spam.py", line 3, in <module>\r'
        error_logs += "    spam.eggs()\r"
        error_logs += '  File "eggs.py", line 42, in eggs\r'
        error_logs += '    return "bacon"\n'

        self.assertEqual(mock_stdout.getvalue(), error_logs)

    @patch("sys.stdout", new_callable=StringIO)
    def test_handle_event_request_fault_exception_logging_notrace(self, mock_stdout):
        def raise_exception_handler(json_input, lambda_context):
            try:
                import invalid_module
            except ImportError as e:
                raise bootstrap.FaultException(
                    "FaultExceptionType", "Fault exception msg", None
                )

        bootstrap.handle_event_request(
            self.lambda_runtime,
            raise_exception_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            bootstrap.StandardLogSink(),
        )
        error_logs = "[ERROR] FaultExceptionType: Fault exception msg\n"

        self.assertEqual(mock_stdout.getvalue(), error_logs)

    @patch("sys.stdout", new_callable=StringIO)
    def test_handle_event_request_fault_exception_logging_nomessage_notrace(
        self, mock_stdout
    ):
        def raise_exception_handler(json_input, lambda_context):
            try:
                import invalid_module
            except ImportError as e:
                raise bootstrap.FaultException("FaultExceptionType", None, None)

        bootstrap.handle_event_request(
            self.lambda_runtime,
            raise_exception_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            bootstrap.StandardLogSink(),
        )
        error_logs = "[ERROR] FaultExceptionType\n"

        self.assertEqual(mock_stdout.getvalue(), error_logs)

    @patch("sys.stdout", new_callable=StringIO)
    def test_handle_event_request_fault_exception_logging_notype_notrace(
        self, mock_stdout
    ):
        def raise_exception_handler(json_input, lambda_context):
            try:
                import invalid_module
            except ImportError as e:
                raise bootstrap.FaultException(None, "Fault exception msg", None)

        bootstrap.handle_event_request(
            self.lambda_runtime,
            raise_exception_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            bootstrap.StandardLogSink(),
        )
        error_logs = "[ERROR] Fault exception msg\n"

        self.assertEqual(mock_stdout.getvalue(), error_logs)

    @patch("sys.stdout", new_callable=StringIO)
    def test_handle_event_request_fault_exception_logging_notype_nomessage(
        self, mock_stdout
    ):
        def raise_exception_handler(json_input, lambda_context):
            try:
                import invalid_module
            except ImportError as e:
                raise bootstrap.FaultException(
                    None,
                    None,
                    traceback.format_list(
                        [
                            ("spam.py", 3, "<module>", "spam.eggs()"),
                            ("eggs.py", 42, "eggs", 'return "bacon"'),
                        ]
                    ),
                )

        bootstrap.handle_event_request(
            self.lambda_runtime,
            raise_exception_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            bootstrap.StandardLogSink(),
        )

        error_logs = "[ERROR]\r"
        error_logs += "Traceback (most recent call last):\r"
        error_logs += '  File "spam.py", line 3, in <module>\r'
        error_logs += "    spam.eggs()\r"
        error_logs += '  File "eggs.py", line 42, in eggs\r'
        error_logs += '    return "bacon"\n'

        self.assertEqual(mock_stdout.getvalue(), error_logs)

    @patch("sys.stdout", new_callable=StringIO)
    @patch("imp.find_module")
    @patch("imp.load_module")
    def test_handle_event_request_fault_exception_logging_syntax_error(
        self, mock_load_module, mock_find_module, mock_stdout
    ):

        try:
            eval("-")
        except SyntaxError as e:
            syntax_error = e

        mock_find_module.return_value = (None, None, ("", "", None))
        mock_load_module.side_effect = syntax_error

        response_handler = bootstrap._get_handler("a.b")

        bootstrap.handle_event_request(
            self.lambda_runtime,
            response_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            bootstrap.StandardLogSink(),
        )

        import sys

        sys.stderr.write(mock_stdout.getvalue())

        error_logs = "[ERROR] Runtime.UserCodeSyntaxError: Syntax error in module 'a': unexpected EOF while parsing (<string>, line 1)\r"
        error_logs += "Traceback (most recent call last):\r"
        error_logs += '  File "<string>" Line 1\r'
        error_logs += "    -\n"

        self.assertEqual(mock_stdout.getvalue(), error_logs)


class TestXrayFault(unittest.TestCase):
    def test_make_xray(self):
        class CustomException(Exception):
            def __init__(self):
                pass

        actual = bootstrap.make_xray_fault(
            CustomException.__name__,
            "test_message",
            "working/dir",
            [["test.py", 28, "test_method", "does_not_matter"]],
        )

        self.assertEqual(actual["working_directory"], "working/dir")
        self.assertEqual(actual["paths"], ["test.py"])
        self.assertEqual(len(actual["exceptions"]), 1)
        self.assertEqual(actual["exceptions"][0]["message"], "test_message")
        self.assertEqual(actual["exceptions"][0]["type"], "CustomException")
        self.assertEqual(len(actual["exceptions"][0]["stack"]), 1)
        self.assertEqual(actual["exceptions"][0]["stack"][0]["label"], "test_method")
        self.assertEqual(actual["exceptions"][0]["stack"][0]["path"], "test.py")
        self.assertEqual(actual["exceptions"][0]["stack"][0]["line"], 28)

    def test_make_xray_with_multiple_tb(self):
        class CustomException(Exception):
            def __init__(self):
                pass

        actual = bootstrap.make_xray_fault(
            CustomException.__name__,
            "test_message",
            "working/dir",
            [
                ["test.py", 28, "test_method", ""],
                ["another_test.py", 2718, "another_test_method", ""],
            ],
        )

        self.assertEqual(len(actual["exceptions"]), 1)
        self.assertEqual(len(actual["exceptions"][0]["stack"]), 2)
        self.assertEqual(actual["exceptions"][0]["stack"][0]["label"], "test_method")
        self.assertEqual(actual["exceptions"][0]["stack"][0]["path"], "test.py")
        self.assertEqual(actual["exceptions"][0]["stack"][0]["line"], 28)
        self.assertEqual(
            actual["exceptions"][0]["stack"][1]["label"], "another_test_method"
        )
        self.assertEqual(actual["exceptions"][0]["stack"][1]["path"], "another_test.py")
        self.assertEqual(actual["exceptions"][0]["stack"][1]["line"], 2718)


class TestGetEventHandler(unittest.TestCase):
    class FaultExceptionMatcher(BaseException):
        def __init__(self, msg, exception_type=None, trace_pattern=None):
            self.msg = msg
            self.exception_type = exception_type
            self.trace = (
                trace_pattern if trace_pattern is None else re.compile(trace_pattern)
            )

        def __eq__(self, other):
            trace_matches = True
            if self.trace is not None:
                # Validate that trace is an array
                if not isinstance(other.trace, list):
                    trace_matches = False
                elif not self.trace.match("".join(other.trace)):
                    trace_matches = False

            return (
                self.msg in other.msg
                and self.exception_type == other.exception_type
                and trace_matches
            )

    def test_get_event_handler_bad_handler(self):
        handler_name = "bad_handler"
        response_handler = bootstrap._get_handler(handler_name)
        with self.assertRaises(FaultException) as cm:
            response_handler()

        returned_exception = cm.exception
        self.assertEqual(
            self.FaultExceptionMatcher(
                "Bad handler 'bad_handler': not enough values to unpack (expected 2, got 1)",
                "Runtime.MalformedHandlerName",
            ),
            returned_exception,
        )

    def test_get_event_handler_import_error(self):
        handler_name = "no_module.handler"
        response_handler = bootstrap._get_handler(handler_name)
        with self.assertRaises(FaultException) as cm:
            response_handler()
        returned_exception = cm.exception
        self.assertEqual(
            self.FaultExceptionMatcher(
                "Unable to import module 'no_module': No module named 'no_module'",
                "Runtime.ImportModuleError",
            ),
            returned_exception,
        )

    def test_get_event_handler_syntax_error(self):
        tmp_file = tempfile.NamedTemporaryFile(suffix=".py", dir=".", delete=False)
        tmp_file.write(
            b"def syntax_error()\n\tprint('syntax error, no colon after function')"
        )
        tmp_file.close()
        filename_w_ext = os.path.basename(tmp_file.name)
        filename, _ = os.path.splitext(filename_w_ext)
        handler_name = "{}.syntax_error".format(filename)
        response_handler = bootstrap._get_handler(handler_name)

        with self.assertRaises(FaultException) as cm:
            response_handler()
        returned_exception = cm.exception
        self.assertEqual(
            self.FaultExceptionMatcher(
                "Syntax error in",
                "Runtime.UserCodeSyntaxError",
                ".*File.*\\.py.*Line 1.*",
            ),
            returned_exception,
        )
        if os.path.exists(tmp_file.name):
            os.remove(tmp_file.name)

    def test_get_event_handler_missing_error(self):
        tmp_file = tempfile.NamedTemporaryFile(suffix=".py", dir=".", delete=False)
        tmp_file.write(b"def wrong_handler_name():\n\tprint('hello')")
        tmp_file.close()
        filename_w_ext = os.path.basename(tmp_file.name)
        filename, _ = os.path.splitext(filename_w_ext)
        handler_name = "{}.my_handler".format(filename)
        response_handler = bootstrap._get_handler(handler_name)
        with self.assertRaises(FaultException) as cm:
            response_handler()
        returned_exception = cm.exception
        self.assertEqual(
            self.FaultExceptionMatcher(
                "Handler 'my_handler' missing on module '{}'".format(filename),
                "Runtime.HandlerNotFound",
            ),
            returned_exception,
        )
        if os.path.exists(tmp_file.name):
            os.remove(tmp_file.name)

    @patch("imp.find_module")
    def test_get_event_handler_build_in_conflict(self, mock_find_module):
        handler_name = "sys.hello"
        mock_find_module.return_value = (None, None, ("", "", C_BUILTIN))
        response_handler = bootstrap._get_handler(handler_name)
        with self.assertRaises(FaultException) as cm:
            response_handler()
        returned_exception = cm.exception
        self.assertEqual(
            self.FaultExceptionMatcher(
                "Cannot use built-in module sys as a handler module",
                "Runtime.BuiltInModuleConflict",
            ),
            returned_exception,
        )


class TestContentType(unittest.TestCase):
    def setUp(self):
        self.lambda_runtime = Mock()
        self.lambda_runtime.marshaller = LambdaMarshaller()

    def test_application_json(self):
        bootstrap.handle_event_request(
            lambda_runtime_client=self.lambda_runtime,
            request_handler=lambda event, ctx: {"response": event["msg"]},
            invoke_id="invoke-id",
            event_body=b'{"msg":"foo"}',
            content_type="application/json",
            client_context_json=None,
            cognito_identity_json=None,
            invoked_function_arn="invocation-arn",
            epoch_deadline_time_in_ms=1415836801003,
            log_sink=bootstrap.StandardLogSink(),
        )

        self.lambda_runtime.post_invocation_result.assert_called_once_with(
            "invoke-id", '{"response": "foo"}', "application/json"
        )

    def test_binary_request_binary_response(self):
        event_body = b"\x89PNG\r\n\x1a\n\x00\x00\x00"
        bootstrap.handle_event_request(
            lambda_runtime_client=self.lambda_runtime,
            request_handler=lambda event, ctx: event,
            invoke_id="invoke-id",
            event_body=event_body,
            content_type="image/png",
            client_context_json=None,
            cognito_identity_json=None,
            invoked_function_arn="invocation-arn",
            epoch_deadline_time_in_ms=1415836801003,
            log_sink=bootstrap.StandardLogSink(),
        )

        self.lambda_runtime.post_invocation_result.assert_called_once_with(
            "invoke-id", event_body, "application/unknown"
        )

    def test_json_request_binary_response(self):
        binary_data = b"\x89PNG\r\n\x1a\n\x00\x00\x00"
        bootstrap.handle_event_request(
            lambda_runtime_client=self.lambda_runtime,
            request_handler=lambda event, ctx: binary_data,
            invoke_id="invoke-id",
            event_body=b'{"msg":"ignored"}',
            content_type="application/json",
            client_context_json=None,
            cognito_identity_json=None,
            invoked_function_arn="invocation-arn",
            epoch_deadline_time_in_ms=1415836801003,
            log_sink=bootstrap.StandardLogSink(),
        )

        self.lambda_runtime.post_invocation_result.assert_called_once_with(
            "invoke-id", binary_data, "application/unknown"
        )

    def test_binary_with_application_json(self):
        bootstrap.handle_event_request(
            lambda_runtime_client=self.lambda_runtime,
            request_handler=lambda event, ctx: event,
            invoke_id="invoke-id",
            event_body=b"\x89PNG\r\n\x1a\n\x00\x00\x00",
            content_type="application/json",
            client_context_json=None,
            cognito_identity_json=None,
            invoked_function_arn="invocation-arn",
            epoch_deadline_time_in_ms=1415836801003,
            log_sink=bootstrap.StandardLogSink(),
        )

        self.lambda_runtime.post_invocation_result.assert_not_called()
        self.lambda_runtime.post_invocation_error.assert_called_once()

        (
            invoke_id,
            error_result,
            xray_fault,
        ), _ = self.lambda_runtime.post_invocation_error.call_args
        error_dict = json.loads(error_result)

        self.assertEqual("invoke-id", invoke_id)
        self.assertEqual("Runtime.UnmarshalError", error_dict["errorType"])


class TestLogError(unittest.TestCase):
    @patch("sys.stdout", new_callable=StringIO)
    def test_log_error_standard_log_sink(self, mock_stdout):
        err_to_log = bootstrap.make_error("Error message", "ErrorType", None)
        bootstrap.log_error(err_to_log, bootstrap.StandardLogSink())

        expected_logged_error = "[ERROR] ErrorType: Error message\n"
        self.assertEqual(mock_stdout.getvalue(), expected_logged_error)

    def test_log_error_framed_log_sink(self):
        with NamedTemporaryFile() as temp_file:
            with bootstrap.FramedTelemetryLogSink(temp_file.name) as log_sink:
                err_to_log = bootstrap.make_error("Error message", "ErrorType", None)
                bootstrap.log_error(err_to_log, log_sink)

            expected_logged_error = "[ERROR] ErrorType: Error message"

            with open(temp_file.name, "rb") as f:
                content = f.read()

                frame_type = int.from_bytes(content[:4], "big")
                self.assertEqual(frame_type, 0xA55A0001)

                length = int.from_bytes(content[4:8], "big")
                self.assertEqual(length, len(expected_logged_error.encode("utf8")))

                actual_message = content[8:].decode()
                self.assertEqual(actual_message, expected_logged_error)

    @patch("sys.stdout", new_callable=StringIO)
    def test_log_error_indentation_standard_log_sink(self, mock_stdout):
        err_to_log = bootstrap.make_error(
            "Error message", "ErrorType", ["  line1  ", "  line2  ", "  "]
        )
        bootstrap.log_error(err_to_log, bootstrap.StandardLogSink())

        expected_logged_error = "[ERROR] ErrorType: Error message\rTraceback (most recent call last):\r\xa0\xa0line1  \r\xa0\xa0line2  \r\xa0\xa0\n"
        self.assertEqual(mock_stdout.getvalue(), expected_logged_error)

    def test_log_error_indentation_framed_log_sink(self):
        with NamedTemporaryFile() as temp_file:
            with bootstrap.FramedTelemetryLogSink(temp_file.name) as log_sink:
                err_to_log = bootstrap.make_error(
                    "Error message", "ErrorType", ["  line1  ", "  line2  ", "  "]
                )
                bootstrap.log_error(err_to_log, log_sink)

            expected_logged_error = "[ERROR] ErrorType: Error message\nTraceback (most recent call last):\n\xa0\xa0line1  \n\xa0\xa0line2  \n\xa0\xa0"

            with open(temp_file.name, "rb") as f:
                content = f.read()

                frame_type = int.from_bytes(content[:4], "big")
                self.assertEqual(frame_type, 0xA55A0001)

                length = int.from_bytes(content[4:8], "big")
                self.assertEqual(length, len(expected_logged_error.encode("utf8")))

                actual_message = content[8:].decode()
                self.assertEqual(actual_message, expected_logged_error)

    @patch("sys.stdout", new_callable=StringIO)
    def test_log_error_empty_stacktrace_line_standard_log_sink(self, mock_stdout):
        err_to_log = bootstrap.make_error(
            "Error message", "ErrorType", ["line1", "", "line2"]
        )
        bootstrap.log_error(err_to_log, bootstrap.StandardLogSink())

        expected_logged_error = "[ERROR] ErrorType: Error message\rTraceback (most recent call last):\rline1\r\rline2\n"
        self.assertEqual(mock_stdout.getvalue(), expected_logged_error)

    def test_log_error_empty_stacktrace_line_framed_log_sink(self):
        with NamedTemporaryFile() as temp_file:
            with bootstrap.FramedTelemetryLogSink(temp_file.name) as log_sink:
                err_to_log = bootstrap.make_error(
                    "Error message", "ErrorType", ["line1", "", "line2"]
                )
                bootstrap.log_error(err_to_log, log_sink)

            expected_logged_error = "[ERROR] ErrorType: Error message\nTraceback (most recent call last):\nline1\n\nline2"

            with open(temp_file.name, "rb") as f:
                content = f.read()

                frame_type = int.from_bytes(content[:4], "big")
                self.assertEqual(frame_type, 0xA55A0001)

                length = int.from_bytes(content[4:8], "big")
                self.assertEqual(length, len(expected_logged_error))

                actual_message = content[8:].decode()
                self.assertEqual(actual_message, expected_logged_error)


class TestUnbuffered(unittest.TestCase):
    def test_write(self):
        mock_stream = MagicMock()
        unbuffered = bootstrap.Unbuffered(mock_stream)

        unbuffered.write("YOLO!")

        mock_stream.write.assert_called_once_with("YOLO!")
        mock_stream.flush.assert_called_once()

    def test_writelines(self):
        mock_stream = MagicMock()
        unbuffered = bootstrap.Unbuffered(mock_stream)

        unbuffered.writelines(["YOLO!"])

        mock_stream.writelines.assert_called_once_with(["YOLO!"])
        mock_stream.flush.assert_called_once()


class TestLogSink(unittest.TestCase):
    @patch("sys.stdout", new_callable=StringIO)
    def test_create_unbuffered_log_sinks(self, mock_stdout):
        if "_LAMBDA_TELEMETRY_LOG_FD" in os.environ:
            del os.environ["_LAMBDA_TELEMETRY_LOG_FD"]

        actual = bootstrap.create_log_sink()

        self.assertIsInstance(actual, bootstrap.StandardLogSink)
        actual.log("log")
        self.assertEqual(mock_stdout.getvalue(), "log")

    def test_create_framed_telemetry_log_sinks(self):
        fd = "test_fd"
        os.environ["_LAMBDA_TELEMETRY_LOG_FD"] = fd

        actual = bootstrap.create_log_sink()

        self.assertIsInstance(actual, bootstrap.FramedTelemetryLogSink)
        self.assertEqual(actual.filename, "/proc/self/fd/" + fd)
        self.assertFalse("_LAMBDA_TELEMETRY_LOG_FD" in os.environ)

    def test_single_frame(self):
        with NamedTemporaryFile() as temp_file:
            message = "hello world\nsomething on a new line!\n"
            with bootstrap.FramedTelemetryLogSink(temp_file.name) as ls:
                ls.log(message)
            with open(temp_file.name, "rb") as f:
                content = f.read()

                frame_type = int.from_bytes(content[:4], "big")
                self.assertEqual(frame_type, 0xA55A0001)

                length = int.from_bytes(content[4:8], "big")
                self.assertEqual(length, len(message))

                actual_message = content[8:].decode()
                self.assertEqual(actual_message, message)

    def test_multiple_frame(self):
        with NamedTemporaryFile() as temp_file:
            first_message = "hello world\nsomething on a new line!"
            second_message = "hello again\nhere's another message\n"

            with bootstrap.FramedTelemetryLogSink(temp_file.name) as ls:
                ls.log(first_message)
                ls.log(second_message)

            with open(temp_file.name, "rb") as f:
                content = f.read()
                pos = 0
                for message in [first_message, second_message]:
                    frame_type = int.from_bytes(content[pos : pos + 4], "big")
                    self.assertEqual(frame_type, 0xA55A0001)
                    pos += 4

                    length = int.from_bytes(content[pos : pos + 4], "big")
                    self.assertEqual(length, len(message))
                    pos += 4

                    actual_message = content[pos : pos + len(message)].decode()
                    self.assertEqual(actual_message, message)
                    pos += len(message)

                self.assertEqual(content[pos:], b"")


class TestBootstrapModule(unittest.TestCase):
    @patch("awslambdaric.bootstrap.handle_event_request")
    @patch("awslambdaric.bootstrap.LambdaRuntimeClient")
    def test_run(self, mock_runtime_client, mock_handle_event_request):
        expected_app_root = "/tmp/test/app_root"
        expected_handler = "app.my_test_handler"
        expected_lambda_runtime_api_addr = "test_addr"

        mock_event_request = MagicMock()
        mock_event_request.x_amzn_trace_id = "123"

        mock_runtime_client.return_value.wait_next_invocation.side_effect = [
            mock_event_request,
            MagicMock(),
        ]

        with self.assertRaises(TypeError) as cm:
            bootstrap.run(
                expected_app_root, expected_handler, expected_lambda_runtime_api_addr
            )
        returned_exception = cm.exception

        mock_handle_event_request.assert_called_once()

    @patch(
        "awslambdaric.bootstrap.LambdaLoggerHandler",
        Mock(side_effect=Exception("Boom!")),
    )
    @patch("awslambdaric.bootstrap.build_fault_result", MagicMock())
    @patch("awslambdaric.bootstrap.log_error", MagicMock())
    @patch("awslambdaric.bootstrap.LambdaRuntimeClient", MagicMock())
    @patch("awslambdaric.bootstrap.sys")
    def test_run_exception(self, mock_sys):
        class TestException(Exception):
            pass

        expected_app_root = "/tmp/test/app_root"
        expected_handler = "app.my_test_handler"
        expected_lambda_runtime_api_addr = "test_addr"

        mock_sys.exit.side_effect = TestException("Boom!")

        with self.assertRaises(TestException) as cm:
            bootstrap.run(
                expected_app_root, expected_handler, expected_lambda_runtime_api_addr
            )
        returned_exception = cm.exception

        mock_sys.exit.assert_called_once_with(1)


if __name__ == "__main__":
    unittest.main()
