Python botocore.stub.Stubber() Examples

def test_map_parameter_grouping():
    """
    Test special parameters that have shape of map are recorded
    as a list of keys based on `para_whitelist.json`
    """
    ddb = session.create_client('dynamodb', region_name='us-west-2')
    response = {
        'ResponseMetadata': {
            'RequestId': REQUEST_ID,
            'HTTPStatusCode': 500,
        }
    }

    with Stubber(ddb) as stubber:
        stubber.add_response('batch_write_item', response, {'RequestItems': ANY})
        ddb.batch_write_item(RequestItems={'table1': [{}], 'table2': [{}]})

    subsegment = xray_recorder.current_segment().subsegments[0]
    assert subsegment.fault
    assert subsegment.http['response']['status'] == 500

    aws_meta = subsegment.aws
    assert sorted(aws_meta['table_names']) == ['table1', 'table2'] 

def test_create_fails_existing_rule_set():
    ses = botocore.session.get_session().create_client("ses")
    stubber = Stubber(ses)
    stubber.add_response(
        "describe_active_receipt_rule_set", active_receipt_rule_set_response
    )
    stubber.activate()
    provider._ses = ses
    request = Request("Create", "lists.binx.io")
    response = handler(request, {})
    assert response["Status"] == "FAILED", response["Reason"]
    assert (
        response["Reason"]
        == "active receipt rule set is already set in region eu-west-1 - lists.binx.io"
    )
    stubber.assert_no_pending_responses() 

def test_delete_non_existing(self, mock_client):
        """Test that trying to delete a s3 that doesn't exist raises an error
        """
        bucket, key = 'glance', 'no_exist'
        fake_s3_client = botocore.session.get_session().create_client('s3')

        with stub.Stubber(fake_s3_client) as stubber:
            stubber.add_client_error(method='head_object',
                                     service_error_code='404',
                                     service_message='''
                                     The specified key does not exist.
                                     ''',
                                     expected_params={
                                         'Bucket': bucket,
                                         'Key': key
                                     })
            fake_s3_client.head_bucket = mock.MagicMock()
            mock_client.return_value = fake_s3_client

            uri = "s3+https://user:key@auth_address/%s/%s" % (bucket, key)
            loc = location.get_location_from_uri_and_backend(uri,
                                                             's3_region1',
                                                             conf=self.conf)
            self.assertRaises(exceptions.NotFound, self.store.delete, loc) 

def test_no_such_identity():
    ses = botocore.session.get_session().create_client("ses", region_name="eu-west-1")
    stubber = Stubber(ses)
    stubber.add_response(
        "get_identity_verification_attributes",
        GetIdentityVerificationAttributesReponse(),
        {"Identities": ["lists.binx.io"]},
    )
    stubber.activate()
    provider._ses = ses
    counter = Counter()
    provider.invoke_lambda = counter.increment

    request = Request("Create", "lists.binx.io", "eu-west-1")
    response = handler(request, ())
    assert response["Status"] == "FAILED", response["Reason"]
    assert not provider.asynchronous
    stubber.assert_no_pending_responses() 

def test_create_cluster_valid_parameters(cluster_provisioner):
    """Test that the parameters passed down to run_job_flow are valid"""

    stubber = Stubber(cluster_provisioner.emr)
    response = {"JobFlowId": "job-flow-id"}
    stubber.add_response("run_job_flow", response)

    emr_release = "5.0.0"
    with stubber:
        jobflow_id = cluster_provisioner.start(
            user_username="user",
            user_email="user@example.com",
            identifier="cluster",
            emr_release=emr_release,
            size=3,
            public_key="public-key",
        )

    assert jobflow_id == response["JobFlowId"] 

def test_spark_job_add(notebook_maker, spark_job_provisioner):
    notebook = notebook_maker()
    identifier = "test-identifier"
    key = "jobs/%s/%s" % (identifier, notebook.name)

    stubber = Stubber(spark_job_provisioner.s3)
    response = {"Expiration": "whatever", "ETag": "12345", "VersionId": "1.0"}
    expected_params = {
        "Body": notebook,
        "Bucket": settings.AWS_CONFIG["CODE_BUCKET"],
        "Key": key,
    }
    stubber.add_response("put_object", response, expected_params)

    with stubber:
        result = spark_job_provisioner.add(
            identifier=identifier, notebook_file=notebook
        )
        assert result == key 

def test_decrypt_invalid_ciphertext_error_no_context(self):
        kms_client = boto3.client('kms')
        key_id = "test"
        encryption_context = {}
        with Stubber(kms_client) as stubber:
            stubber.add_client_error(
                'decrypt',
                'InvalidCiphertextException',
                'The request was rejected because the specified ciphertext, or additional authenticated data incorporated into the ciphertext, such as the encryption context, is corrupted, missing, or otherwise invalid.',
                400, 
                expected_params = {
                    'CiphertextBlob': 'encoded_key',
                    'EncryptionContext': encryption_context
            })
            key_service = KeyService(kms_client, key_id, encryption_context)
            with self.assertRaises(KmsError) as e:
                msg = ("Could not decrypt hmac key with KMS. The credential may "
                        "require that an encryption context be provided to decrypt "
                        "it.")
                response = key_service.decrypt('encoded_key')
                self.assertEqual(e, KmsError(msg)) 

def test_create_no_existing_rule_set():
    ses = botocore.session.get_session().create_client("ses")
    stubber = Stubber(ses)
    stubber.add_response(
        "describe_active_receipt_rule_set", no_active_receipt_rule_set_response
    )
    stubber.add_response(
        "set_active_receipt_rule_set",
        no_active_receipt_rule_set_response,
        {"RuleSetName": "lists.binx.io"},
    )
    stubber.activate()
    provider._ses = ses
    request = Request("Create", "lists.binx.io")
    response = handler(request, {})
    assert response["Status"] == "SUCCESS", response["Reason"]
    stubber.assert_no_pending_responses() 

def test_refusal_to_overwrite_settings():
    ses = botocore.session.get_session().create_client("ses", region_name="eu-west-1")
    stubber = Stubber(ses)
    provider._ses = ses

    stubber.add_response(
        "get_identity_notification_attributes",
        GetIdentityNotificationAttributesResponse(attributes),
        {"Identities": ["lists.binx.io"]},
    )
    stubber.activate()
    provider._ses = ses

    request = Request("Create", attributes)
    response = handler(request, ())
    assert response["Status"] == "FAILED"
    stubber.assert_no_pending_responses() 

def test_update_receipt_change_region_with_existing_rule_set():
    ses = botocore.session.get_session().create_client("ses")
    stubber = Stubber(ses)
    stubber.add_response(
        "describe_active_receipt_rule_set", active_receipt_rule_set_response
    )
    stubber.activate()
    provider._ses = ses
    request = Request("Update", "lists.xebia.com")
    request["OldResourceProperties"] = {"Region": "us-east-1"}
    response = handler(request, {})
    assert response["Status"] == "FAILED"
    assert (
        response["Reason"]
        == "active receipt rule set is already set in region eu-west-1 - lists.binx.io"
    )
    stubber.assert_no_pending_responses() 

def test_create_missing_topic(self):
        app.current_request = Request(
            push,
            {
                "X-GitHub-Event": "comment",
                "X-Hub-Signature": "sha=whatever",
                "X-GitHub-Delivery": uuid.uuid4().hex,
            },
        )

        stubber = Stubber(SNS)
        stubber.add_response(
            "list_topics", {"Topics": [{"TopicArn": "arn:foo:bar:push"}]}
        )
        stubber.add_response("create_topic", {"TopicArn": "arn:foo:baz:123_push"})
        stubber.add_response("publish", {"MessageId": "1234"})

        with stubber:
            response = index("123")
        assert response == {"Code": "Ok", "Message": "Webhook received."} 

def test_adjusts_nothing_when_equal(cluster_response):
    ecs = lambda_function.ecs_client()
    stubber = Stubber(ecs)

    describe_services_response = {
        'services': [
            {
                'serviceArn': AGENT_SERVICE_ARN,
                'serviceName': 'AgentService',
                'clusterArn': CLUSTER_ARN,
                'desiredCount': 3,
            }
        ]
    }
    expected_params = {'cluster': 'cluster1', 'services': [AGENT_SERVICE_ARN]}
    stubber.add_response('describe_services', describe_services_response, expected_params)

    expected_params = {'clusters': ['cluster1']}
    stubber.add_response('describe_clusters', cluster_response, expected_params)

    with stubber:
        response = lambda_function.adjust_service_desired_count(ecs, 'cluster1', AGENT_SERVICE_ARN)
        assert response is None 

def test__create_bucket_fails(self, mock_consul):
        """
        Test s3 provisioning fails on bucket creation, and retries up to 4 times
        """
        instance = OpenEdXInstanceFactory()
        instance.s3_access_key = 'test'
        instance.s3_secret_access_key = 'test'
        instance.s3_bucket_name = 'test'
        max_tries = 4
        stubber = Stubber(s3_client)
        for _ in range(max_tries):
            stubber.add_client_error('create_bucket')
        with self.assertLogs('instance.models.instance', level='INFO') as cm:
            with stubber:
                with self.assertRaises(ClientError):
                    instance._create_bucket(max_tries=max_tries)

            base_log_text = (
                'INFO:instance.models.instance:instance={} ({!s:.15}) | Retrying bucket creation'
                ' due to "", attempt %s of {}.'.format(instance.ref.pk, instance.ref.name, max_tries)
            )
            self.assertEqual(
                cm.output,
                [base_log_text % i for i in range(1, max_tries + 1)]
            ) 

def test_get_non_existing(self, mock_client):
        """Test that trying to retrieve a s3 that doesn't exist raises an error
        """
        bucket, key = 'glance', 'no_exist'
        fake_s3_client = botocore.session.get_session().create_client('s3')

        with stub.Stubber(fake_s3_client) as stubber:
            stubber.add_client_error(method='head_object',
                                     service_error_code='404',
                                     service_message='''
                                     The specified key does not exist.
                                     ''',
                                     expected_params={
                                         'Bucket': bucket,
                                         'Key': key
                                     })
            mock_client.return_value = fake_s3_client

            uri = "s3://user:key@auth_address/%s/%s" % (bucket, key)
            loc = location.get_location_from_uri(uri, conf=self.conf)
            self.assertRaises(exceptions.NotFound, self.store.get, loc) 

def test_add_with_verifier(self, mock_client):
        """Assert 'verifier.update' is called when verifier is provided"""
        expected_image_id = str(uuid.uuid4())
        expected_s3_size = FIVE_KB
        expected_s3_contents = b"*" * expected_s3_size
        image_s3 = six.BytesIO(expected_s3_contents)

        fake_s3_client = botocore.session.get_session().create_client('s3')
        verifier = mock.MagicMock(name='mock_verifier')

        with stub.Stubber(fake_s3_client) as stubber:
            stubber.add_response(method='head_bucket', service_response={})
            stubber.add_client_error(method='head_object',
                                     service_error_code='404',
                                     service_message='')
            stubber.add_response(method='put_object', service_response={})

            mock_client.return_value = fake_s3_client
            self.store.add(expected_image_id, image_s3, expected_s3_size,
                           self.hash_algo, verifier=verifier)
        verifier.update.assert_called_with(expected_s3_contents) 

def test_get_non_existing(self, mock_client):
        """Test that trying to retrieve a s3 that doesn't exist raises an
        error
        """
        bucket, key = 'glance', 'no_exist'
        fake_s3_client = botocore.session.get_session().create_client('s3')

        with stub.Stubber(fake_s3_client) as stubber:
            stubber.add_client_error(method='head_object',
                                     service_error_code='404',
                                     service_message='''
                                     The specified key does not exist.
                                     ''',
                                     expected_params={
                                         'Bucket': bucket,
                                         'Key': key
                                     })
            mock_client.return_value = fake_s3_client

            uri = "s3+https://user:key@auth_address/%s/%s" % (bucket, key)
            loc = location.get_location_from_uri_and_backend(uri,
                                                             's3_region1',
                                                             conf=self.conf)
            self.assertRaises(exceptions.NotFound, self.store.get, loc) 

def test_add_with_verifier(self, mock_client):
        """Assert 'verifier.update' is called when verifier is provided"""
        expected_image_id = str(uuid.uuid4())
        expected_s3_size = FIVE_KB
        expected_s3_contents = b"*" * expected_s3_size
        image_s3 = six.BytesIO(expected_s3_contents)

        fake_s3_client = botocore.session.get_session().create_client('s3')
        verifier = mock.MagicMock(name='mock_verifier')

        with stub.Stubber(fake_s3_client) as stubber:
            stubber.add_response(method='head_bucket', service_response={})
            stubber.add_client_error(method='head_object',
                                     service_error_code='404',
                                     service_message='')
            stubber.add_response(method='put_object', service_response={})

            mock_client.return_value = fake_s3_client
            self.store.add(expected_image_id, image_s3, expected_s3_size,
                           self.hash_algo, verifier=verifier)
        verifier.update.assert_called_with(expected_s3_contents) 

def test_generate_key_data_success(self):
        kms_client = boto3.client('kms')
        key_id = "test"
        encryption_context = {}
        with Stubber(kms_client) as stubber:
            stubber.add_response('generate_data_key', {
            'CiphertextBlob': b'ciphertext',
            'Plaintext': b'plaintext',
            'KeyId': 'string'
        }, expected_params = {
            'KeyId': key_id,
            'EncryptionContext': encryption_context,
            'NumberOfBytes': 1
        })
            key_service = KeyService(kms_client, key_id, encryption_context)
            response = key_service.generate_key_data(1)
        self.assertEqual(response[0], b'plaintext')
        self.assertEqual(response[1], b'ciphertext') 

def test_md5_from_s3_tags_no_md5(self):
        tag_set = {"TagSet": []}

        s3_stubber = Stubber(self.s3_client)
        get_object_tagging_response = tag_set
        get_object_tagging_expected_params = {
            "Bucket": self.s3_bucket_name,
            "Key": self.s3_key_name,
        }
        s3_stubber.add_response(
            "get_object_tagging",
            get_object_tagging_response,
            get_object_tagging_expected_params,
        )
        with s3_stubber:
            md5_hash = md5_from_s3_tags(
                self.s3_client, self.s3_bucket_name, self.s3_key_name
            )
            self.assertEquals("", md5_hash) 

def test_md5_from_s3_tags_has_md5(self):
        expected_md5_hash = "d41d8cd98f00b204e9800998ecf8427e"
        tag_set = {"TagSet": [{"Key": "md5", "Value": expected_md5_hash}]}

        s3_stubber = Stubber(self.s3_client)
        get_object_tagging_response = tag_set
        get_object_tagging_expected_params = {
            "Bucket": self.s3_bucket_name,
            "Key": self.s3_key_name,
        }
        s3_stubber.add_response(
            "get_object_tagging",
            get_object_tagging_response,
            get_object_tagging_expected_params,
        )
        with s3_stubber:
            md5_hash = md5_from_s3_tags(
                self.s3_client, self.s3_bucket_name, self.s3_key_name
            )
            self.assertEquals(expected_md5_hash, md5_hash) 

def test_time_from_s3(self):

        expected_s3_time = datetime.datetime(2019, 1, 1)

        s3_stubber = Stubber(self.s3_client)
        head_object_response = {"LastModified": expected_s3_time}
        head_object_expected_params = {
            "Bucket": self.s3_bucket_name,
            "Key": self.s3_key_name,
        }
        s3_stubber.add_response(
            "head_object", head_object_response, head_object_expected_params
        )
        with s3_stubber:
            s3_time = time_from_s3(
                self.s3_client, self.s3_bucket_name, self.s3_key_name
            )
            self.assertEquals(expected_s3_time, s3_time) 

def test_verify_s3_object_versioning_not_enabled(self):
        s3_obj = self.s3.Object(self.s3_bucket_name, self.s3_key_name)

        # Set up responses
        get_bucket_versioning_response = {"Status": "Disabled"}
        get_bucket_versioning_expected_params = {"Bucket": self.s3_bucket_name}
        s3_stubber_resource = Stubber(self.s3.meta.client)
        s3_stubber_resource.add_response(
            "get_bucket_versioning",
            get_bucket_versioning_response,
            get_bucket_versioning_expected_params,
        )
        with self.assertRaises(Exception) as cm:
            with s3_stubber_resource:
                verify_s3_object_version(self.s3, s3_obj)
            self.assertEquals(
                cm.exception.message,
                "Object versioning is not enabled in bucket {}".format(
                    self.s3_bucket_name
                ),
            ) 

def test_describe_table(loop, recorder):
    segment = recorder.begin_segment('name')

    req_id = '1234'
    response = {'ResponseMetadata': {'RequestId': req_id, 'HTTPStatusCode': 403}}

    session = aiobotocore.get_session()
    async with session.create_client('dynamodb', region_name='eu-west-2') as client:
        with Stubber(client) as stubber:
            stubber.add_response('describe_table', response, {'TableName': 'mytable'})
            await client.describe_table(TableName='mytable')

    subsegment = segment.subsegments[0]
    assert subsegment.error
    assert subsegment.http['response']['status'] == 403

    aws_meta = subsegment.aws
    assert aws_meta['table_name'] == 'mytable'
    assert aws_meta['request_id'] == req_id
    assert aws_meta['region'] == 'eu-west-2'
    assert aws_meta['operation'] == 'DescribeTable' 

def test_s3_parameter_capture(loop, recorder):
    segment = recorder.begin_segment('name')

    bucket_name = 'mybucket'
    key = 'mykey'
    version_id = 'myversionid'
    response = {'ResponseMetadata': {'RequestId': '1234', 'HTTPStatusCode': 200}}

    session = aiobotocore.get_session()
    async with session.create_client('s3', region_name='eu-west-2') as client:
        with Stubber(client) as stubber:
            stubber.add_response('get_object', response,
                                 {'Bucket': bucket_name, 'Key': key, 'VersionId': version_id})
            await client.get_object(Bucket=bucket_name, Key=key,
                                    VersionId=version_id)

    subsegment = segment.subsegments[0]
    aws_meta = subsegment.aws

    assert aws_meta['bucket_name'] == bucket_name
    assert aws_meta['key'] == key
    assert aws_meta['version_id'] == version_id
    assert aws_meta['operation'] == 'GetObject' 

def test_ddb_table_name():
    ddb = session.create_client('dynamodb', region_name='us-west-2')
    response = {
        'ResponseMetadata': {
            'RequestId': REQUEST_ID,
            'HTTPStatusCode': 403,
        }
    }

    with Stubber(ddb) as stubber:
        stubber.add_response('describe_table', response, {'TableName': 'mytable'})
        ddb.describe_table(TableName='mytable')

    subsegment = xray_recorder.current_segment().subsegments[0]
    assert subsegment.error
    assert subsegment.http['response']['status'] == 403

    aws_meta = subsegment.aws
    assert aws_meta['table_name'] == 'mytable'
    assert aws_meta['request_id'] == REQUEST_ID
    assert aws_meta['region'] == 'us-west-2'
    assert aws_meta['operation'] == 'DescribeTable' 

def test_pass_through_on_context_missing():
    """
    The built-in patcher or subsegment capture logic should not throw
    any error when a `None` subsegment created from `LOG_ERROR` missing context.
    """
    xray_recorder.configure(context_missing='LOG_ERROR')
    xray_recorder.clear_trace_entities()

    ddb = session.create_client('dynamodb', region_name='us-west-2')
    response = {
        'ResponseMetadata': {
            'RequestId': REQUEST_ID,
            'HTTPStatusCode': 200,
        }
    }

    with Stubber(ddb) as stubber:
        stubber.add_response('describe_table', response, {'TableName': 'mytable'})
        result = ddb.describe_table(TableName='mytable')
    assert result is not None

    xray_recorder.configure(context_missing='RUNTIME_ERROR') 

def test_put_item():
    """Test put_item."""
    table_name = 'test-table'
    id_val = 'my_id'
    target = 'my_target'
    table: Table = boto3.resource('dynamodb').Table(table_name)
    stubber = Stubber(table.meta.client)

    stubber.add_response(
        'put_item',
        {
            'Attributes': {
                'id': {'S': id_val},
                'target': {'S': target}
            }
        }
    )

    with stubber:
        assert not put_item(table, id_val, target) 

def test_validate_account_alias(caplog):
    """Test validate_account_alias."""
    caplog.set_level(logging.INFO, logger='runway')
    alias = 'test-alias'
    iam_client = boto3.client('iam')
    stubber = Stubber(iam_client)

    stubber.add_response('list_account_aliases', {'AccountAliases': [alias]})
    stubber.add_response('list_account_aliases', {'AccountAliases': ['no-match']})

    with stubber:
        assert not validate_account_alias(iam_client, alias)
        with pytest.raises(SystemExit) as excinfo:
            assert validate_account_alias(iam_client, alias)
        assert excinfo.value.code == 1
    stubber.assert_no_pending_responses()
    assert caplog.messages == [
        'Verified current AWS account alias matches required alias {}.'.format(alias),
        'Current AWS account aliases "{}" do not match required account'
        ' alias {} in Runway config.'.format('no-match', alias)
    ] 

def test_account_id(self, monkeypatch):
        """Test account_id."""
        account_id = '123456789012'
        client = boto3.client('sts')
        stubber = Stubber(client)
        mock_session = MockBoto3Session(clients={'sts.us-east-1': client},
                                        region_name='us-east-1')
        monkeypatch.setattr(Context, 'get_session', lambda self_: mock_session)
        context = Context(env_name='test',
                          env_region='us-east-1',
                          env_root='./',
                          env_vars={})

        stubber.add_response(method='get_caller_identity',
                             service_response={'UserId': 'test-user',
                                               'Account': account_id,
                                               'Arn': 'arn:aws:iam::{}:'
                                                      'user/test-user'
                                                      .format(account_id)})

        with stubber as stub:
            assert context.account_id == account_id
        stub.assert_no_pending_responses() 

def test_resolve_cfn_outputs(self, kwargs, stack_info, expected):
        """Test resolve_cfn_outputs."""
        client = boto3.client('cloudformation')
        stubber = Stubber(client)
        for stack, outputs in stack_info.items():
            for key, val in outputs.items():
                stubber.add_response(
                    'describe_stacks',
                    {
                        'Stacks': [{
                            'StackName': stack,
                            'CreationTime': datetime.now(),
                            'StackStatus': 'CREATE_COMPLETE',
                            'Outputs': [{
                                'OutputKey': key,
                                'OutputValue': val
                            }]
                        }]
                    }
                )
        with stubber:
            assert TerraformBackendConfig.resolve_cfn_outputs(
                client, **kwargs
            ) == expected
        stubber.assert_no_pending_responses()