# 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

from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
from aws_lambda_powertools.utilities.data_classes import APIGatewayProxyEvent

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


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

    Parameters
    ----------

    Returns
    -------
    Callable
        Decorated Lambda Handler
    """
    with TestKeyEnvironment(event.get_header_value(CLOUDZERO_TEST_KEY_HEADER)):
        return handler(event, context, **kwargs)
