Python aws_hook.AwsHook类代码示例

 def expand_role(self):
     if 'Model' not in self.config:
         return
     hook = AwsHook(self.aws_conn_id)
     config = self.config['Model']
     if 'ExecutionRoleArn' in config:
         config['ExecutionRoleArn'] = hook.expand_role(config['ExecutionRoleArn'])

    def test_get_client_type_returns_a_boto3_client_of_the_requested_type(self):
        self._create_clusters()
        hook = AwsHook(aws_conn_id='aws_default')
        client_from_hook = hook.get_client_type('redshift')

        clusters = client_from_hook.describe_clusters()['Clusters']
        self.assertEqual(len(clusters), 2)

    def test_get_resource_type_returns_a_boto3_resource_of_the_requested_type(self):

        hook = AwsHook(aws_conn_id='aws_default')
        resource_from_hook = hook.get_resource_type('dynamodb')

        # this table needs to be created in production
        table = resource_from_hook.create_table(
            TableName='test_airflow',
            KeySchema=[
                {
                    'AttributeName': 'id',
                    'KeyType': 'HASH'
                },
            ],
            AttributeDefinitions=[
                {
                    'AttributeName': 'name',
                    'AttributeType': 'S'
                }
            ],
            ProvisionedThroughput={
                'ReadCapacityUnits': 10,
                'WriteCapacityUnits': 10
            }
        )

        table.meta.client.get_waiter(
            'table_exists').wait(TableName='test_airflow')

        self.assertEqual(table.item_count, 0)

 def test_expand_role(self):
     conn = boto3.client('iam', region_name='us-east-1')
     conn.create_role(RoleName='test-role', AssumeRolePolicyDocument='some policy')
     hook = AwsHook()
     arn = hook.expand_role('test-role')
     expect_arn = conn.get_role(RoleName='test-role').get('Role').get('Arn')
     self.assertEqual(arn, expect_arn)

    def test_get_client_type_returns_a_boto3_client_of_the_requested_type(self):
        client = boto3.client("emr", region_name="us-east-1")
        if len(client.list_clusters()["Clusters"]):
            raise ValueError("AWS not properly mocked")

        hook = AwsHook(aws_conn_id="aws_default")
        client_from_hook = hook.get_client_type("emr")

        self.assertEqual(client_from_hook.list_clusters()["Clusters"], [])

 def test_get_credentials_from_login(self, mock_get_connection):
     mock_connection = Connection(login='aws_access_key_id',
                                  password='aws_secret_access_key')
     mock_get_connection.return_value = mock_connection
     hook = AwsHook()
     credentials_from_hook = hook.get_credentials()
     self.assertEqual(credentials_from_hook.access_key, 'aws_access_key_id')
     self.assertEqual(credentials_from_hook.secret_key, 'aws_secret_access_key')
     self.assertIsNone(credentials_from_hook.token)

 def test_get_credentials_from_extra(self, mock_get_connection):
     mock_connection = Connection(
         extra='{"aws_access_key_id": "aws_access_key_id",'
         '"aws_secret_access_key": "aws_secret_access_key"}'
     )
     mock_get_connection.return_value = mock_connection
     hook = AwsHook()
     credentials_from_hook = hook.get_credentials()
     self.assertEqual(credentials_from_hook.access_key, 'aws_access_key_id')
     self.assertEqual(credentials_from_hook.secret_key, 'aws_secret_access_key')
     self.assertIsNone(credentials_from_hook.token)

    def _inject_aws_credentials(self):
        if TRANSFER_SPEC not in self.body or AWS_S3_DATA_SOURCE not in self.body[TRANSFER_SPEC]:
            return

        aws_hook = AwsHook(self.aws_conn_id)
        aws_credentials = aws_hook.get_credentials()
        aws_access_key_id = aws_credentials.access_key
        aws_secret_access_key = aws_credentials.secret_key
        self.body[TRANSFER_SPEC][AWS_S3_DATA_SOURCE][AWS_ACCESS_KEY] = {
            ACCESS_KEY_ID: aws_access_key_id,
            SECRET_ACCESS_KEY: aws_secret_access_key,
        }

 def test_get_credentials_from_role_arn(self, mock_get_connection):
     mock_connection = Connection(
         extra='{"role_arn":"arn:aws:iam::123456:role/role_arn"}')
     mock_get_connection.return_value = mock_connection
     hook = AwsHook()
     credentials_from_hook = hook.get_credentials()
     self.assertEqual(credentials_from_hook.access_key, 'AKIAIOSFODNN7EXAMPLE')
     self.assertEqual(credentials_from_hook.secret_key,
                      'aJalrXUtnFEMI/K7MDENG/bPxRfiCYzEXAMPLEKEY')
     self.assertEqual(credentials_from_hook.token,
                      'BQoEXAMPLEH4aoAH0gNCAPyJxz4BlCFFxWNE1OPTgk5TthT+FvwqnKwRcOIfrRh'
                      '3c/LTo6UDdyJwOOvEVPvLXCrrrUtdnniCEXAMPLE/IvU1dYUg2RVAJBanLiHb4I'
                      'gRmpRV3zrkuWJOgQs8IZZaIv2BXIa2R4OlgkBN9bkUDNCJiBeb/AXlzBBko7b15'
                      'fjrBs2+cTQtpZ3CYWFXG8C5zqx37wnOE49mRl/+OtkIKGO7fAE')

 def test_get_credentials_from_extra_with_s3_config_and_profile(
     self, mock_get_connection, mock_parse_s3_config
 ):
     mock_connection = Connection(
         extra='{"s3_config_format": "aws", '
               '"profile": "test", '
               '"s3_config_file": "aws-credentials", '
               '"region_name": "us-east-1"}')
     mock_get_connection.return_value = mock_connection
     hook = AwsHook()
     hook._get_credentials(region_name=None)
     mock_parse_s3_config.assert_called_with(
         'aws-credentials',
         'aws',
         'test'
     )

 def expand_role(self):
     if 'TrainingJobDefinition' in self.config:
         config = self.config['TrainingJobDefinition']
         if 'RoleArn' in config:
             hook = AwsHook(self.aws_conn_id)
             config['RoleArn'] = hook.expand_role(config['RoleArn'])

 def expand_role(self):
     if 'ExecutionRoleArn' in self.config:
         hook = AwsHook(self.aws_conn_id)
         self.config['ExecutionRoleArn'] = hook.expand_role(self.config['ExecutionRoleArn'])