test_logger.py

import functools
import inspect
import io
import json
import logging
import random
import re
import secrets
import string
import sys
from collections import namedtuple
from datetime import datetime, timezone
from typing import Any, Callable, Dict, Iterable, List, Optional, Union

import pytest

from aws_lambda_powertools import Logger
from aws_lambda_powertools.logging import correlation_paths
from aws_lambda_powertools.logging.exceptions import InvalidLoggerSamplingRateError
from aws_lambda_powertools.logging.formatter import (
    BasePowertoolsFormatter,
    LambdaPowertoolsFormatter,
)
from aws_lambda_powertools.shared import constants
from aws_lambda_powertools.utilities.data_classes import S3Event, event_source


@pytest.fixture
def stdout():
    return io.StringIO()


@pytest.fixture
def lambda_context():
    lambda_context = {
        "function_name": "test",
        "memory_limit_in_mb": 128,
        "invoked_function_arn": "arn:aws:lambda:eu-west-1:809313241:function:test",
        "aws_request_id": "52fdfc07-2182-154f-163f-5f0f9a621d72",
    }

    return namedtuple("LambdaContext", lambda_context.keys())(*lambda_context.values())


@pytest.fixture
def lambda_event():
    return {"greeting": "hello"}


@pytest.fixture
def service_name():
    chars = string.ascii_letters + string.digits
    return "".join(random.SystemRandom().choice(chars) for _ in range(15))


def capture_logging_output(stdout):
    return json.loads(stdout.getvalue().strip())


def capture_multiple_logging_statements_output(stdout):
    return [json.loads(line.strip()) for line in stdout.getvalue().split("\n") if line]


def test_setup_service_name(stdout, service_name):
    # GIVEN Logger is initialized
    # WHEN service is explicitly defined
    logger = Logger(service=service_name, stream=stdout)

    logger.info("Hello")

    # THEN service field should be equals service given
    log = capture_logging_output(stdout)
    assert service_name == log["service"]


def test_setup_service_env_var(monkeypatch, stdout, service_name):
    # GIVEN Logger is initialized
    # WHEN service is explicitly defined via POWERTOOLS_SERVICE_NAME env
    monkeypatch.setenv("POWERTOOLS_SERVICE_NAME", service_name)
    logger = Logger(stream=stdout)

    logger.info("Hello")

    # THEN service field should be equals POWERTOOLS_SERVICE_NAME value
    log = capture_logging_output(stdout)
    assert service_name == log["service"]


def test_setup_sampling_rate_env_var(monkeypatch, stdout, service_name):
    # GIVEN Logger is initialized
    # WHEN samping rate is explicitly set to 100% via POWERTOOLS_LOGGER_SAMPLE_RATE env
    sampling_rate = "1"
    monkeypatch.setenv("POWERTOOLS_LOGGER_SAMPLE_RATE", sampling_rate)
    logger = Logger(service=service_name, stream=stdout)
    logger.debug("I am being sampled")

    # THEN sampling rate should be equals POWERTOOLS_LOGGER_SAMPLE_RATE value
    # log level should be DEBUG
    # and debug log statements should be in stdout
    log = capture_logging_output(stdout)
    assert sampling_rate == log["sampling_rate"]
    assert "DEBUG" == log["level"]
    assert "I am being sampled" == log["message"]


def test_inject_lambda_context(lambda_context, stdout, service_name):
    # GIVEN Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # WHEN a lambda function is decorated with logger
    @logger.inject_lambda_context
    def handler(event, context):
        logger.info("Hello")

    handler({}, lambda_context)

    # THEN lambda contextual info should always be in the logs
    log = capture_logging_output(stdout)
    expected_logger_context_keys = (
        "function_name",
        "function_memory_size",
        "function_arn",
        "function_request_id",
    )
    for key in expected_logger_context_keys:
        assert key in log


def test_inject_lambda_context_log_event_request(lambda_context, stdout, lambda_event, service_name):
    # GIVEN Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # WHEN a lambda function is decorated with logger instructed to log event
    @logger.inject_lambda_context(log_event=True)
    def handler(event, context):
        logger.info("Hello")

    handler(lambda_event, lambda_context)

    # THEN logger should log event received from Lambda
    logged_event, _ = capture_multiple_logging_statements_output(stdout)
    assert logged_event["message"] == lambda_event


def test_inject_lambda_context_log_event_request_env_var(
    monkeypatch,
    lambda_context,
    stdout,
    lambda_event,
    service_name,
):
    # GIVEN Logger is initialized
    monkeypatch.setenv("POWERTOOLS_LOGGER_LOG_EVENT", "true")
    logger = Logger(service=service_name, stream=stdout)

    # WHEN a lambda function is decorated with logger instructed to log event
    # via POWERTOOLS_LOGGER_LOG_EVENT env
    @logger.inject_lambda_context
    def handler(event, context):
        logger.info("Hello")

    handler(lambda_event, lambda_context)

    # THEN logger should log event received from Lambda
    logged_event, _ = capture_multiple_logging_statements_output(stdout)
    assert logged_event["message"] == lambda_event


def test_inject_lambda_context_log_no_request_by_default(
    monkeypatch,
    lambda_context,
    stdout,
    lambda_event,
    service_name,
):
    # GIVEN Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # WHEN a lambda function is decorated with logger
    @logger.inject_lambda_context
    def handler(event, context):
        logger.info("Hello")

    handler(lambda_event, lambda_context)

    # THEN logger should not log event received by lambda handler
    log = capture_logging_output(stdout)
    assert log["message"] != lambda_event


def test_inject_lambda_cold_start(lambda_context, stdout, service_name):
    # cold_start can be false as it's a global variable in Logger module
    # so we reset it to simulate the correct behaviour
    # since Lambda will only import our logger lib once per concurrent execution
    from aws_lambda_powertools.logging import logger

    logger.is_cold_start = True

    # GIVEN Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # WHEN a lambda function is decorated with logger, and called twice
    @logger.inject_lambda_context
    def handler(event, context):
        logger.info("Hello")

    handler({}, lambda_context)
    handler({}, lambda_context)

    # THEN cold_start key should only be true in the first call
    first_log, second_log = capture_multiple_logging_statements_output(stdout)
    assert first_log["cold_start"] is True
    assert second_log["cold_start"] is False


def test_logger_append_duplicated(stdout, service_name):
    # GIVEN Logger is initialized with request_id field
    logger = Logger(service=service_name, stream=stdout, request_id="value")

    # WHEN `request_id` is appended to the existing structured log
    # using a different value
    logger.structure_logs(append=True, request_id="new_value")
    logger.info("log")

    # THEN subsequent log statements should have the latest value
    log = capture_logging_output(stdout)
    assert "new_value" == log["request_id"]


def test_logger_honors_given_exception_keys(stdout, service_name):
    # GIVEN Logger is initialized with exception and exception_name fields
    logger = Logger(service=service_name, stream=stdout)

    # WHEN log level info
    logger.info("log", exception="exception_value", exception_name="exception_name_value")

    # THEN log statements should have these keys
    log = capture_logging_output(stdout)
    assert "exception_value" == log["exception"]
    assert "exception_name_value" == log["exception_name"]


def test_logger_invalid_sampling_rate(service_name):
    # GIVEN Logger is initialized
    # WHEN sampling_rate non-numeric value
    # THEN we should raise InvalidLoggerSamplingRateError
    with pytest.raises(InvalidLoggerSamplingRateError):
        Logger(service=service_name, stream=stdout, sampling_rate="TEST")


def test_inject_lambda_context_with_structured_log(lambda_context, stdout, service_name):
    # GIVEN Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # WHEN structure_logs has been used to add an additional key upfront
    # and a lambda function is decorated with logger.inject_lambda_context
    logger.structure_logs(append=True, additional_key="test")

    @logger.inject_lambda_context
    def handler(event, context):
        logger.info("Hello")

    handler({}, lambda_context)

    # THEN lambda contextual info should always be in the logs
    log = capture_logging_output(stdout)
    expected_logger_context_keys = (
        "function_name",
        "function_memory_size",
        "function_arn",
        "function_request_id",
        "additional_key",
    )
    for key in expected_logger_context_keys:
        assert key in log


def test_logger_children_propagate_changes(stdout, service_name):
    # GIVEN Loggers are initialized
    # create child logger before parent to mimick
    # importing logger from another module/file
    # as loggers are created in global scope
    child = Logger(stream=stdout, service=service_name, child=True)
    parent = Logger(stream=stdout, service=service_name)

    # WHEN a child Logger adds an additional key
    child.structure_logs(append=True, customer_id="value")

    # THEN child Logger changes should propagate to parent
    # and subsequent log statements should have the latest value
    parent.info("Hello parent")
    child.info("Hello child")

    parent_log, child_log = capture_multiple_logging_statements_output(stdout)
    assert "customer_id" in parent_log
    assert "customer_id" in child_log
    assert child.parent.name == service_name


def test_logger_child_not_set_returns_same_logger(stdout):
    # GIVEN two Loggers are initialized with the same service name
    # WHEN child param isn't set
    logger_one = Logger(service="something", stream=stdout)
    logger_two = Logger(service="something", stream=stdout)

    # THEN we should have two Logger instances
    # however inner logger wise should be the same
    assert id(logger_one) != id(logger_two)
    assert logger_one._logger is logger_two._logger
    assert logger_one.name is logger_two.name

    # THEN we should also not see any duplicated logs
    logger_one.info("One - Once")
    logger_two.info("Two - Once")

    logs = list(capture_multiple_logging_statements_output(stdout))
    assert len(logs) == 2


def test_logger_level_case_insensitive(service_name):
    # GIVEN a Loggers is initialized
    # WHEN log level is set as "info" instead of "INFO"
    logger = Logger(service=service_name, level="info")

    # THEN we should correctly set log level as INFO
    assert logger.level == logging.INFO


def test_logger_level_not_set(service_name):
    # GIVEN a Loggers is initialized
    # WHEN no log level was passed
    logger = Logger(service=service_name)

    # THEN we should default to INFO
    assert logger.level == logging.INFO


def test_logger_level_as_int(service_name):
    # GIVEN a Loggers is initialized
    # WHEN log level is int
    logger = Logger(service=service_name, level=logging.INFO)

    # THEN we should be expected int (20, in this case)
    assert logger.level == logging.INFO


def test_logger_level_env_var_as_int(monkeypatch, service_name):
    # GIVEN Logger is initialized
    # WHEN log level is explicitly defined via POWERTOOLS_LOG_LEVEL env as int
    # THEN Logger should propagate ValueError
    # since env vars can only be string
    # and '50' is not a correct log level
    monkeypatch.setenv(constants.POWERTOOLS_LOG_LEVEL_ENV, 50)
    with pytest.raises(ValueError, match="Unknown level: '50'"):
        Logger(service=service_name)


def test_logger_level_env_var_legacy_as_int(monkeypatch, service_name):
    # GIVEN Logger is initialized
    # WHEN log level is explicitly defined via legacy env LOG_LEVEL env as int
    # THEN Logger should propagate ValueError
    # since env vars can only be string
    # and '50' is not a correct log level
    monkeypatch.setenv(constants.POWERTOOLS_LOG_LEVEL_LEGACY_ENV, 50)
    with pytest.raises(ValueError, match="Unknown level: '50'"):
        Logger(service=service_name)


def test_logger_switch_between_levels(stdout, service_name):
    # GIVEN a Loggers is initialized with INFO level
    logger = Logger(service=service_name, level="INFO", stream=stdout)
    logger.info("message info")

    # WHEN we switch to DEBUG level
    logger.setLevel(level="DEBUG")
    logger.debug("message debug")

    # THEN we must have different levels and messages in stdout
    log_output = capture_multiple_logging_statements_output(stdout)
    assert log_output[0]["level"] == "INFO"
    assert log_output[0]["message"] == "message info"
    assert log_output[1]["level"] == "DEBUG"
    assert log_output[1]["message"] == "message debug"


def test_logger_record_caller_location(stdout, service_name):
    # GIVEN Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # WHEN log statement is run
    logger.info("log")

    # THEN 'location' field should have
    # the correct caller resolution
    caller_fn_name = inspect.currentframe().f_code.co_name
    log = capture_logging_output(stdout)
    assert caller_fn_name in log["location"]


def test_logger_do_not_log_twice_when_root_logger_is_setup(stdout, service_name):
    # GIVEN Lambda configures the root logger with a handler
    root_logger = logging.getLogger()
    root_logger.addHandler(logging.StreamHandler(stream=stdout))

    # WHEN we create a new Logger and child Logger
    logger = Logger(service=service_name, stream=stdout)
    child_logger = Logger(service=service_name, child=True, stream=stdout)
    logger.info("PARENT")
    child_logger.info("CHILD")
    root_logger.info("ROOT")

    # THEN it should only contain only two log entries
    # since child's log records propagated to root logger should be rejected
    logs = list(stdout.getvalue().strip().split("\n"))
    assert len(logs) == 2


def test_logger_extra_kwargs(stdout, service_name):
    # GIVEN Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # WHEN `request_id` is an extra field in a log message to the existing structured log
    fields = {"request_id": "blah"}

    logger.info("with extra fields", extra=fields)
    logger.info("without extra fields")

    extra_fields_log, no_extra_fields_log = capture_multiple_logging_statements_output(stdout)

    # THEN first log should have request_id field in the root structure
    assert "request_id" in extra_fields_log

    # THEN second log should not have request_id in the root structure
    assert "request_id" not in no_extra_fields_log


def test_logger_arbitrary_fields_as_kwargs(stdout, service_name):
    # GIVEN Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # WHEN `request_id` is an arbitrary field in a log message to the existing structured log
    fields = {"request_id": "blah"}

    logger.info("with arbitrary fields", **fields)
    logger.info("without extra fields")

    extra_fields_log, no_extra_fields_log = capture_multiple_logging_statements_output(stdout)

    # THEN first log should have request_id field in the root structure
    assert "request_id" in extra_fields_log

    # THEN second log should not have request_id in the root structure
    assert "request_id" not in no_extra_fields_log


def test_logger_log_twice_when_log_filter_isnt_present_and_root_logger_is_setup(monkeypatch, stdout, service_name):
    # GIVEN Lambda configures the root logger with a handler
    root_logger = logging.getLogger()
    root_logger.addHandler(logging.StreamHandler(stream=stdout))

    # WHEN we create a new Logger and child Logger
    # and log deduplication filter for child messages are disabled
    # see #262 for more details on why this is needed for Pytest Live Log feature
    monkeypatch.setenv(constants.LOGGER_LOG_DEDUPLICATION_ENV, "true")
    logger = Logger(service=service_name, stream=stdout)
    child_logger = Logger(service=service_name, child=True, stream=stdout)
    logger.info("PARENT")
    child_logger.info("CHILD")

    # THEN it should only contain only two log entries
    # since child's log records propagated to root logger should be rejected
    logs = list(stdout.getvalue().strip().split("\n"))
    assert len(logs) == 4


def test_logger_exception_extract_exception_name(stdout, service_name):
    # GIVEN Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # WHEN calling a logger.exception with a ValueError
    try:
        raise ValueError("something went wrong")
    except Exception:
        logger.exception("Received an exception")

    # THEN we expect a "exception_name" to be "ValueError"
    log = capture_logging_output(stdout)
    assert "ValueError" == log["exception_name"]


def test_logger_set_correlation_id(lambda_context, stdout, service_name):
    # GIVEN
    logger = Logger(service=service_name, stream=stdout)
    request_id = "xxx-111-222"
    mock_event = {"requestContext": {"requestId": request_id}}

    def handler(event, _):
        logger.set_correlation_id(event["requestContext"]["requestId"])
        logger.info("Foo")

    # WHEN
    handler(mock_event, lambda_context)

    # THEN
    log = capture_logging_output(stdout)
    assert request_id == log["correlation_id"]


def test_logger_get_correlation_id(lambda_context, stdout, service_name):
    # GIVEN a logger with a correlation_id set
    logger = Logger(service=service_name, stream=stdout)
    logger.set_correlation_id("foo")

    # WHEN calling get_correlation_id
    correlation_id = logger.get_correlation_id()

    # THEN it should return the correlation_id
    assert "foo" == correlation_id


def test_logger_set_correlation_id_path(lambda_context, stdout, service_name):
    # GIVEN
    logger = Logger(service=service_name, stream=stdout)
    request_id = "xxx-111-222"
    mock_event = {"requestContext": {"requestId": request_id}}

    @logger.inject_lambda_context(correlation_id_path=correlation_paths.API_GATEWAY_REST)
    def handler(event, context):
        logger.info("Foo")

    # WHEN
    handler(mock_event, lambda_context)

    # THEN
    log = capture_logging_output(stdout)
    assert request_id == log["correlation_id"]


def test_logger_set_correlation_id_path_custom_functions(lambda_context, stdout, service_name):
    # GIVEN an initialized Logger
    # AND a Lambda handler decorated with a JMESPath expression using Powertools custom functions
    logger = Logger(service=service_name, stream=stdout)

    @logger.inject_lambda_context(correlation_id_path="Records[*].powertools_json(body).id")
    def handler(event, context): ...

    # WHEN handler is called
    request_id = "xxx-111-222"
    mock_event = {"Records": [{"body": json.dumps({"id": request_id})}]}
    handler(mock_event, lambda_context)

    # THEN there should be no exception and correlation ID should match
    assert logger.get_correlation_id() == [request_id]


def test_logger_append_remove_keys(stdout, service_name):
    # GIVEN a Logger is initialized
    logger = Logger(service=service_name, stream=stdout)
    extra_keys = {"request_id": "id", "context": "value"}

    # WHEN keys are updated
    logger.append_keys(**extra_keys)
    logger.info("message with new keys")

    # And removed
    logger.remove_keys(extra_keys.keys())
    logger.info("message after keys being removed")

    # THEN additional keys should only be present in the first log statement
    extra_keys_log, keys_removed_log = capture_multiple_logging_statements_output(stdout)

    assert extra_keys.items() <= extra_keys_log.items()
    assert (extra_keys.items() <= keys_removed_log.items()) is False


def test_logger_append_and_show_current_keys(stdout, service_name):
    # GIVEN a Logger is initialized
    logger = Logger(service=service_name, stream=stdout)
    extra_keys = {"request_id": "id", "context": "value"}

    # WHEN keys are updated
    logger.append_keys(**extra_keys)

    # THEN appended keys must be present in logger
    current_keys = logger.get_current_keys()
    assert "request_id" in current_keys
    assert "context" in current_keys


def test_logger_formatter_without_get_current_keys_method(stdout, service_name):
    class CustomFormatter(BasePowertoolsFormatter):
        def append_keys(self, **additional_keys):
            # Fake method
            pass

        def clear_state(self) -> None:
            # Fake method
            pass

    custom_formater = CustomFormatter()

    # GIVEN a Logger is initialized with a Logger Formatter from scratch

    logger = Logger(service=service_name, stream=stdout, logger_formatter=custom_formater)
    extra_keys = {"request_id": "id", "context": "value"}

    # WHEN keys are updated
    logger.append_keys(**extra_keys)

    # THEN the appended keys will not persist
    # unless the customer implements the required methods and persists the log_format
    assert logger.get_current_keys() == {}


def test_logger_custom_formatter(stdout, service_name, lambda_context):
    class CustomFormatter(BasePowertoolsFormatter):
        custom_format = {}

        def append_keys(self, **additional_keys):
            self.custom_format.update(additional_keys)

        def remove_keys(self, keys: Iterable[str]):
            for key in keys:
                self.custom_format.pop(key, None)

        def clear_state(self):
            self.custom_format.clear()

        def format(self, record: logging.LogRecord) -> str:  # noqa: A003
            return json.dumps(
                {
                    "message": super().format(record),
                    "timestamp": self.formatTime(record),
                    "my_default_key": "test",
                    **self.custom_format,
                },
            )

    custom_formatter = CustomFormatter()

    # GIVEN a Logger is initialized with a custom formatter
    logger = Logger(service=service_name, stream=stdout, logger_formatter=custom_formatter)

    # WHEN a lambda function is decorated with logger
    @logger.inject_lambda_context(correlation_id_path="foo")
    def handler(event, context):
        logger.info("Hello")

    handler({"foo": "value"}, lambda_context)

    lambda_context_keys = (
        "function_name",
        "function_memory_size",
        "function_arn",
        "function_request_id",
    )

    log = capture_logging_output(stdout)

    # THEN custom key should always be present
    # and lambda contextual info should also be in the logs
    # and get_correlation_id should return None
    assert "my_default_key" in log
    assert all(k in log for k in lambda_context_keys)
    assert log["correlation_id"] == "value"
    assert logger.get_correlation_id() is None


def test_logger_custom_powertools_formatter_clear_state(stdout, service_name, lambda_context):
    class CustomFormatter(LambdaPowertoolsFormatter):
        def __init__(
            self,
            json_serializer: Optional[Callable[[Dict], str]] = None,
            json_deserializer: Optional[Callable[[Union[Dict, str, bool, int, float]], str]] = None,
            json_default: Optional[Callable[[Any], Any]] = None,
            datefmt: Optional[str] = None,
            use_datetime_directive: bool = False,
            log_record_order: Optional[List[str]] = None,
            utc: bool = False,
            **kwargs,
        ):
            super().__init__(
                json_serializer,
                json_deserializer,
                json_default,
                datefmt,
                use_datetime_directive,
                log_record_order,
                utc,
                **kwargs,
            )

    custom_formatter = CustomFormatter()

    # GIVEN a Logger is initialized with a custom formatter
    logger = Logger(service=service_name, stream=stdout, logger_formatter=custom_formatter)

    # WHEN a lambda function is decorated with logger
    # and state is to be cleared in the next invocation
    @logger.inject_lambda_context(clear_state=True)
    def handler(event, context):
        if event.get("add_key"):
            logger.append_keys(my_key="value")
        logger.info("Hello")

    handler({"add_key": True}, lambda_context)
    handler({}, lambda_context)

    lambda_context_keys = (
        "function_name",
        "function_memory_size",
        "function_arn",
        "function_request_id",
    )

    first_log, second_log = capture_multiple_logging_statements_output(stdout)

    # THEN my_key should only present once
    # and lambda contextual info should also be in both logs
    assert "my_key" in first_log
    assert "my_key" not in second_log
    assert all(k in first_log for k in lambda_context_keys)
    assert all(k in second_log for k in lambda_context_keys)


def test_logger_custom_formatter_has_standard_and_custom_keys(stdout, service_name, lambda_context):
    class CustomFormatter(LambdaPowertoolsFormatter): ...

    # GIVEN a Logger is initialized with a custom formatter
    logger = Logger(service=service_name, stream=stdout, logger_formatter=CustomFormatter(), my_key="value")

    # WHEN a lambda function is decorated with logger
    @logger.inject_lambda_context
    def handler(event, context):
        logger.info("Hello")

    handler({}, lambda_context)

    standard_keys = (
        "level",
        "location",
        "message",
        "timestamp",
        "service",
        "cold_start",
        "function_name",
        "function_memory_size",
        "function_arn",
        "function_request_id",
    )

    log = capture_logging_output(stdout)

    # THEN all standard keys should be available
    assert all(k in log for k in standard_keys)
    assert "my_key" in log


def test_logger_custom_handler(lambda_context, service_name, tmp_path):
    # GIVEN a Logger is initialized with a FileHandler
    log_file = tmp_path / "log.json"
    handler = logging.FileHandler(filename=log_file)
    logger = Logger(service=service_name, logger_handler=handler)

    # WHEN a log statement happens
    @logger.inject_lambda_context
    def handler(event, context):
        logger.info("custom handler")

    handler({}, lambda_context)

    # THEN we should output to a file not stdout
    log = log_file.read_text()
    assert "custom handler" in log


def test_clear_state_on_inject_lambda_context(lambda_context, stdout, service_name):
    # GIVEN
    logger = Logger(service=service_name, stream=stdout)

    # WHEN clear_state is set and a key was conditionally added in the first invocation
    @logger.inject_lambda_context(clear_state=True)
    def handler(event, context):
        if event.get("add_key"):
            logger.append_keys(my_key="value")
        logger.info("Foo")

    # THEN custom key should only exist in the first log
    handler({"add_key": True}, lambda_context)
    handler({}, lambda_context)

    first_log, second_log = capture_multiple_logging_statements_output(stdout)
    assert "my_key" in first_log
    assert "my_key" not in second_log


def test_clear_state_keeps_standard_keys(lambda_context, stdout, service_name):
    # GIVEN
    logger = Logger(service=service_name, stream=stdout)
    standard_keys = ["level", "location", "message", "timestamp", "service"]

    # WHEN clear_state is set
    @logger.inject_lambda_context(clear_state=True)
    def handler(event, context):
        logger.info("Foo")

    # THEN all standard keys should be available as usual
    handler({}, lambda_context)
    handler({}, lambda_context)

    first_log, second_log = capture_multiple_logging_statements_output(stdout)
    for key in standard_keys:
        assert key in first_log
        assert key in second_log


def test_clear_state_keeps_custom_keys(lambda_context, stdout, service_name):
    # GIVEN
    location_format = "%(module)s.%(funcName)s:clear_state"
    logger = Logger(service=service_name, stream=stdout, location=location_format, custom_key="foo")

    # WHEN clear_state is set
    @logger.inject_lambda_context(clear_state=True)
    def handler(event, context):
        logger.info("Foo")

    # THEN all standard keys should be available as usual
    handler({}, lambda_context)
    handler({}, lambda_context)

    first_log, second_log = capture_multiple_logging_statements_output(stdout)
    for log in (first_log, second_log):
        assert "foo" == log["custom_key"]
        assert "test_logger.handler:clear_state" == log["location"]


def test_clear_state_keeps_exception_keys(lambda_context, stdout, service_name):
    # GIVEN
    logger = Logger(service=service_name, stream=stdout)

    # WHEN clear_state is set and an exception was logged
    @logger.inject_lambda_context(clear_state=True)
    def handler(event, context):
        try:
            raise ValueError("something went wrong")
        except Exception:
            logger.exception("Received an exception")

    # THEN we expect a "exception_name" to be "ValueError"
    handler({}, lambda_context)
    log = capture_logging_output(stdout)
    assert "ValueError" == log["exception_name"]


def test_inject_lambda_context_allows_handler_with_kwargs(lambda_context, stdout, service_name):
    # GIVEN
    logger = Logger(service=service_name, stream=stdout)

    # WHEN
    @logger.inject_lambda_context(clear_state=True)
    def handler(event, context, my_custom_option=None):
        pass

    # THEN
    handler({}, lambda_context, my_custom_option="blah")


@pytest.mark.parametrize("utc", [False, True])
def test_use_datetime(stdout, service_name, utc):
    # GIVEN
    logger = Logger(
        service=service_name,
        stream=stdout,
        datefmt="custom timestamp: milliseconds=%F microseconds=%f timezone=%z",
        use_datetime_directive=True,
        utc=utc,
    )

    # WHEN a log statement happens
    logger.info({})

    # THEN the timestamp has the appropriate formatting
    log = capture_logging_output(stdout)

    expected_tz = datetime.now().astimezone(timezone.utc if utc else None).strftime("%z")
    assert re.fullmatch(
        f"custom timestamp: milliseconds=[0-9]+ microseconds=[0-9]+ timezone={re.escape(expected_tz)}",
        log["timestamp"],
    )


@pytest.mark.parametrize("utc", [False, True])
def test_use_rfc3339_iso8601(stdout, service_name, utc):
    # GIVEN
    logger = Logger(service=service_name, stream=stdout, use_rfc3339=True, utc=utc)
    RFC3339_REGEX = r"^((?:(\d{4}-\d{2}-\d{2})T(\d{2}:\d{2}:\d{2}(?:\.\d+)?))(Z|[\+-]\d{2}:\d{2})?)$"

    # WHEN a log statement happens
    logger.info({})

    # THEN the timestamp has the appropriate formatting
    log = capture_logging_output(stdout)

    assert re.fullmatch(RFC3339_REGEX, log["timestamp"])  # "2022-10-27T17:42:26.841+0200"


def test_inject_lambda_context_log_event_request_data_classes(lambda_context, stdout, lambda_event, service_name):
    # GIVEN Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # WHEN a lambda function is decorated with logger instructed to log event
    # AND the event is an event source data class
    @event_source(data_class=S3Event)
    @logger.inject_lambda_context(log_event=True)
    def handler(event, context):
        logger.info("Hello")

    handler(lambda_event, lambda_context)

    # THEN logger should log event received from Lambda
    logged_event, _ = capture_multiple_logging_statements_output(stdout)
    assert logged_event["message"] == lambda_event


def test_inject_lambda_context_with_additional_args(lambda_context, stdout, service_name):
    # GIVEN Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # AND a handler that use additional parameters
    @logger.inject_lambda_context
    def handler(event, context, planet, str_end="."):
        logger.info(f"Hello {planet}{str_end}")

    handler({}, lambda_context, "World", str_end="!")

    # THEN the decorator should included them
    log = capture_logging_output(stdout)

    assert log["message"] == "Hello World!"


def test_logger_log_uncaught_exceptions(service_name, stdout):
    # GIVEN an initialized Logger is set with log_uncaught_exceptions
    logger = Logger(service=service_name, stream=stdout, log_uncaught_exceptions=True)

    # WHEN Python's exception hook is inspected
    exception_hook = sys.excepthook

    # THEN it should contain our custom exception hook with a copy of our logger
    assert isinstance(exception_hook, functools.partial)
    assert exception_hook.keywords.get("logger") == logger


def test_stream_defaults_to_stdout(service_name, capsys):
    # GIVEN Logger is initialized without any explicit stream
    logger = Logger(service=service_name)
    msg = "testing stdout"

    # WHEN logging statements are issued
    logger.info(msg)

    # THEN we should default to standard output, not standard error.
    # NOTE: we can't assert on capsys.readouterr().err due to a known bug: https://github.com/pytest-dev/pytest/issues/5997
    log = json.loads(capsys.readouterr().out.strip())
    assert log["message"] == msg


def test_logger_logs_stack_trace_with_default_value(service_name, stdout):
    # GIVEN a Logger instance with serialize_stacktrace default value = True
    logger = Logger(service=service_name, stream=stdout)

    # WHEN invoking a Lambda
    def handler(event, context):
        try:
            raise ValueError("something went wrong")
        except Exception:
            logger.exception("Received an exception")

    # THEN we expect a "stack_trace" in log
    handler({}, lambda_context)
    log = capture_logging_output(stdout)
    assert "stack_trace" in log


def test_logger_logs_stack_trace_with_non_default_value(service_name, stdout):
    # GIVEN a Logger instance with serialize_stacktrace = False
    logger = Logger(service=service_name, stream=stdout, serialize_stacktrace=False)

    # WHEN invoking a Lambda
    def handler(event, context):
        try:
            raise ValueError("something went wrong")
        except Exception:
            logger.exception("Received an exception")

    # THEN we expect a "stack_trace" not in log
    handler({}, lambda_context)
    log = capture_logging_output(stdout)
    assert "stack_trace" not in log


def test_define_log_level_via_advanced_logging_controls(
    monkeypatch,
    stdout,
    service_name,
):
    # GIVEN Logger is initialized with log level set to None
    # GIVEN AWS Lambda Advanced Logging Controls is set to DEBUG
    monkeypatch.setenv(constants.LAMBDA_LOG_LEVEL_ENV, "DEBUG")
    logger = Logger(service=service_name, stream=stdout)

    # WHEN logging statements are issued
    logger.debug("Hello")

    # THEN Lambda log level must be DEBUG
    log = capture_logging_output(stdout)
    assert log["level"] == "DEBUG"


def test_log_level_advanced_logging_controler_preference_over_powertools_log_level(
    monkeypatch,
    stdout,
    service_name,
):
    # GIVEN Logger is initialized with log level set to INFO
    # GIVEN AWS Lambda Advanced Logging Controls is set to DEBUG
    monkeypatch.setenv(constants.LAMBDA_LOG_LEVEL_ENV, "DEBUG")
    logger = Logger(service=service_name, stream=stdout, level="INFO")

    # WHEN logging statements are issued
    logger.debug("Hello")

    # THEN Lambda log level must be DEBUG because it takes precedence over POWERTOOLS_LOG_LEVEL
    log = capture_logging_output(stdout)
    assert log["level"] == "DEBUG"


def test_log_level_advanced_logging_controler_warning_different_log_levels_using_constructor(
    monkeypatch,
    service_name,
    stdout,
):
    # GIVEN AWS Lambda Advanced Logging Controls is set to INFO
    monkeypatch.setenv(constants.LAMBDA_LOG_LEVEL_ENV, "INFO")

    # WHEN Logger is initialized with log level set to DEBUG
    # THEN Logger should propagate a warning
    with pytest.warns(UserWarning):
        logger = Logger(service=service_name, stream=stdout, level="DEBUG")

    # THEN Logger must be INFO because it takes precedence over POWERTOOLS_LOG_LEVEL
    assert logger.log_level == logging.INFO


def test_log_level_advanced_logging_controler_warning_different_log_levels_using_set_level(
    monkeypatch,
    service_name,
    stdout,
):
    # GIVEN AWS Lambda Advanced Logging Controls is set to INFO
    # GIVEN Logger is initialized with log level set to None
    monkeypatch.setenv(constants.LAMBDA_LOG_LEVEL_ENV, "INFO")
    logger = Logger(service=service_name, stream=stdout)

    # WHEN Logger setLevel is set to DEBUG
    # THEN Logger should propagate a warning
    with pytest.warns(UserWarning):
        logger.setLevel(level="DEBUG")

    # THEN Logger must be INFO because it takes precedence over POWERTOOLS_LOG_LEVEL
    assert logger.log_level == logging.INFO


def test_logger_add_remove_filter(stdout, service_name):
    # GIVEN a Logger with a custom logging filter
    class ApiKeyFilter(logging.Filter):  # NOSONAR  # need filter to test actual impl.
        def filter(self, record):
            if getattr(record, "api_key", None):
                record.api_key = "REDACTED"

            return True

    redact_api_key_filter = ApiKeyFilter()
    logger = Logger(service=service_name, stream=stdout)
    logger.addFilter(redact_api_key_filter)

    # WHEN a new log statement is issued
    # AND another log statement is issued after filter is removed
    logger.info("filtered", api_key=secrets.token_urlsafe())
    logger.removeFilter(redact_api_key_filter)
    logger.info("unfiltered", api_key=secrets.token_urlsafe())

    # THEN logging filter should be called and mutate the log record accordingly
    log = capture_multiple_logging_statements_output(stdout)
    assert log[0]["api_key"] == "REDACTED"
    assert log[1]["api_key"] != "REDACTED"


def test_logger_json_unicode(stdout, service_name):
    # GIVEN Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # WHEN all non-ascii chars are logged as messages
    # AND non-ascii is also used as ephemeral fields
    # latest: https://www.unicode.org/versions/Unicode15.1.0/#Summary
    non_ascii_chars = [chr(i) for i in range(128, 111_411_1)]
    japanese_field = "47業レルし化"
    japanese_string = "スコビルデモ２"

    logger.info(non_ascii_chars, **{japanese_field: japanese_string})

    # THEN JSON logs should not try to escape them
    log = capture_logging_output(stdout)

    assert log["message"] == non_ascii_chars
    assert log[japanese_field] == japanese_string


def test_append_context_keys_adds_and_removes_keys(stdout, service_name):
    # GIVEN a Logger is initialized
    logger = Logger(service=service_name, stream=stdout)
    test_keys = {"user_id": "123", "operation": "test"}

    # WHEN context keys are added
    with logger.append_context_keys(**test_keys):
        logger.info("message with context keys")
    logger.info("message without context keys")

    # THEN context keys should only be present in the first log statement
    with_context_log, without_context_log = capture_multiple_logging_statements_output(stdout)

    assert "user_id" in with_context_log
    assert test_keys["user_id"] == with_context_log["user_id"]
    assert "user_id" not in without_context_log


def test_append_context_keys_handles_empty_dict(stdout, service_name):
    # GIVEN a Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # WHEN context is added with no keys
    with logger.append_context_keys():
        logger.info("message with empty context")

    # THEN log should contain only default keys
    log_output = capture_logging_output(stdout)
    assert set(log_output.keys()) == {"service", "timestamp", "level", "message", "location"}


def test_append_context_keys_handles_exception(stdout, service_name):
    # GIVEN a Logger is initialized
    logger = Logger(service=service_name, stream=stdout)
    test_user_id = "128"

    # WHEN an exception occurs within the context
    exception_raised = False
    try:
        with logger.append_context_keys(user_id=test_user_id):
            logger.info("message before exception")
            raise ValueError("Test exception")
    except ValueError:
        exception_raised = True
        logger.info("message after exception")

    # THEN verify the exception was raised and handled
    assert exception_raised, "Expected ValueError to be raised"


def test_append_context_keys_nested_contexts(stdout, service_name):
    # GIVEN a Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # WHEN nested contexts are used
    with logger.append_context_keys(level1="outer"):
        logger.info("outer context message")
        with logger.append_context_keys(level2="inner"):
            logger.info("nested context message")
        logger.info("back to outer context message")
    logger.info("no context message")

    # THEN logs should contain appropriate context keys
    outer, nested, back_outer, no_context = capture_multiple_logging_statements_output(stdout)

    assert outer["level1"] == "outer"
    assert "level2" not in outer

    assert nested["level1"] == "outer"
    assert nested["level2"] == "inner"

    assert back_outer["level1"] == "outer"
    assert "level2" not in back_outer

    assert "level1" not in no_context
    assert "level2" not in no_context


def test_append_context_keys_with_formatter(stdout, service_name):
    # GIVEN a Logger is initialized with a custom formatter
    class CustomFormatter(BasePowertoolsFormatter):
        def append_keys(self, **additional_keys):
            pass

        def clear_state(self) -> None:
            pass

        def remove_keys(self, keys: Iterable[str]) -> None:
            pass

    custom_formatter = CustomFormatter()
    logger = Logger(service=service_name, stream=stdout, logger_formatter=custom_formatter)
    test_keys = {"request_id": "id", "context": "value"}

    # WHEN context keys are added
    with logger.append_context_keys(**test_keys):
        logger.info("message with context")

    # THEN the context keys should not persist
    current_keys = logger.get_current_keys()
    assert current_keys == {}


def test_logger_change_level_child_logger(stdout, service_name):
    # GIVEN a new Logger and child Logger
    logger = Logger(service=service_name, stream=stdout)
    child_logger = Logger(service=service_name, child=True, stream=stdout, level="DEBUG")

    # WHEN we emit logs for both in DEBUG level
    logger.debug("PARENT")
    child_logger.debug("CHILD")

    # THEN only child log must emit log due to level
    logs = list(stdout.getvalue().strip().split("\n"))
    assert len(logs) == 1
    assert "service" in logs[0]


def test_clear_state_with_append_keys():
    # GIVEN a Logger is initialized
    logger = Logger(service="service_name", stream=stdout)

    # WHEN append keys are added
    logger.append_keys(custom_key="custom_key")
    logger.info("message with appended keys")
    logger.clear_state()

    # THEN context keys should be cleared
    assert "custom_key" not in logger.get_current_keys()


def test_clear_state(stdout, service_name):
    # GIVEN a Logger is initialized
    logger = Logger(service=service_name, stream=stdout)
    logger.info("message for the user")

    # WHEN the clear_state method is called
    logger.clear_state()

    # THEN the logger's current keys should be reset to their default values
    expected_keys = {
        "level": "%(levelname)s",
        "location": "%(funcName)s:%(lineno)d",
        "message": None,
        "timestamp": "%(asctime)s",
        "service": service_name,
        "sampling_rate": None,
    }
    assert logger.get_current_keys() == expected_keys


def test_clear_state_log_output(stdout, service_name):
    # GIVEN a Logger is initialized
    logger = Logger(service=service_name, stream=stdout)

    # WHEN we append a custom key and log
    logger.append_keys(custom_key="test_value")
    logger.info("first message")

    # AND we clear the state and log again
    logger.clear_state()
    logger.info("second message")

    # THEN the first log should contain the custom key
    # AND the second log should not contain the custom key
    first_log, second_log = capture_multiple_logging_statements_output(stdout)

    assert "custom_key" in first_log
    assert first_log["custom_key"] == "test_value"
    assert "custom_key" not in second_log