Untitled

import logging

import boto3

from aurora_performance_data_service.helpers.brazil_path_utils import is_apollo_env
from aurora_performance_data_service.helpers.constants import (
    AURORA_PERF_ANALYSIS_ACCT,
    AWS_DEFAULT_REGION,
)
from aurora_performance_data_service.helpers.services.aws_credentials import (
    get_aws_credential_provider,
)
from aurora_performance_data_service.helpers.workflow.workflow_utils import get_app_configs

logger = logging.getLogger(__name__)


class __AWSClient:
    """
    Client class to connect to AWS services using IAM credentials

    # TODO: Need to add AWS auth for test and production deployment
    """

    def __init__(self, profile_name, service_name, region, iam_auth=False):
        """
        Initializer for AWSClient class
        :param profile_name: Profile name in the AWS credentials file.
        :param service_name: The AWS service with which we want to interact
        :param region: The region for the service specified by :service_name
        :param iam_auth: Use IAM credentials over Turtle.
        """
        if not service_name:
            err_msg = f"service_name parameters expected. Found {service_name}"
            logger.warning(err_msg)
            raise Exception(err_msg)

        self.__service_name = service_name
        self.__region = region
        self.__use_iam_role_auth = iam_auth
        self.__profile_name = profile_name
        self.__session = None
        self.__client = None
        self.__aws_credentials = None
        self.__app_config = get_app_configs()

    @property
    def aws_credentials(self):
        """ """
        if not self.__aws_credentials:
            if self.__app_config.turtle_cred_file_path or is_apollo_env():
                if self.__app_config.turtle_cred_file_path:
                    self.__aws_credentials = get_aws_credential_provider()
        return self.__aws_credentials

    def __get_boto_session(self):
        return boto3.Session(profile_name=self.__profile_name)

    @property
    def client(self, domain=None, region=None):
        if not self.__client:
            self.__session = self.__get_boto_session()
            if self.__use_iam_role_auth:
                self.__client = self.__session.client(
                    service_name=self.__service_name,
                    region_name=self.__region,
                )
                return self.__client

            if self.aws_credentials:
                self.__client = self.__session.client(
                    service_name=self.__service_name,
                    region_name=self.__region,
                    aws_access_key_id=self.aws_credentials.aws_access_key,
                    aws_secret_access_key=self.aws_credentials.aws_secret_access_key,
                    aws_session_token=self.aws_credentials.aws_session_token,
                )

        return self.__client


def get_aws_client(
    profile_name="default",
    service_name=None,
    region=AWS_DEFAULT_REGION,
    iam_auth=True,
):
    """Returns the AWS client object"""

    return __AWSClient(
        profile_name=profile_name,
        service_name=service_name,
        region=region,
        iam_auth=iam_auth,
    ).client


def get_turtle_credentials(
    profile_name="default",
    service_name=None,
    region=AWS_DEFAULT_REGION,
    iam_auth=False,
):
    return __AWSClient(
        profile_name=profile_name,
        service_name=service_name,
        region=region,
        iam_auth=iam_auth,
    ).aws_credentials


def get_instance_type(instance_id, region):
    client = get_aws_client(service_name="ec2", region=region)
    response = client.describe_instances()
    for reservation in response["Reservations"]:
        for instance in reservation["Instances"]:
            return instance["InstanceType"] if instance_id == instance["InstanceId"] else ""

1. I want to write a function in this code to test my code