atst/tests/mock_boto3.py

import pytest
from unittest.mock import Mock

from atst.domain.csp.cloud import AWSCloudProvider


AWS_CONFIG = {
    "AWS_ACCESS_KEY_ID": "",
    "AWS_SECRET_KEY": "",
    "AWS_REGION_NAME": "us-fake-1",
}
AUTH_CREDENTIALS = {
    "aws_access_key_id": AWS_CONFIG["AWS_ACCESS_KEY_ID"],
    "aws_secret_access_key": AWS_CONFIG["AWS_SECRET_KEY"],
}


def mock_boto_organizations(_config=None, **kwargs):
    describe_create_account_status = (
        "SUCCEEDED"
        if _config.get("organizations.describe_create_account.failure", False) == False
        else "FAILED"
    )

    import boto3

    mock = Mock(wraps=boto3.client("organizations", **kwargs))

    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/organizations.html#Organizations.Client.create_account
    mock.create_account = Mock(
        return_value={
            "CreateAccountStatus": {
                "Id": "create-account-status-id",
                "AccountName": "account-name",
                "AccountId": "account-id",
                "State": "SUCCEEDED",
            }
        }
    )

    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/organizations.html#Organizations.Client.describe_create_account_status
    mock.describe_create_account_status = Mock(
        return_value={
            "CreateAccountStatus": {
                "Id": "create-account-status-id",
                "AccountName": "account-name",
                "AccountId": "account-id",
                "State": describe_create_account_status,
            }
        }
    )
    return mock


def mock_boto_iam(_config=None, **kwargs):
    user_already_exists = _config.get("iam.create_user.already_exists", False)
    policy_already_exists = _config.get("iam.create_policy.already_exists", False)

    def _raise_entity_already_exists(**kwargs):
        raise real_iam_client.exceptions.EntityAlreadyExistsException(
            {"Error": {}}, "operation-name"
        )

    import boto3

    real_iam_client = boto3.client("iam", **kwargs)
    mock = Mock(wraps=real_iam_client)
    mock.exceptions.EntityAlreadyExistsException = (
        real_iam_client.exceptions.EntityAlreadyExistsException
    )

    mock.put_user_policy = Mock(return_value={"ResponseMetadata": {}})

    if user_already_exists:
        mock.create_user = Mock(side_effect=_raise_entity_already_exists)
    else:
        mock.create_user = Mock(
            return_value={
                "User": {
                    "UserId": "user-id",
                    "Arn": "user-arn",
                    "UserName": "user-name",
                }
            }
        )

    mock.get_user = Mock(
        return_value={
            "User": {"UserId": "user-id", "Arn": "user-arn", "UserName": "user-name"}
        }
    )

    mock.create_access_key = Mock(
        return_value={
            "AccessKey": {
                "AccessKeyId": "access-key-id",
                "SecretAccessKey": "secret-access-key",
            }
        }
    )

    if policy_already_exists:
        mock.create_policy = Mock(side_effect=_raise_entity_already_exists)
    else:
        mock.create_policy = Mock(return_value={"Policy": {"Arn": "policy-arn"}})

    return mock


def mock_boto_sts(_config=None, **kwargs):
    import boto3

    mock = Mock(wraps=boto3.client("sts", **kwargs))
    mock.assume_role = Mock(
        return_value={
            "Credentials": {
                "AccessKeyId": "access-key-id",
                "SecretAccessKey": "secret-access-key",
                "SessionToken": "session-token",
            }
        }
    )

    return mock


class MockBoto3:
    CLIENTS = {
        "organizations": mock_boto_organizations,
        "iam": mock_boto_iam,
        "sts": mock_boto_sts,
    }

    def __init__(self, config=None):
        self.config = config or {}
        self.client_instances = {}

    def client(self, client_name, **kwargs):
        """
        Return a new mock client for the given `client_name`, either by
        retrieving it from the `client_instances` cache or by instantiating
        it for the first time.

        Params should be the same ones you'd pass to `boto3.client`.
        """

        if client_name in self.client_instances:
            return self.client_instances[client_name]

        try:
            client_fn = self.CLIENTS[client_name]
            client_instance = client_fn(**kwargs, _config=self.config)
            self.client_instances[client_name] = client_instance
            return client_instance
        except KeyError:
            raise ValueError(f"MockBoto3: {client_name} client is not yet implemented.")


@pytest.fixture(scope="function")
def mock_boto3(request):
    marks = request.node.get_closest_marker("mock_boto3")
    if marks:
        mock_config = marks.args[0] if len(marks.args) else {}
    else:
        mock_config = {}
    return MockBoto3(mock_config)


@pytest.fixture(scope="function")
def mock_aws(mock_boto3):
    return AWSCloudProvider(AWS_CONFIG, boto3=mock_boto3)