# 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 unittest.mock import Mock

import pytest

from czc import aws_environment
from czc.aws_environment import (
    DEV_CORE_ACCOUNT_ID,
    DEV_REACTOR_ACCOUNT_ID,
    PROD_CORE_ACCOUNT_ID,
    PROD_REACTOR_ACCOUNT_ID,
    Account,
    AccountInfo,
    Environment,
)
from czt.unittest import Context, MockEx


@pytest.fixture(scope='function')
def context(mocker, environ):
    with Context(mocker, aws_environment, ['boto3'], extra=dict(mock_sts=Mock())) as context:
        MockEx(context.mock_boto3.client).if_called_with('sts').return_value(context.mock_sts)
        aws_environment.resolve_account_info.cache_clear()
        yield context


@pytest.mark.unit
@pytest.mark.parametrize(
    'account,expected_account_info',
    [
        (DEV_CORE_ACCOUNT_ID, AccountInfo(Environment.dev, Account.core)),
        (DEV_REACTOR_ACCOUNT_ID, AccountInfo(Environment.dev, Account.reactor)),
        (PROD_CORE_ACCOUNT_ID, AccountInfo(Environment.prod, Account.core)),
        (PROD_REACTOR_ACCOUNT_ID, AccountInfo(Environment.prod, Account.reactor)),
    ],
)
def test_resolve_account_info(context, account, expected_account_info):
    MockEx(context.mock_sts.get_caller_identity).if_called_with().return_value(dict(Account=account))
    account_info = aws_environment.resolve_account_info()
    assert account_info == expected_account_info
