# Copyright (c) 2022 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

from typing import Mapping

import pytest
from toolz.dicttoolz import dissoc

import czc.logging as logging


class TestLogger:
    __test__ = False
    context: Mapping

    def __init__(self) -> None:
        self.context = {}

    def append_keys(self, **kwargs):
        self.context = {
            **self.context,
            **kwargs
        }

    def remove_keys(self, keys: Mapping):
        self.context = dissoc(self.context, *keys.keys())


@pytest.mark.unit
def test_log_context_add_event_data():
    logger = TestLogger()

    @logging.log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}

    event = {
        'field1': 'value1',
        'field2': {
            'field3': 'value3'
        }
    }

    context = test_fn(event)

    assert context == event
    assert logger.context == {}


@pytest.mark.unit
def test_log_context_nested():
    logger = TestLogger()

    inner_event = {
        'field4': 'value4'
    }

    @logging.log_context_from_event(logger)
    def test_inner_fn(event):
        return {**logger.context}

    @logging.log_context_from_event(logger)
    def test_fn(event):
        return {**logger.context}, test_inner_fn(inner_event)

    event = {
        'field1': 'value1',
        'field2': {
            'field3': 'value3'
        }
    }

    context, inner_context = test_fn(event)

    assert context == event
    assert inner_context == {**event, **inner_event}
    assert logger.context == {}
