# 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

import os
from typing import Optional

from aws_lambda_powertools import Logger

from czc.api.v2.constants import CLOUDZERO_TEST_KEY

logger = Logger()


def _set_test_key(test_key_value):
    os.environ[CLOUDZERO_TEST_KEY] = test_key_value
    logger.info(f'Test Key is present! Using {test_key_value}')


def _unset_test_key():
    if CLOUDZERO_TEST_KEY in os.environ:
        logger.info(f'Removing {CLOUDZERO_TEST_KEY} from environment.')
        os.environ.pop(CLOUDZERO_TEST_KEY)


class TestKeyEnvironment:
    def __init__(self, cloudzero_test_key_value: Optional[str]):
        self._cloudzero_test_key_value = cloudzero_test_key_value

    def __enter__(self):
        if self._cloudzero_test_key_value:
            _set_test_key(self._cloudzero_test_key_value)
        else:
            logger.info(f'Test Key is not present! Resetting {CLOUDZERO_TEST_KEY}.')
            _unset_test_key()

    def __exit__(self, exc_type, exc_val, exc_tb):
        _unset_test_key()
