# 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_CONNECTIONS_ACCOUNT_ID,
    DEV_CORE_ACCOUNT_ID,
    DEV_REACTOR_ACCOUNT_ID,
    PROD_CONNECTIONS_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()
        aws_environment.aws_service.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)),
        (DEV_CONNECTIONS_ACCOUNT_ID, AccountInfo(Environment.dev, Account.connections)),
        (PROD_CORE_ACCOUNT_ID, AccountInfo(Environment.prod, Account.core)),
        (PROD_REACTOR_ACCOUNT_ID, AccountInfo(Environment.prod, Account.reactor)),
        (PROD_CONNECTIONS_ACCOUNT_ID, AccountInfo(Environment.prod, Account.connections)),
    ],
)
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


@pytest.mark.unit
@pytest.mark.parametrize(
    'property_func,service_name',
    [
        (lambda s: s.apigateway_client, 'apigateway'),
        (lambda s: s.cloudformation_client, 'cloudformation'),
        (lambda s: s.ec2_client, 'ec2'),
        (lambda s: s.events_client, 'events'),
        (lambda s: s.iam_client, 'iam'),
        (lambda s: s.lambda_client, 'lambda'),
        (lambda s: s.logs_client, 'logs'),
        (lambda s: s.s3_client, 's3'),
        (lambda s: s.secretsmanager_client, 'secretsmanager'),
        (lambda s: s.sqs_client, 'sqs'),
        (lambda s: s.stepfunctions_client, 'stepfunctions'),
        (lambda s: s.sts_client, 'sts'),
        (lambda s: s.tagging_api_client, 'resourcegroupstaggingapi'),
    ],
)
def test_aws_service_clients_call_boto3_correctly(context, property_func, service_name):
    """Test that each aws_service client property calls boto3.client with the correct service name."""
    mock_client = Mock()
    MockEx(context.mock_boto3.client).if_called_with(service_name).return_value(mock_client)

    service = aws_environment._AwsService()
    result = property_func(service)

    assert result == mock_client
    assert MockEx(context.mock_boto3.client).was_called_once_with(service_name)


@pytest.mark.unit
@pytest.mark.parametrize(
    'property_func,service_name',
    [
        (lambda s: s.cloudformation_resource, 'cloudformation'),
        (lambda s: s.dynamodb_resource, 'dynamodb'),
        (lambda s: s.s3_resource, 's3'),
        (lambda s: s.sqs_resource, 'sqs'),
    ],
)
def test_aws_service_resources_call_boto3_correctly(context, property_func, service_name):
    """Test that each aws_service resource property calls boto3.resource with the correct service name."""
    mock_resource = Mock()
    MockEx(context.mock_boto3.resource).if_called_with(service_name).return_value(mock_resource)

    service = aws_environment._AwsService()
    result = property_func(service)

    assert result == mock_resource
    assert MockEx(context.mock_boto3.resource).was_called_once_with(service_name)


@pytest.mark.unit
@pytest.mark.parametrize(
    'boto_func,service_func',
    [
        (lambda b: b.client, lambda s: s.s3_client),
        (lambda b: b.resource, lambda s: s.dynamodb_resource),
        (lambda b: b.client, lambda s: s.client('ec2')),
        (lambda b: b.resource, lambda s: s.resource('sqs')),
    ],
)
def test_aws_service_caching_and_cache_clear(context, boto_func, service_func):
    service_mock = boto_func(context.mock_boto3)
    service_mock.side_effect = lambda name: Mock()

    aws_service = aws_environment.aws_service
    result1 = service_func(aws_service)
    result2 = service_func(aws_service)

    assert result1 is result2
    assert service_mock.call_count == 1

    aws_service.cache_clear()

    result3 = service_func(aws_service)

    assert result3 is not result1
    assert service_mock.call_count == 2
