Python five.Queue类代码示例

    def test_mediator_body_exception(self):
        ready_queue = Queue()

        def mycallback(value):
            raise KeyError('foo')

        m = Mediator(ready_queue, mycallback)
        ready_queue.put(MockTask('Elaine M. Benes'))

        m.body()

    def test_run(self):
        ready_queue = Queue()

        condition = [None]

        def mycallback(value):
            condition[0].set()

        m = Mediator(ready_queue, mycallback)
        condition[0] = m._is_shutdown
        ready_queue.put(MockTask('Elaine M. Benes'))

        m.run()
        self.assertTrue(m._is_shutdown.isSet())
        self.assertTrue(m._is_stopped.isSet())

    def setup(self):
        self.buffer = FastQueue()
        self.timer = Timer()

        @self.app.task(shared=False)
        def foo_task(x, y, z):
            return x * y * z
        self.foo_task = foo_task

    def test_mediator_body(self):
        ready_queue = Queue()
        got = {}

        def mycallback(value):
            got['value'] = value.value

        m = Mediator(ready_queue, mycallback)
        ready_queue.put(MockTask('George Costanza'))

        m.body()

        self.assertEqual(got['value'], 'George Costanza')

        ready_queue.put(MockTask('Jerry Seinfeld'))
        m._does_debug = False
        m.body()
        self.assertEqual(got['value'], 'Jerry Seinfeld')

    def test_poll_result(self):

        results = Queue()

        class Message(object):

            def __init__(self, **merge):
                self.payload = dict({'status': states.STARTED,
                                     'result': None}, **merge)

        class MockBinding(object):

            def __init__(self, *args, **kwargs):
                pass

            def __call__(self, *args, **kwargs):
                return self

            def declare(self):
                pass

            def get(self, no_ack=False):
                try:
                    return results.get(block=False)
                except Empty:
                    pass

        class MockBackend(AMQPBackend):
            Queue = MockBinding

        backend = MockBackend()

        # FFWD's to the latest state.
        results.put(Message(status=states.RECEIVED, seq=1))
        results.put(Message(status=states.STARTED, seq=2))
        results.put(Message(status=states.FAILURE, seq=3))
        r1 = backend.get_task_meta(uuid())
        self.assertDictContainsSubset({'status': states.FAILURE,
                                       'seq': 3}, r1,
                                       'FFWDs to the last state')

        # Caches last known state.
        results.put(Message())
        tid = uuid()
        backend.get_task_meta(tid)
        self.assertIn(tid, backend._cache, 'Caches last known state')

        # Returns cache if no new states.
        results.queue.clear()
        assert not results.qsize()
        backend._cache[tid] = 'hello'
        self.assertEqual(backend.get_task_meta(tid), 'hello',
                         'Returns cache if no new states')

    def test_mediator_crash(self, _exit):
        ms = [None]

        class _Mediator(Mediator):

            def body(self):
                try:
                    raise KeyError('foo')
                finally:
                    ms[0]._is_shutdown.set()

        ready_queue = Queue()
        ms[0] = m = _Mediator(ready_queue, None)
        ready_queue.put(MockTask('George Constanza'))

        stderr = Mock()
        p, sys.stderr = sys.stderr, stderr
        try:
            m.run()
        finally:
            sys.stderr = p
        self.assertTrue(_exit.call_count)
        self.assertTrue(stderr.write.call_count)

class test_Consumer(AppCase):

    def setup(self):
        self.buffer = FastQueue()
        self.timer = Timer()

        @self.app.task(shared=False)
        def foo_task(x, y, z):
            return x * y * z
        self.foo_task = foo_task

    def teardown(self):
        self.timer.stop()

    def test_info(self):
        l = MyKombuConsumer(self.buffer.put, timer=self.timer, app=self.app)
        l.task_consumer = Mock()
        l.qos = QoS(l.task_consumer.qos, 10)
        l.connection = Mock()
        l.connection.info.return_value = {'foo': 'bar'}
        l.controller = l.app.WorkController()
        l.controller.pool = Mock()
        l.controller.pool.info.return_value = [Mock(), Mock()]
        l.controller.consumer = l
        info = l.controller.stats()
        self.assertEqual(info['prefetch_count'], 10)
        self.assertTrue(info['broker'])

    def test_start_when_closed(self):
        l = MyKombuConsumer(self.buffer.put, timer=self.timer, app=self.app)
        l.blueprint.state = CLOSE
        l.start()

    def test_connection(self):
        l = MyKombuConsumer(self.buffer.put, timer=self.timer, app=self.app)

        l.blueprint.start(l)
        self.assertIsInstance(l.connection, Connection)

        l.blueprint.state = RUN
        l.event_dispatcher = None
        l.blueprint.restart(l)
        self.assertTrue(l.connection)

        l.blueprint.state = RUN
        l.shutdown()
        self.assertIsNone(l.connection)
        self.assertIsNone(l.task_consumer)

        l.blueprint.start(l)
        self.assertIsInstance(l.connection, Connection)
        l.blueprint.restart(l)

        l.stop()
        l.shutdown()
        self.assertIsNone(l.connection)
        self.assertIsNone(l.task_consumer)

    def test_close_connection(self):
        l = MyKombuConsumer(self.buffer.put, timer=self.timer, app=self.app)
        l.blueprint.state = RUN
        step = find_step(l, consumer.Connection)
        conn = l.connection = Mock()
        step.shutdown(l)
        self.assertTrue(conn.close.called)
        self.assertIsNone(l.connection)

        l = MyKombuConsumer(self.buffer.put, timer=self.timer, app=self.app)
        eventer = l.event_dispatcher = Mock()
        eventer.enabled = True
        heart = l.heart = MockHeart()
        l.blueprint.state = RUN
        Events = find_step(l, consumer.Events)
        Events.shutdown(l)
        Heart = find_step(l, consumer.Heart)
        Heart.shutdown(l)
        self.assertTrue(eventer.close.call_count)
        self.assertTrue(heart.closed)

    @patch('celery.worker.consumer.warn')
    def test_receive_message_unknown(self, warn):
        l = _MyKombuConsumer(self.buffer.put, timer=self.timer, app=self.app)
        l.blueprint.state = RUN
        l.steps.pop()
        backend = Mock()
        m = create_message(backend, unknown={'baz': '!!!'})
        l.event_dispatcher = Mock()
        l.node = MockNode()

        callback = self._get_on_message(l)
        callback(m.decode(), m)
        self.assertTrue(warn.call_count)

    @patch('celery.worker.strategy.to_timestamp')
    def test_receive_message_eta_OverflowError(self, to_timestamp):
        to_timestamp.side_effect = OverflowError()
        l = _MyKombuConsumer(self.buffer.put, timer=self.timer, app=self.app)
        l.blueprint.state = RUN
        l.steps.pop()
        m = create_message(Mock(), task=self.foo_task.name,
#.........这里部分代码省略.........

 def __init__(self, fill_rate, queue=None, capacity=1):
     self._bucket = TokenBucket(fill_rate, capacity)
     self.queue = queue
     if not self.queue:
         self.queue = Queue()

class TokenBucketQueue(object):
    """Queue with rate limited get operations.

    This uses the token bucket algorithm to rate limit the queue on get
    operations.

    :param fill_rate: The rate in tokens/second that the bucket will
                      be refilled.
    :keyword capacity: Maximum number of tokens in the bucket.
                       Default is 1.

    """
    RateLimitExceeded = RateLimitExceeded

    def __init__(self, fill_rate, queue=None, capacity=1):
        self._bucket = TokenBucket(fill_rate, capacity)
        self.queue = queue
        if not self.queue:
            self.queue = Queue()

    def put(self, item, block=True):
        """Put an item onto the queue."""
        self.queue.put(item, block=block)

    def put_nowait(self, item):
        """Put an item into the queue without blocking.

        :raises Queue.Full: If a free slot is not immediately available.

        """
        return self.put(item, block=False)

    def get(self, block=True):
        """Remove and return an item from the queue.

        :raises RateLimitExceeded: If a token could not be consumed from the
                                   token bucket (consuming from the queue
                                   too fast).
        :raises Queue.Empty: If an item is not immediately available.

        """
        get = block and self.queue.get or self.queue.get_nowait

        if not block and not self.items:
            raise Empty()

        if not self._bucket.can_consume(1):
            raise RateLimitExceeded()

        return get()

    def get_nowait(self):
        """Remove and return an item from the queue without blocking.

        :raises RateLimitExceeded: If a token could not be consumed from the
                                   token bucket (consuming from the queue
                                   too fast).
        :raises Queue.Empty: If an item is not immediately available.

        """
        return self.get(block=False)

    def qsize(self):
        """Returns the size of the queue."""
        return self.queue.qsize()

    def empty(self):
        """Returns :const:`True` if the queue is empty."""
        return self.queue.empty()

    def clear(self):
        """Delete all data in the queue."""
        return self.items.clear()

    def wait(self, block=False):
        """Wait until a token can be retrieved from the bucket and return
        the next item."""
        get = self.get
        expected_time = self.expected_time
        while 1:
            remaining = expected_time()
            if not remaining:
                return get(block=block)
            sleep(remaining)

    def expected_time(self, tokens=1):
        """Returns the expected time in seconds of when a new token should be
        available."""
        if not self.items:
            return 0
        return self._bucket.expected_time(tokens)

    @property
    def items(self):
        """Underlying data.  Do not modify."""
        return self.queue.queue

 def __init__(self):
     self._buffer = Queue()
     self._count = count(1)
     self._tref = None
     self._pool = None

class Batches(Task):
    abstract = True

    #: Maximum number of message in buffer.
    flush_every = 10

    #: Timeout in seconds before buffer is flushed anyway.
    flush_interval = 30

    def __init__(self):
        self._buffer = Queue()
        self._count = count(1)
        self._tref = None
        self._pool = None

    def run(self, requests):
        raise NotImplementedError('must implement run(requests)')

    def Strategy(self, task, app, consumer):
        self._pool = consumer.pool
        hostname = consumer.hostname
        eventer = consumer.event_dispatcher
        Req = Request
        connection_errors = consumer.connection_errors
        timer = consumer.timer
        put_buffer = self._buffer.put
        flush_buffer = self._do_flush
        body_can_be_buffer = consumer.pool.body_can_be_buffer

        def task_message_handler(message, body, ack, reject, callbacks, **kw):
            if body is None:
                body, headers, decoded, utc = (
                    message.body, message.headers, False, True,
                )
                if not body_can_be_buffer:
                    body = bytes(body) if isinstance(body, buffer_t) else body
            else:
                body, headers, decoded, utc = proto1_to_proto2(message, body)

            request = Req(
                message,
                on_ack=ack, on_reject=reject, app=app, hostname=hostname,
                eventer=eventer, task=task,
                body=body, headers=headers, decoded=decoded, utc=utc,
                connection_errors=connection_errors,
            )
            put_buffer(request)

            if self._tref is None:     # first request starts flush timer.
                self._tref = timer.call_repeatedly(
                    self.flush_interval, flush_buffer,
                )

            if not next(self._count) % self.flush_every:
                flush_buffer()

        return task_message_handler

    def flush(self, requests):
        return self.apply_buffer(requests, ([SimpleRequest.from_request(r)
                                             for r in requests],))

    def _do_flush(self):
        logger.debug('Batches: Wake-up to flush buffer...')
        requests = None
        if self._buffer.qsize():
            requests = list(consume_queue(self._buffer))
            if requests:
                logger.debug('Batches: Buffer complete: %s', len(requests))
                self.flush(requests)
        if not requests:
            logger.debug('Batches: Canceling timer: Nothing in buffer.')
            if self._tref:
                self._tref.cancel()  # cancel timer.
            self._tref = None

    def apply_buffer(self, requests, args=(), kwargs={}):
        acks_late = [], []
        [acks_late[r.task.acks_late].append(r) for r in requests]
        assert requests and (acks_late[True] or acks_late[False])

        def on_accepted(pid, time_accepted):
            [req.acknowledge() for req in acks_late[False]]

        def on_return(result):
            [req.acknowledge() for req in acks_late[True]]

        return self._pool.apply_async(
            apply_batches_task,
            (self, args, 0, None),
            accept_callback=on_accepted,
            callback=acks_late[True] and on_return or noop,
        )

 def setup(self):
     self.buffer = FastQueue()
     self.timer = Timer()

    def test_poll_result(self):

        results = Queue()

        class Message(object):

            def __init__(self, **merge):
                self.payload = dict({'status': states.STARTED,
                                     'result': None}, **merge)
                self.body = pickle.dumps(self.payload)
                self.content_type = 'application/x-python-serialize'
                self.content_encoding = 'binary'

        class MockBinding(object):

            def __init__(self, *args, **kwargs):
                self.channel = Mock()

            def __call__(self, *args, **kwargs):
                return self

            def declare(self):
                pass

            def get(self, no_ack=False):
                try:
                    return results.get(block=False)
                except Empty:
                    pass

            def is_bound(self):
                return True

        class MockBackend(AMQPBackend):
            Queue = MockBinding

        backend = MockBackend()
        backend._republish = Mock()

        # FFWD's to the latest state.
        results.put(Message(status=states.RECEIVED, seq=1))
        results.put(Message(status=states.STARTED, seq=2))
        results.put(Message(status=states.FAILURE, seq=3))
        r1 = backend.get_task_meta(uuid())
        self.assertDictContainsSubset({'status': states.FAILURE,
                                       'seq': 3}, r1,
                                      'FFWDs to the last state')

        # Caches last known state.
        results.put(Message())
        tid = uuid()
        backend.get_task_meta(tid)
        self.assertIn(tid, backend._cache, 'Caches last known state')

        self.assertTrue(backend._republish.called)

        # Returns cache if no new states.
        results.queue.clear()
        assert not results.qsize()
        backend._cache[tid] = 'hello'
        self.assertEqual(backend.get_task_meta(tid), 'hello',
                         'Returns cache if no new states')

class test_Consumer(ConsumerCase):

    def setup(self):
        self.buffer = FastQueue()
        self.timer = Timer()

        @self.app.task(shared=False)
        def foo_task(x, y, z):
            return x * y * z
        self.foo_task = foo_task

    def teardown(self):
        self.timer.stop()

    def LoopConsumer(self, buffer=None, controller=None, timer=None, app=None,
                     without_mingle=True, without_gossip=True,
                     without_heartbeat=True, **kwargs):
        if controller is None:
            controller = Mock(name='.controller')
        buffer = buffer if buffer is not None else self.buffer.put
        timer = timer if timer is not None else self.timer
        app = app if app is not None else self.app
        c = Consumer(
            buffer,
            timer=timer,
            app=app,
            controller=controller,
            without_mingle=without_mingle,
            without_gossip=without_gossip,
            without_heartbeat=without_heartbeat,
            **kwargs
        )
        c.task_consumer = Mock(name='.task_consumer')
        c.qos = QoS(c.task_consumer.qos, 10)
        c.connection = Mock(name='.connection')
        c.controller = c.app.WorkController()
        c.heart = Mock(name='.heart')
        c.controller.consumer = c
        c.pool = c.controller.pool = Mock(name='.controller.pool')
        c.node = Mock(name='.node')
        c.event_dispatcher = mock_event_dispatcher()
        return c

    def NoopConsumer(self, *args, **kwargs):
        c = self.LoopConsumer(*args, **kwargs)
        c.loop = Mock(name='.loop')
        return c

    def test_info(self):
        c = self.NoopConsumer()
        c.connection.info.return_value = {'foo': 'bar'}
        c.controller.pool.info.return_value = [Mock(), Mock()]
        info = c.controller.stats()
        assert info['prefetch_count'] == 10
        assert info['broker']

    def test_start_when_closed(self):
        c = self.NoopConsumer()
        c.blueprint.state = CLOSE
        c.start()

    def test_connection(self):
        c = self.NoopConsumer()

        c.blueprint.start(c)
        assert isinstance(c.connection, Connection)

        c.blueprint.state = RUN
        c.event_dispatcher = None
        c.blueprint.restart(c)
        assert c.connection

        c.blueprint.state = RUN
        c.shutdown()
        assert c.connection is None
        assert c.task_consumer is None

        c.blueprint.start(c)
        assert isinstance(c.connection, Connection)
        c.blueprint.restart(c)

        c.stop()
        c.shutdown()
        assert c.connection is None
        assert c.task_consumer is None

    def test_close_connection(self):
        c = self.NoopConsumer()
        c.blueprint.state = RUN
        step = find_step(c, consumer.Connection)
        connection = c.connection
        step.shutdown(c)
        connection.close.assert_called()
        assert c.connection is None

    def test_close_connection__heart_shutdown(self):
        c = self.NoopConsumer()
        event_dispatcher = c.event_dispatcher
        heart = c.heart
        c.event_dispatcher.enabled = True
#.........这里部分代码省略.........

class Batches(Task):
    abstract = True

    #: Maximum number of message in buffer.
    flush_every = 10

    #: Timeout in seconds before buffer is flushed anyway.
    flush_interval = 30

    def __init__(self):
        self._buffer = Queue()
        self._count = count(1)
        self._tref = None
        self._pool = None

    def run(self, requests):
        raise NotImplementedError("must implement run(requests)")

    def Strategy(self, task, app, consumer):
        self._pool = consumer.pool
        hostname = consumer.hostname
        eventer = consumer.event_dispatcher
        Req = Request
        connection_errors = consumer.connection_errors
        timer = consumer.timer
        put_buffer = self._buffer.put
        flush_buffer = self._do_flush

        def task_message_handler(message, body, ack):
            request = Req(
                body,
                on_ack=ack,
                app=app,
                hostname=hostname,
                events=eventer,
                task=task,
                connection_errors=connection_errors,
                delivery_info=message.delivery_info,
            )
            put_buffer(request)

            if self._tref is None:  # first request starts flush timer.
                self._tref = timer.apply_interval(self.flush_interval * 1000.0, flush_buffer)

            if not next(self._count) % self.flush_every:
                flush_buffer()

        return task_message_handler

    def flush(self, requests):
        return self.apply_buffer(requests, ([SimpleRequest.from_request(r) for r in requests],))

    def _do_flush(self):
        logger.debug("Batches: Wake-up to flush buffer...")
        requests = None
        if self._buffer.qsize():
            requests = list(consume_queue(self._buffer))
            if requests:
                logger.debug("Batches: Buffer complete: %s", len(requests))
                self.flush(requests)
        if not requests:
            logger.debug("Batches: Cancelling timer: Nothing in buffer.")
            self._tref.cancel()  # cancel timer.
            self._tref = None

    def apply_buffer(self, requests, args=(), kwargs={}):
        acks_late = [], []
        [acks_late[r.task.acks_late].append(r) for r in requests]
        assert requests and (acks_late[True] or acks_late[False])

        def on_accepted(pid, time_accepted):
            [req.acknowledge() for req in acks_late[False]]

        def on_return(result):
            [req.acknowledge() for req in acks_late[True]]

        return self._pool.apply_async(
            apply_batches_task,
            (self, args, 0, None),
            accept_callback=on_accepted,
            callback=acks_late[True] and on_return or None,
        )