# Copyright (c) 2021 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 Callable

import pyfaaster.common.utils as utils
from aws_lambda_powertools.middleware_factory import lambda_handler_decorator

from czc.api.environment import TestKeyEnvironment
from czc.api.v2.constants import CLOUDZERO_TEST_KEY_HEADER


@lambda_handler_decorator(trace_execution=True)
def web_test_key_aware(handler: Callable, event, context, **kwargs):
    """Make this Web API Lambda Handler Test Key Aware

    Parameters
    ----------

    Returns
    -------
    Callable
        Decorated Lambda Handler
    """
    cloudzero_test_key_value = utils.deep_get(event, 'headers', CLOUDZERO_TEST_KEY_HEADER, ignore_case=True)
    with TestKeyEnvironment(cloudzero_test_key_value):
        return handler(event, context, **kwargs)
