test_tags.py

#!/usr/bin/env python
import argparse
import logging
import string
import random

import boto3
from botocore.client import Config

import pytest

def make_random_string(size):
    return ''.join(random.choice(string.ascii_letters) for _ in range(size))


class S3():
    def __init__(self, access, secret, endpoint, bucket, debug=True, region='us-east-1'):
        self.s3 = boto3.client(
            's3',region,
            endpoint_url = endpoint,
            aws_access_key_id = access,
            aws_secret_access_key = secret,
        )
        self.bucket = bucket
        self.logger = logging.getLogger(__name__)
        if debug:
            boto3.set_stream_logger(name='botocore',level=logging.DEBUG)
            self.logger.setLevel(logging.DEBUG)
        else:
            self.logger.setLevel(logging.INFO)
        self.obj_set = set()

    @staticmethod
    def make_tags(tags):
        tag_list = []
        for k,v in tags.items():
            tag_list.append({'Key':k,
                             'Value':v if v is not None else ''})

        return {'TagSet':tag_list}

    def init_bucket_obj(self,bucket,obj):
        try:
            self.s3.head_bucket(Bucket=bucket)
            self.logger.debug('skipping bucket create as it already exists')
        except:
            self.logger.debug('creating bucket %s' % bucket)
            self.s3.create_bucket(Bucket=bucket)

        try:
            self.s3.head_object(Bucket=bucket, Key=obj)
            self.logger.debug('skipping object create as it already exists')
        except:
            self.logger.debug('creating obj')
            self.s3.put_object(Bucket=bucket, Key=obj, Body='abcdefg', Tagging="foo&bar=baZ")


    def test_put_get_tags(self, obj, tagset):
        self.init_bucket_obj(self.bucket,obj)
        self.obj_set.add(obj)
        try:
            r = self.s3.put_object_tagging(Bucket=self.bucket, Key=obj,
                                           Tagging=tagset)
        except:
            self.logger.info('put tags %s failed' % (tagset))
            return -1

        r = self.s3.get_object_tagging(Bucket=self.bucket, Key=obj)
        self.logger.info('get tags returned %s' % r)
        return 0

    def test_object_tags(self, obj='obj1'):

        tagset = self.make_tags({'foo':None,
                            'bar':'baz',
                            'key1':'val1'})
        self.test_put_get_tags(obj,tagset)

    def test_empty_tags(self, obj='obj2'):
        self.test_put_get_tags(obj, self.make_tags({}))

    def test_max_tags(self, obj='maxtags'):
        d = {}
        for i in range(10):
            k = make_random_string(128)
            v = make_random_string(256)
            d[k]=v # potential collision?

        tagset = self.make_tags(d)
        r = self.test_put_get_tags(obj, tagset)

    def test_exceeding_tags_size(self, obj='maxtags2'):
        d = {}
        for i in range(11):
            k = make_random_string(128)
            v = make_random_string(256)
            d[k]=v # potential collision?

        tagset = self.make_tags(d)
        r = self.test_put_get_tags(obj, tagset)

    def test_exceeding_tags_kv(self, obj='maxtags3'):
        d = {}
        for i in range(10):
            k = make_random_string(129)
            v = make_random_string(257)
            d[k]=v # potential collision?

        tagset = self.make_tags(d)
        r = self.test_put_get_tags(obj, tagset)

    def cleanup(self):
        r = self.s3.list_objects_v2(Bucket=self.bucket)
        for key in r['Contents']:
            self.s3.delete_object(Bucket=self.bucket, Key=key['Key'])

        self.s3.delete_bucket(Bucket=self.bucket)

@pytest.fixture
def s3test():
    return S3("access","secret","http://localhost:8000","testtagbucket")

def test_object_tags(s3test):
    assert s3test.test_object_tags() == 0


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='test s3 tags')
    parser.add_argument('--access', type=str, default='access',
                        help='AWS access key id')
    parser.add_argument('--secret', type=str, default='secret',
                        help='AWS Secret Access Key')
    parser.add_argument('--endpoint', type=str, default='http://localhost:8000',
                        help='S3 host')
    parser.add_argument('--bucket', type=str, default = 'bucket1',
                        help='bucket name')
    parser.add_argument('--cleanup', type=bool, default = False,
                        help='delete bucket and objects after test (default: False)')
    args = parser.parse_args()
    s3test = S3(args.access, args.secret, args.endpoint, args.bucket)
    s3test.test_object_tags()
    s3test.test_empty_tags()
    s3test.test_max_tags()
    s3test.test_exceeding_tags_size()
    s3test.test_exceeding_tags_kv()
    if args.cleanup:
        s3test.cleanup()