s3-tests/s3tests/common/greenlets.py

import bunch
import collections
import gevent
import random
import time

from ..common import context, get_next_key
from ..common.results import TransferGreenletResult
from ..realistic import FileVerifier


# Make sure context has somewhere to store what we need
context.update(bunch.Bunch(
    needs_first_read=collections.deque(),
    all_keys=[],
    files_iter=None,
))


class SafeTransferGreenlet(gevent.Greenlet):
    def __init__(self, timeout=120):
        gevent.Greenlet.__init__(self)
        self.timeout = timeout
        self.result = None

    def _run(self):
        result = self.result = TransferGreenletResult(self.type)
        result.markStarted()

        try:
            with gevent.Timeout(self.timeout, False):
                result.success = self._doit()
        except gevent.GreenletExit:
            return
        except:
            result.setError(show_traceback=True)

        result.markFinished()


class ReaderGreenlet(SafeTransferGreenlet):
    type = 'reader'

    def _doit(self):
        if context.needs_first_read:
            key = context.needs_first_read.popleft()
        elif context.all_keys:
            key = random.choice(context.all_keys)
        else:
            time.sleep(1)
            return self.result.setError('No available keys to test with reader. Try again later.')

        # Copynew the key object
        key = key.bucket.new_key(key.name)
        self.result.setKey(key)

        fp = FileVerifier()

        key.get_contents_to_file(fp)

        self.result.request_finish = time.time()
        self.result.request_start = fp.created_at
        self.result.chunks = fp.chunks
        self.result.size = fp.size

        if not fp.valid():
            return self.result.setError('Failed to validate key {name!s}'.format(name=key.name))

        return True


class WriterGreenlet(SafeTransferGreenlet):
    type = 'writer'

    def _doit(self):
        key = get_next_key(context.bucket)
        self.result.setKey(key)

        fp = next(context.files_iter)
        self.result.size = fp.size

        key.set_contents_from_file(fp)

        self.result.request_finish = time.time()
        self.result.request_start = fp.start_time
        self.result.chunks = fp.last_chunks

        # And at the end, add to needs_first_read and shuffle
        context.needs_first_read.append(key)
        context.all_keys.append(key)

        return True