Python consumer.Consumer类代码示例

 def test_on_task_no_events(self):
     l = MainConsumer(self.ready_queue, self.eta_schedule, send_events=False)
     task = Mock()
     task.revoked.return_value = False
     l.event_dispatcher = Mock()
     l.event_dispatcher.enabled = False
     task.eta = None
     l._does_info = False
     l.on_task(task)

 def test_on_task_no_events(self):
     l = Consumer(self.ready_queue, timer=self.timer)
     task = Mock()
     task.revoked.return_value = False
     l.event_dispatcher = Mock()
     l.event_dispatcher.enabled = False
     task.eta = None
     l._does_info = False
     l.on_task(task)

    def test_on_decode_error(self, crit):
        l = Consumer(self.ready_queue, timer=self.timer)

        class MockMessage(Mock):
            content_type = 'application/x-msgpack'
            content_encoding = 'binary'
            body = 'foobarbaz'

        message = MockMessage()
        l.on_decode_error(message, KeyError('foo'))
        self.assertTrue(message.ack.call_count)
        self.assertIn("Can't decode message body", crit.call_args[0][0])

    def test_receieve_message(self):
        l = Consumer(self.ready_queue, timer=self.timer)
        m = create_message(Mock(), task=foo_task.name,
                           args=[2, 4, 8], kwargs={})
        l.update_strategies()
        callback = self._get_on_message(l)
        callback(m.decode(), m)

        in_bucket = self.ready_queue.get_nowait()
        self.assertIsInstance(in_bucket, Request)
        self.assertEqual(in_bucket.name, foo_task.name)
        self.assertEqual(in_bucket.execute(), 2 * 4 * 8)
        self.assertTrue(self.timer.empty())

 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 get_consumer(self, no_hub=False, **kwargs):
     consumer = Consumer(
         on_task_request=Mock(),
         init_callback=Mock(),
         pool=Mock(),
         app=self.app,
         timer=Mock(),
         controller=Mock(),
         hub=None if no_hub else Mock(),
         **kwargs
     )
     consumer.blueprint = Mock()
     consumer._restart_state = Mock()
     consumer.connection = Mock()
     consumer.connection_errors = (socket.error, OSError, )
     return consumer

 def __init__(self, app, heartbeat=None, on_task=None):
     (
         self.obj,
         self.connection,
         self.consumer,
         self.blueprint,
         self.hub,
         self.qos,
         self.heartbeat,
         self.clock,
     ) = self.args = [Mock(name='obj'),
                      Mock(name='connection'),
                      Mock(name='consumer'),
                      Mock(name='blueprint'),
                      Mock(name='Hub'),
                      Mock(name='qos'),
                      heartbeat,
                      Mock(name='clock')]
     self.connection.supports_heartbeats = True
     self.consumer.callbacks = []
     self.obj.strategies = {}
     self.connection.connection_errors = (socket.error, )
     #hent = self.Hub.__enter__ = Mock(name='Hub.__enter__')
     #self.Hub.__exit__ = Mock(name='Hub.__exit__')
     #self.hub = hent.return_value = Mock(name='hub_context')
     self.hub.on_task = on_task or []
     self.hub.readers = {}
     self.hub.writers = {}
     self.hub.consolidate = set()
     self.hub.fire_timers.return_value = 1.7
     self.Hub = self.hub
     # need this for create_task_handler
     _consumer = Consumer(Mock(), timer=Mock(), app=app)
     self.obj.create_task_handler = _consumer.create_task_handler
     self.on_unknown_message = self.obj.on_unknown_message = Mock(
         name='on_unknown_message',
     )
     _consumer.on_unknown_message = self.on_unknown_message
     self.on_unknown_task = self.obj.on_unknown_task = Mock(
         name='on_unknown_task',
     )
     _consumer.on_unknown_task = self.on_unknown_task
     self.on_invalid_task = self.obj.on_invalid_task = Mock(
         name='on_invalid_task',
     )
     _consumer.on_invalid_task = self.on_invalid_task
     _consumer.strategies = self.obj.strategies

    def test_receieve_message_ack_raises(self, logger, warn):
        l = Consumer(self.ready_queue, timer=self.timer)
        backend = Mock()
        m = create_message(backend, args=[2, 4, 8], kwargs={})

        l.event_dispatcher = Mock()
        l.connection_errors = (socket.error, )
        m.reject = Mock()
        m.reject.side_effect = socket.error('foo')
        callback = self._get_on_message(l)
        self.assertFalse(callback(m.decode(), m))
        self.assertTrue(warn.call_count)
        with self.assertRaises(Empty):
            self.ready_queue.get_nowait()
        self.assertTrue(self.timer.empty())
        m.reject.assert_called_with()
        self.assertTrue(logger.critical.call_count)

    def test_loop(self):

        class Connection(current_app.connection().__class__):
            obj = None

            def drain_events(self, **kwargs):
                self.obj.connection = None

        l = Consumer(self.ready_queue, timer=self.timer)
        l.connection = Connection()
        l.connection.obj = l
        l.task_consumer = Mock()
        l.qos = QoS(l.task_consumer, 10)

        l.loop(*l.loop_args())
        l.loop(*l.loop_args())
        self.assertTrue(l.task_consumer.consume.call_count)
        l.task_consumer.qos.assert_called_with(prefetch_count=10)
        l.task_consumer.qos = Mock()
        self.assertEqual(l.qos.value, 10)
        l.qos.decrement_eventually()
        self.assertEqual(l.qos.value, 9)
        l.qos.update()
        self.assertEqual(l.qos.value, 9)
        l.task_consumer.qos.assert_called_with(prefetch_count=9)

    def test_loop_when_socket_error(self):

        class Connection(current_app.connection().__class__):
            obj = None

            def drain_events(self, **kwargs):
                self.obj.connection = None
                raise socket.error('foo')

        l = Consumer(self.ready_queue, timer=self.timer)
        l.namespace.state = RUN
        c = l.connection = Connection()
        l.connection.obj = l
        l.task_consumer = Mock()
        l.qos = QoS(l.task_consumer, 10)
        with self.assertRaises(socket.error):
            l.loop(*l.loop_args())

        l.namespace.state = CLOSE
        l.connection = c
        l.loop(*l.loop_args())

 def __init__(self, app, heartbeat=None, on_task_message=None,
              transport_driver_type=None):
     hub = Hub()
     (
         self.obj,
         self.connection,
         self.consumer,
         self.blueprint,
         self.hub,
         self.qos,
         self.heartbeat,
         self.clock,
     ) = self.args = [Mock(name='obj'),
                      Mock(name='connection'),
                      Mock(name='consumer'),
                      Mock(name='blueprint'),
                      hub,
                      Mock(name='qos'),
                      heartbeat,
                      Mock(name='clock')]
     self.connection.supports_heartbeats = True
     self.connection.get_heartbeat_interval.side_effect = (
         lambda: self.heartbeat
     )
     self.consumer.callbacks = []
     self.obj.strategies = {}
     self.connection.connection_errors = (socket.error,)
     if transport_driver_type:
         self.connection.transport.driver_type = transport_driver_type
     self.hub.readers = {}
     self.hub.writers = {}
     self.hub.consolidate = set()
     self.hub.timer = Mock(name='hub.timer')
     self.hub.timer._queue = [Mock()]
     self.hub.fire_timers = Mock(name='hub.fire_timers')
     self.hub.fire_timers.return_value = 1.7
     self.hub.poller = Mock(name='hub.poller')
     self.hub.close = Mock(name='hub.close()')  # asynloop calls hub.close
     self.Hub = self.hub
     self.blueprint.state = RUN
     # need this for create_task_handler
     _consumer = Consumer(Mock(), timer=Mock(), app=app)
     _consumer.on_task_message = on_task_message or []
     self.obj.create_task_handler = _consumer.create_task_handler
     self.on_unknown_message = self.obj.on_unknown_message = Mock(
         name='on_unknown_message',
     )
     _consumer.on_unknown_message = self.on_unknown_message
     self.on_unknown_task = self.obj.on_unknown_task = Mock(
         name='on_unknown_task',
     )
     _consumer.on_unknown_task = self.on_unknown_task
     self.on_invalid_task = self.obj.on_invalid_task = Mock(
         name='on_invalid_task',
     )
     _consumer.on_invalid_task = self.on_invalid_task
     _consumer.strategies = self.obj.strategies

 def __init__(self, app, heartbeat=None, on_task_message=None):
     hub = Hub()
     (
         self.obj,
         self.connection,
         self.consumer,
         self.blueprint,
         self.hub,
         self.qos,
         self.heartbeat,
         self.clock,
     ) = self.args = [
         Mock(name="obj"),
         Mock(name="connection"),
         Mock(name="consumer"),
         Mock(name="blueprint"),
         hub,
         Mock(name="qos"),
         heartbeat,
         Mock(name="clock"),
     ]
     self.connection.supports_heartbeats = True
     self.connection.get_heartbeat_interval.side_effect = lambda: self.heartbeat
     self.consumer.callbacks = []
     self.obj.strategies = {}
     self.connection.connection_errors = (socket.error,)
     self.hub.readers = {}
     self.hub.writers = {}
     self.hub.consolidate = set()
     self.hub.timer = Mock(name="hub.timer")
     self.hub.timer._queue = [Mock()]
     self.hub.fire_timers = Mock(name="hub.fire_timers")
     self.hub.fire_timers.return_value = 1.7
     self.hub.poller = Mock(name="hub.poller")
     self.hub.close = Mock(name="hub.close()")  # asynloop calls hub.close
     self.Hub = self.hub
     self.blueprint.state = RUN
     # need this for create_task_handler
     _consumer = Consumer(Mock(), timer=Mock(), app=app)
     _consumer.on_task_message = on_task_message or []
     self.obj.create_task_handler = _consumer.create_task_handler
     self.on_unknown_message = self.obj.on_unknown_message = Mock(name="on_unknown_message")
     _consumer.on_unknown_message = self.on_unknown_message
     self.on_unknown_task = self.obj.on_unknown_task = Mock(name="on_unknown_task")
     _consumer.on_unknown_task = self.on_unknown_task
     self.on_invalid_task = self.obj.on_invalid_task = Mock(name="on_invalid_task")
     _consumer.on_invalid_task = self.on_invalid_task
     _consumer.strategies = self.obj.strategies

 def test_on_task_revoked(self):
     l = Consumer(self.ready_queue, timer=self.timer)
     task = Mock()
     task.revoked.return_value = True
     l.on_task(task)

    def test_reset_connection_with_no_node(self):

        l = MainConsumer(self.ready_queue, self.eta_schedule, self.logger)
        self.assertEqual(None, l.pool)
        l.reset_connection()

import sys

@task(accept_magic_kwargs=False)
def T():
    pass

tid = uuid()
P = TaskPool()
hostname = socket.gethostname()
task = {'task': T.name, 'args': (), 'kwargs': {}, 'id': tid, 'flags': 0}
app = current_app._get_current_object()
ready_queue = Queue()

def on_put(req):
    req.execute_using_pool(P)

def on_ack(*a): pass


m = Message(None, {}, {}, task)

ready_queue.put = on_put
x = Consumer(ready_queue, hostname=hostname, app=app)
x.update_strategies()
name = T.name
ts = time()
for i in range(100000):
    x.strategies[name](m, m.body, on_ack)
print(time() - ts)