test_bootstrap.py

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

import importlib
import json
import logging
import logging.config
import os
import re
import tempfile
import time
import traceback
import unittest
from io import StringIO
from tempfile import NamedTemporaryFile
from unittest.mock import MagicMock, Mock, patch, ANY

import awslambdaric.bootstrap as bootstrap
from awslambdaric.lambda_runtime_exception import FaultException
from awslambdaric.lambda_runtime_log_utils import (
    LogFormat,
    _get_log_level_from_env_var,
    JsonFormatter,
)
from awslambdaric.lambda_runtime_marshaller import LambdaMarshaller
from awslambdaric.lambda_literals import (
    lambda_unhandled_exception_warning_message,
)
import snapshot_restore_py


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()

        logging.getLogger().handlers.clear()

    def tearDown(self) -> None:
        logging.getLogger().handlers.clear()
        logging.getLogger().level = logging.NOTSET

        return super().tearDown()

    @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,
            "tenant_id",
            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,
            "tenant_id",
            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,
            "tenant_id",
            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,
            "tenant_id",
            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,
            "tenant_id",
            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,
            "tenant_id",
            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.assertTrue(
            xray_fault["exceptions"][0]["stack"][0]["path"].endswith(
                os.path.relpath(__file__)
            )
        )
        self.assertEqual(len(xray_fault["paths"]), 1)
        self.assertTrue(xray_fault["paths"][0].endswith(os.path.relpath(__file__)))

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

            raise MyError("")

        expected_response = {"errorType": "MyError", "errorMessage": ""}
        bootstrap.handle_event_request(
            self.lambda_runtime,
            raise_exception_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            "tenant_id",
            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.assertTrue(
            xray_fault["exceptions"][0]["stack"][0]["path"].endswith(
                os.path.relpath(__file__)
            )
        )
        self.assertEqual(len(xray_fault["paths"]), 1)
        self.assertTrue(xray_fault["paths"][0].endswith(os.path.relpath(__file__)))

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

        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,
            "tenant_id",
            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  # noqa: F401
            except ImportError:
                raise FaultException(
                    "FaultExceptionType",
                    "Fault exception msg",
                    ["trace_line1\ntrace_line2", "trace_line3\ntrace_line4"],
                )

        expected_response = {
            "errorType": "FaultExceptionType",
            "errorMessage": "Fault exception msg",
            "requestId": "invoke_id",
            "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,
            "tenant_id",
            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  # noqa: F401
            except ImportError:
                raise bootstrap.FaultException(
                    "FaultExceptionType",
                    "Fault exception msg",
                    traceback.format_list(
                        [
                            ("spam.py", 3, "<module>", "spam.eggs()"),
                            ("eggs.py", 42, "eggs", 'return "bacon"'),
                        ]
                    ),
                )

        logging.getLogger().addHandler(logging.StreamHandler(mock_stdout))

        bootstrap.handle_event_request(
            self.lambda_runtime,
            raise_exception_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            "tenant_id",
            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  # noqa: F401
            except ImportError:
                raise bootstrap.FaultException(
                    "FaultExceptionType", "Fault exception msg", None
                )

        logging.getLogger().addHandler(logging.StreamHandler(mock_stdout))

        bootstrap.handle_event_request(
            self.lambda_runtime,
            raise_exception_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            "tenant_id",
            bootstrap.StandardLogSink(),
        )
        error_logs = "[ERROR] FaultExceptionType: Fault exception msg\rTraceback (most recent call last):\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  # noqa: F401
            except ImportError:
                raise bootstrap.FaultException("FaultExceptionType", None, None)

        logging.getLogger().addHandler(logging.StreamHandler(mock_stdout))

        bootstrap.handle_event_request(
            self.lambda_runtime,
            raise_exception_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            "tenant_id",
            bootstrap.StandardLogSink(),
        )
        error_logs = "[ERROR] FaultExceptionType\rTraceback (most recent call last):\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  # noqa: F401
            except ImportError:
                raise bootstrap.FaultException(None, "Fault exception msg", None)

        logging.getLogger().addHandler(logging.StreamHandler(mock_stdout))

        bootstrap.handle_event_request(
            self.lambda_runtime,
            raise_exception_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            "tenant_id",
            bootstrap.StandardLogSink(),
        )
        error_logs = "[ERROR] Fault exception msg\rTraceback (most recent call last):\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  # noqa: F401
            except ImportError:
                raise bootstrap.FaultException(
                    None,
                    None,
                    traceback.format_list(
                        [
                            ("spam.py", 3, "<module>", "spam.eggs()"),
                            ("eggs.py", 42, "eggs", 'return "bacon"'),
                        ]
                    ),
                )

        logging.getLogger().addHandler(logging.StreamHandler(mock_stdout))

        bootstrap.handle_event_request(
            self.lambda_runtime,
            raise_exception_handler,
            "invoke_id",
            self.event_body,
            "application/json",
            {},
            {},
            "invoked_function_arn",
            0,
            "tenant_id",
            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)
    def test_handle_event_request_fault_exception_logging_in_json(self, mock_stdout):
        def raise_exception_handler(json_input, lambda_context):
            try:
                import invalid_module  # noqa: F401
            except ImportError:
                raise bootstrap.FaultException("FaultExceptionType", None, None)

        logging_handler = logging.StreamHandler(mock_stdout)
        logging_handler.setFormatter(JsonFormatter())
        logging.getLogger().addHandler(logging_handler)

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

        stdout_value = mock_stdout.getvalue()

        # this line is not in json because of the way the test runtime is bootstrapped
        error_logs = "[ERROR] FaultExceptionType\rTraceback (most recent call last):\n"

        self.assertEqual(stdout_value, 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"
        with self.assertRaises(FaultException) as cm:
            response_handler = bootstrap._get_handler(handler_name)
        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"
        with self.assertRaises(FaultException) as cm:
            response_handler = bootstrap._get_handler(handler_name)
        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):
        importlib.invalidate_caches()
        with tempfile.NamedTemporaryFile(
            suffix=".py", dir=".", delete=False
        ) as tmp_file:
            tmp_file.write(
                b"def syntax_error()\n\tprint('syntax error, no colon after function')"
            )
            tmp_file.flush()

            filename_w_ext = os.path.basename(tmp_file.name)
            filename, _ = os.path.splitext(filename_w_ext)
            handler_name = "{}.syntax_error".format(filename)

            with self.assertRaises(FaultException) as cm:
                response_handler = bootstrap._get_handler(handler_name)
            returned_exception = cm.exception
            self.assertEqual(
                self.FaultExceptionMatcher(
                    "Syntax error in",
                    "Runtime.UserCodeSyntaxError",
                    ".*File.*\\.py.*Line 1.*",
                ),
                returned_exception,
            )

    def test_get_event_handler_missing_error(self):
        importlib.invalidate_caches()
        with tempfile.NamedTemporaryFile(
            suffix=".py", dir=".", delete=False
        ) as tmp_file:
            tmp_file.write(b"def wrong_handler_name():\n\tprint('hello')")
            tmp_file.flush()

            filename_w_ext = os.path.basename(tmp_file.name)
            filename, _ = os.path.splitext(filename_w_ext)
            handler_name = "{}.my_handler".format(filename)
            with self.assertRaises(FaultException) as cm:
                response_handler = bootstrap._get_handler(handler_name)
            returned_exception = cm.exception
            self.assertEqual(
                self.FaultExceptionMatcher(
                    "Handler 'my_handler' missing on module '{}'".format(filename),
                    "Runtime.HandlerNotFound",
                ),
                returned_exception,
            )

    def test_get_event_handler_slash(self):
        importlib.invalidate_caches()
        handler_name = "tests/test_handler_with_slash/test_handler.my_handler"
        response_handler = bootstrap._get_handler(handler_name)
        response_handler()

    def test_get_event_handler_build_in_conflict(self):
        with self.assertRaises(FaultException) as cm:
            response_handler = bootstrap._get_handler("sys.hello")
        returned_exception = cm.exception
        self.assertEqual(
            self.FaultExceptionMatcher(
                "Cannot use built-in module sys as a handler module",
                "Runtime.BuiltInModuleConflict",
            ),
            returned_exception,
        )

    def test_get_event_handler_doesnt_throw_build_in_module_name_slash(self):
        response_handler = bootstrap._get_handler(
            "tests/test_built_in_module_name/sys.my_handler"
        )
        response_handler()

    def test_get_event_handler_doent_throw_build_in_module_name(self):
        response_handler = bootstrap._get_handler(
            "tests.test_built_in_module_name.sys.my_handler"
        )
        response_handler()


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,
            tenant_id=None,
            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,
            tenant_id=None,
            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,
            tenant_id=None,
            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,
            tenant_id=None,
            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\rTraceback (most recent call last):\n"
        )
        self.assertEqual(mock_stdout.getvalue(), expected_logged_error)

    def test_log_error_framed_log_sink(self):
        with NamedTemporaryFile() as temp_file:
            before = int(time.time_ns() / 1000)
            with bootstrap.FramedTelemetryLogSink(
                os.open(temp_file.name, os.O_CREAT | os.O_RDWR)
            ) as log_sink:
                err_to_log = bootstrap.make_error("Error message", "ErrorType", None)
                bootstrap.log_error(err_to_log, log_sink)
            after = int(time.time_ns() / 1000)

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

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

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

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

                timestamp = int.from_bytes(content[8:16], "big")
                self.assertTrue(before <= timestamp)
                self.assertTrue(timestamp <= after)

                actual_message = content[16:].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:
            before = int(time.time_ns() / 1000)
            with bootstrap.FramedTelemetryLogSink(
                os.open(temp_file.name, os.O_CREAT | os.O_RDWR)
            ) as log_sink:
                err_to_log = bootstrap.make_error(
                    "Error message", "ErrorType", ["  line1  ", "  line2  ", "  "]
                )
                bootstrap.log_error(err_to_log, log_sink)
            after = int(time.time_ns() / 1000)

            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, 0xA55A0017)

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

                timestamp = int.from_bytes(content[8:16], "big")
                self.assertTrue(before <= timestamp)
                self.assertTrue(timestamp <= after)

                actual_message = content[16:].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:
            before = int(time.time_ns() / 1000)
            with bootstrap.FramedTelemetryLogSink(
                os.open(temp_file.name, os.O_CREAT | os.O_RDWR)
            ) as log_sink:
                err_to_log = bootstrap.make_error(
                    "Error message", "ErrorType", ["line1", "", "line2"]
                )
                bootstrap.log_error(err_to_log, log_sink)
            after = int(time.time_ns() / 1000)

            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, 0xA55A0017)

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

                timestamp = int.from_bytes(content[8:16], "big")
                self.assertTrue(before <= timestamp)
                self.assertTrue(timestamp <= after)

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

    # Just to ensure we are not logging the requestId from error response, just sending in the response
    def test_log_error_invokeId_line_framed_log_sink(self):
        with NamedTemporaryFile() as temp_file:
            before = int(time.time_ns() / 1000)
            with bootstrap.FramedTelemetryLogSink(
                os.open(temp_file.name, os.O_CREAT | os.O_RDWR)
            ) as log_sink:
                err_to_log = bootstrap.make_error(
                    "Error message",
                    "ErrorType",
                    ["line1", "", "line2"],
                    "testrequestId",
                )
                bootstrap.log_error(err_to_log, log_sink)
            after = int(time.time_ns() / 1000)

            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, 0xA55A0017)

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

                timestamp = int.from_bytes(content[8:16], "big")
                self.assertTrue(before <= timestamp)
                self.assertTrue(timestamp <= after)

                actual_message = content[16:].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 = 3
        os.environ["_LAMBDA_TELEMETRY_LOG_FD"] = "3"

        actual = bootstrap.create_log_sink()

        self.assertIsInstance(actual, bootstrap.FramedTelemetryLogSink)
        self.assertEqual(actual.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"
            before = int(time.time_ns() / 1000)
            with bootstrap.FramedTelemetryLogSink(
                os.open(temp_file.name, os.O_CREAT | os.O_RDWR)
            ) as ls:
                ls.log(message)
            after = int(time.time_ns() / 1000)
            with open(temp_file.name, "rb") as f:
                content = f.read()

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

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

                timestamp = int.from_bytes(content[8:16], "big")
                self.assertTrue(before <= timestamp)
                self.assertTrue(timestamp <= after)

                actual_message = content[16:].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"

            before = int(time.time_ns() / 1000)
            with bootstrap.FramedTelemetryLogSink(
                os.open(temp_file.name, os.O_CREAT | os.O_RDWR)
            ) as ls:
                ls.log(first_message)
                ls.log(second_message)
            after = int(time.time_ns() / 1000)

            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, 0xA55A0003)
                    pos += 4

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

                    timestamp = int.from_bytes(content[pos : pos + 8], "big")
                    self.assertTrue(before <= timestamp)
                    self.assertTrue(timestamp <= after)
                    pos += 8

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

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


class TestLoggingSetup(unittest.TestCase):
    def test_log_level(self) -> None:
        test_cases = [
            (LogFormat.JSON, "TRACE", logging.DEBUG),
            (LogFormat.JSON, "DEBUG", logging.DEBUG),
            (LogFormat.JSON, "INFO", logging.INFO),
            (LogFormat.JSON, "WARN", logging.WARNING),
            (LogFormat.JSON, "ERROR", logging.ERROR),
            (LogFormat.JSON, "FATAL", logging.CRITICAL),
            (LogFormat.TEXT, "TRACE", logging.DEBUG),
            (LogFormat.TEXT, "DEBUG", logging.DEBUG),
            (LogFormat.TEXT, "INFO", logging.INFO),
            (LogFormat.TEXT, "WARN", logging.WARN),
            (LogFormat.TEXT, "ERROR", logging.ERROR),
            (LogFormat.TEXT, "FATAL", logging.CRITICAL),
            ("Unknown format", "INFO", logging.INFO),
            # if level is unknown fall back to default
            (LogFormat.JSON, "Unknown level", logging.NOTSET),
        ]
        for fmt, log_level, expected_level in test_cases:
            with self.subTest():
                # Drop previous setup
                logging.getLogger().handlers.clear()
                logging.getLogger().level = logging.NOTSET

                bootstrap._setup_logging(
                    fmt,
                    _get_log_level_from_env_var(log_level),
                    bootstrap.StandardLogSink(),
                )

                self.assertEqual(expected_level, logging.getLogger().level)


class TestLambdaLoggerHandlerSetup(unittest.TestCase):
    @classmethod
    def tearDownClass(cls):
        importlib.reload(bootstrap)
        logging.getLogger().handlers.clear()
        logging.getLogger().level = logging.NOTSET

    def test_handler_setup(self, *_):
        test_cases = [
            (62, 0xA55A0003, 46, {}),
            (133, 0xA55A001A, 117, {"AWS_LAMBDA_LOG_FORMAT": "JSON"}),
            (62, 0xA55A001B, 46, {"AWS_LAMBDA_LOG_LEVEL": "INFO"}),
        ]

        for total_length, header, message_length, env_vars in test_cases:
            with patch.dict(
                os.environ, env_vars, clear=True
            ), NamedTemporaryFile() as temp_file:
                importlib.reload(bootstrap)
                logging.getLogger().handlers.clear()
                logging.getLogger().level = logging.NOTSET

                before = int(time.time_ns() / 1000)
                with bootstrap.FramedTelemetryLogSink(
                    os.open(temp_file.name, os.O_CREAT | os.O_RDWR)
                ) as ls:
                    bootstrap._setup_logging(
                        bootstrap._AWS_LAMBDA_LOG_FORMAT,
                        bootstrap._AWS_LAMBDA_LOG_LEVEL,
                        ls,
                    )
                    logger = logging.getLogger()
                    logger.critical("critical")
                after = int(time.time_ns() / 1000)

                content = open(temp_file.name, "rb").read()
                self.assertEqual(len(content), total_length)

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

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

                timestamp = int.from_bytes(content[pos : pos + 8], "big")
                self.assertTrue(before <= timestamp)
                self.assertTrue(timestamp <= after)


class TestLogging(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        logging.getLogger().handlers.clear()
        logging.getLogger().level = logging.NOTSET
        bootstrap._setup_logging(
            LogFormat.from_str("JSON"), "INFO", bootstrap.StandardLogSink()
        )

    @patch("sys.stderr", new_callable=StringIO)
    def test_json_formatter(self, mock_stderr):
        logger = logging.getLogger("a.b")

        test_cases = [
            (
                logging.ERROR,
                "TEST 1",
                {
                    "level": "ERROR",
                    "logger": "a.b",
                    "message": "TEST 1",
                    "requestId": "",
                },
            ),
            (
                logging.ERROR,
                "test \nwith \nnew \nlines",
                {
                    "level": "ERROR",
                    "logger": "a.b",
                    "message": "test \nwith \nnew \nlines",
                    "requestId": "",
                },
            ),
            (
                logging.CRITICAL,
                "TEST CRITICAL",
                {
                    "level": "CRITICAL",
                    "logger": "a.b",
                    "message": "TEST CRITICAL",
                    "requestId": "",
                },
            ),
        ]
        for level, msg, expected in test_cases:
            with self.subTest(msg):
                with patch("sys.stdout", new_callable=StringIO) as mock_stdout:
                    logger.log(level, msg)

                    data = json.loads(mock_stdout.getvalue())
                    data.pop("timestamp")
                    self.assertEqual(
                        data,
                        expected,
                    )
        self.assertEqual(mock_stderr.getvalue(), "")

    @patch("awslambdaric.bootstrap._GLOBAL_TENANT_ID", "test-tenant-id")
    @patch("sys.stderr", new_callable=StringIO)
    def test_json_formatter_with_tenant_id(self, mock_stderr):
        logger = logging.getLogger("a.b")
        level = logging.INFO
        message = "Test json formatting with tenant id"
        expected = {
            "level": "INFO",
            "logger": "a.b",
            "message": message,
            "requestId": "",
            "tenantId": "test-tenant-id",
        }

        with patch("sys.stdout", new_callable=StringIO) as mock_stdout:
            logger.log(level, message)

            data = json.loads(mock_stdout.getvalue())
            data.pop("timestamp")
            self.assertEqual(
                data,
                expected,
            )
        self.assertEqual(mock_stderr.getvalue(), "")

    @patch("sys.stdout", new_callable=StringIO)
    @patch("sys.stderr", new_callable=StringIO)
    def test_exception(self, mock_stderr, mock_stdout):
        try:
            raise ValueError("error message")
        except ValueError:
            logging.getLogger("test.logger").exception("test exception")

        exception_log = json.loads(mock_stdout.getvalue())
        self.assertIn("location", exception_log)
        self.assertIn("stackTrace", exception_log)
        exception_log.pop("timestamp")
        exception_log.pop("location")
        stack_trace = exception_log.pop("stackTrace")

        self.assertEqual(len(stack_trace), 1)

        self.assertEqual(
            exception_log,
            {
                "errorMessage": "error message",
                "errorType": "ValueError",
                "level": "ERROR",
                "logger": "test.logger",
                "message": "test exception",
                "requestId": "",
            },
        )

        self.assertEqual(mock_stderr.getvalue(), "")

    @patch("sys.stdout", new_callable=StringIO)
    @patch("sys.stderr", new_callable=StringIO)
    def test_log_level(self, mock_stderr, mock_stdout):
        logger = logging.getLogger("test.logger")

        logger.debug("debug message")
        logger.info("info message")

        data = json.loads(mock_stdout.getvalue())
        data.pop("timestamp")

        self.assertEqual(
            data,
            {
                "level": "INFO",
                "logger": "test.logger",
                "message": "info message",
                "requestId": "",
            },
        )
        self.assertEqual(mock_stderr.getvalue(), "")

    @patch("sys.stdout", new_callable=StringIO)
    @patch("sys.stderr", new_callable=StringIO)
    def test_set_log_level_manually(self, mock_stderr, mock_stdout):
        logger = logging.getLogger("test.logger")

        # Changing log level after `bootstrap.setup_logging`
        logging.getLogger().setLevel(logging.CRITICAL)

        logger.debug("debug message")
        logger.info("info message")
        logger.warning("warning message")
        logger.error("error message")
        logger.critical("critical message")

        data = json.loads(mock_stdout.getvalue())
        data.pop("timestamp")

        self.assertEqual(
            data,
            {
                "level": "CRITICAL",
                "logger": "test.logger",
                "message": "critical message",
                "requestId": "",
            },
        )
        self.assertEqual(mock_stderr.getvalue(), "")

    @patch("sys.stdout", new_callable=StringIO)
    @patch("sys.stderr", new_callable=StringIO)
    def test_set_log_level_with_dictConfig(self, mock_stderr, mock_stdout):
        # Changing log level after `bootstrap.setup_logging`
        logging.config.dictConfig(
            {
                "version": 1,
                "disable_existing_loggers": False,
                "formatters": {"simple": {"format": "%(levelname)-8s - %(message)s"}},
                "handlers": {
                    "stdout": {
                        "class": "logging.StreamHandler",
                        "formatter": "simple",
                    },
                },
                "root": {
                    "level": "CRITICAL",
                    "handlers": [
                        "stdout",
                    ],
                },
            }
        )

        logger = logging.getLogger("test.logger")
        logger.debug("debug message")
        logger.info("info message")
        logger.warning("warning message")
        logger.error("error message")
        logger.critical("critical message")

        data = mock_stderr.getvalue()
        self.assertEqual(
            data,
            "CRITICAL - critical message\n",
        )
        self.assertEqual(mock_stdout.getvalue(), "")


class TestBootstrapModule(unittest.TestCase):
    @patch("awslambdaric.bootstrap.LambdaRuntimeClient")
    def test_run(self, mock_runtime_client):
        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(SystemExit) as cm:
            bootstrap.run(
                expected_app_root, expected_handler, expected_lambda_runtime_api_addr
            )

        self.assertEqual(cm.exception.code, 1)

    @patch(
        "awslambdaric.bootstrap.LambdaLoggerHandler",
        Mock(side_effect=Exception("Boom!")),
    )
    @patch("awslambdaric.bootstrap.build_fault_result")
    @patch("awslambdaric.bootstrap.log_error", MagicMock())
    @patch("awslambdaric.bootstrap.LambdaRuntimeClient", MagicMock())
    @patch("awslambdaric.bootstrap.sys")
    def test_run_exception(self, mock_sys, mock_build_fault_result):
        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_build_fault_result.return_value = {}
        mock_sys.exit.side_effect = TestException("Boom!")

        with self.assertRaises(TestException):
            bootstrap.run(
                expected_app_root, expected_handler, expected_lambda_runtime_api_addr
            )

        mock_sys.exit.assert_called_once_with(1)


class TestOnInitComplete(unittest.TestCase):
    def tearDown(self):
        # We are accessing private filed for cleaning up
        snapshot_restore_py._before_snapshot_registry = []
        snapshot_restore_py._after_restore_registry = []

    # We are using ANY over here as the main thing we want to test is teh errorType propogation and stack trace generation
    error_result = {
        "errorMessage": "This is a Dummy type error",
        "errorType": "TypeError",
        "requestId": "",
        "stackTrace": ANY,
    }

    def raise_type_error(self):
        raise TypeError("This is a Dummy type error")

    @patch("awslambdaric.bootstrap.LambdaRuntimeClient")
    def test_before_snapshot_exception(self, mock_runtime_client):
        snapshot_restore_py.register_before_snapshot(self.raise_type_error)

        with self.assertRaises(SystemExit) as cm:
            bootstrap.on_init_complete(
                mock_runtime_client, log_sink=bootstrap.StandardLogSink()
            )

        self.assertEqual(cm.exception.code, 64)
        mock_runtime_client.post_init_error.assert_called_once_with(
            self.error_result,
            FaultException.BEFORE_SNAPSHOT_ERROR,
        )

    @patch("awslambdaric.bootstrap.LambdaRuntimeClient")
    def test_after_restore_exception(self, mock_runtime_client):
        snapshot_restore_py.register_after_restore(self.raise_type_error)

        with self.assertRaises(SystemExit) as cm:
            bootstrap.on_init_complete(
                mock_runtime_client, log_sink=bootstrap.StandardLogSink()
            )

        self.assertEqual(cm.exception.code, 65)
        mock_runtime_client.restore_next.assert_called_once()
        mock_runtime_client.report_restore_error.assert_called_once_with(
            self.error_result
        )


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