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 __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.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 self._consumer = _consumer = Consumer(Mock(), timer=Mock(), controller=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 self.on_decode_error = self.obj.on_decode_error = Mock( name='on_decode_error', ) _consumer.on_decode_error = self.on_decode_error _consumer.strategies = self.obj.strategies
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 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 = _amqp_connection() 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.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_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.qos, 10) with self.assertRaises(socket.error): l.loop(*l.loop_args()) l.namespace.state = CLOSE l.connection = c l.loop(*l.loop_args())
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.qos, 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) 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)
@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() def on_task(req): req.execute_using_pool(P) def on_ack(*a): pass m = Message(None, {}, {}, task) x = Consumer(on_task, 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)
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)
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 = Consumer(self.ready_queue, timer=self.timer) l.steps.pop() self.assertEqual(None, l.pool) l.namespace.start(l)