def test_receieve_message_eta_isoformat(self): class MockConsumer(object): prefetch_count_incremented = False def qos(self, **kwargs): self.prefetch_count_incremented = True l = MyKombuConsumer(self.ready_queue, self.eta_schedule, self.logger, send_events=False) backend = MockBackend() m = create_message(backend, task=foo_task.name, eta=datetime.now().isoformat(), args=[2, 4, 8], kwargs={}) l.task_consumer = MockConsumer() l.qos = QoS(l.task_consumer, l.initial_prefetch_count, l.logger) l.event_dispatcher = MockEventDispatcher() l.receive_message(m.decode(), m) l.eta_schedule.stop() items = [entry[2] for entry in self.eta_schedule.queue] found = 0 for item in items: if item.args[0].task_name == foo_task.name: found = True self.assertTrue(found) self.assertTrue(l.task_consumer.prefetch_count_incremented) l.eta_schedule.stop()
def test_receieve_message_eta_isoformat(self): l = MyKombuConsumer(self.ready_queue, self.eta_schedule, self.logger, send_events=False) m = create_message(Mock(), task=foo_task.name, eta=datetime.now().isoformat(), args=[2, 4, 8], kwargs={}) l.task_consumer = Mock() l.qos = QoS(l.task_consumer, l.initial_prefetch_count, l.logger) l.event_dispatcher = Mock() l.enabled = False l.update_strategies() l.receive_message(m.decode(), m) l.eta_schedule.stop() items = [entry[2] for entry in self.eta_schedule.queue] found = 0 for item in items: if item.args[0].name == foo_task.name: found = True self.assertTrue(found) self.assertTrue(l.task_consumer.qos.call_count) l.eta_schedule.stop()
def test_consumer_decrement_eventually(self): consumer = Mock() qos = QoS(consumer, 10, current_app.log.get_default_logger()) qos.decrement_eventually() self.assertEqual(qos.value, 9) qos.value = 0 qos.decrement_eventually() self.assertEqual(qos.value, 0)
def test_consumer_decrement_eventually(self): consumer = Mock() qos = QoS(consumer, 10) qos.decrement_eventually() self.assertEqual(qos.value, 9) qos.value = 0 qos.decrement_eventually() self.assertEqual(qos.value, 0)
def test_info(self): l = MyKombuConsumer(self.ready_queue, timer=self.timer) l.qos = QoS(l.task_consumer, 10) info = l.info self.assertEqual(info['prefetch_count'], 10) self.assertFalse(info['broker']) l.connection = current_app.connection() info = l.info self.assertTrue(info['broker'])
def test_info(self): l = MyKombuConsumer(self.ready_queue, self.eta_schedule, self.logger, send_events=False) l.qos = QoS(l.task_consumer, 10, l.logger) info = l.info self.assertEqual(info["prefetch_count"], 10) self.assertFalse(info["broker"]) l.connection = current_app.broker_connection() info = l.info self.assertTrue(info["broker"])
def test_apply_eta_task(self): from celery.worker import state l = MyKombuConsumer(self.ready_queue, timer=self.timer) l.qos = QoS(None, 10) task = object() qos = l.qos.value l.apply_eta_task(task) self.assertIn(task, state.reserved_requests) self.assertEqual(l.qos.value, qos - 1) self.assertIs(self.ready_queue.get_nowait(), task)
def test_apply_eta_task(self): from celery.worker import state l = MyKombuConsumer(self.ready_queue, self.eta_schedule, self.logger, send_events=False) l.qos = QoS(None, 10, l.logger) task = object() qos = l.qos.value l.apply_eta_task(task) self.assertIn(task, state.reserved_requests) self.assertEqual(l.qos.value, qos - 1) self.assertIs(self.ready_queue.get_nowait(), task)
def test_decrement(self): consumer = self.MockConsumer() qos = QoS(consumer, 10, app_or_default().log.get_default_logger()) qos.update() self.assertEqual(int(qos.value), 10) self.assertEqual(consumer.prefetch_count, 10) qos.decrement() self.assertEqual(int(qos.value), 9) self.assertEqual(consumer.prefetch_count, 9) qos.decrement_eventually() self.assertEqual(int(qos.value), 8) self.assertEqual(consumer.prefetch_count, 9)
def test_consume_messages_ignores_socket_timeout(self): class Connection(current_app.connection().__class__): obj = None def drain_events(self, **kwargs): self.obj.connection = None raise socket.timeout(10) l = MyKombuConsumer(self.ready_queue, timer=self.timer) l.connection = Connection() l.task_consumer = Mock() l.connection.obj = l l.qos = QoS(l.task_consumer, 10) l.consume_messages()
def test_consume_messages_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 = MyKombuConsumer(self.ready_queue, timer=self.timer) l._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.consume_messages() l._state = CLOSE l.connection = c l.consume_messages()
def test_consume_messages(self): class Connection(current_app.broker_connection().__class__): obj = None def drain_events(self, **kwargs): self.obj.connection = None l = MyKombuConsumer(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.consume_messages() l.consume_messages() self.assertTrue(l.task_consumer.consume.call_count) l.task_consumer.qos.assert_called_with(prefetch_count=10) l.qos.decrement() l.consume_messages() l.task_consumer.qos.assert_called_with(prefetch_count=9)
def test_consume_messages_when_socket_error(self): class Connection(current_app.broker_connection().__class__): obj = None def drain_events(self, **kwargs): self.obj.connection = None raise socket.error("foo") l = MyKombuConsumer(self.ready_queue, self.eta_schedule, self.logger, send_events=False) l._state = RUN c = l.connection = Connection() l.connection.obj = l l.task_consumer = Mock() l.qos = QoS(l.task_consumer, 10, l.logger) with self.assertRaises(socket.error): l.consume_messages() l._state = CLOSE l.connection = c l.consume_messages()
def test_consume_messages(self): class Connection(current_app.broker_connection().__class__): obj = None def drain_events(self, **kwargs): self.obj.connection = None class Consumer(object): consuming = False prefetch_count = 0 def consume(self): self.consuming = True def qos(self, prefetch_size=0, prefetch_count=0, apply_global=False): self.prefetch_count = prefetch_count l = MyKombuConsumer(self.ready_queue, self.eta_schedule, self.logger, send_events=False) l.connection = Connection() l.connection.obj = l l.task_consumer = Consumer() l.qos = QoS(l.task_consumer, 10, l.logger) l.consume_messages() l.consume_messages() self.assertTrue(l.task_consumer.consuming) self.assertEqual(l.task_consumer.prefetch_count, 10) l.qos.decrement() l.consume_messages() self.assertEqual(l.task_consumer.prefetch_count, 9)
def test_exceeds_short(self): qos = QoS(Mock(), PREFETCH_COUNT_MAX - 1) qos.update() self.assertEqual(qos.value, PREFETCH_COUNT_MAX - 1) qos.increment() self.assertEqual(qos.value, PREFETCH_COUNT_MAX) qos.increment() self.assertEqual(qos.value, PREFETCH_COUNT_MAX + 1) qos.decrement() self.assertEqual(qos.value, PREFETCH_COUNT_MAX) qos.decrement() self.assertEqual(qos.value, PREFETCH_COUNT_MAX - 1)
def test_consumer_increment_decrement(self): consumer = self.MockConsumer() qos = QoS(consumer, 10, current_app.log.get_default_logger()) qos.update() self.assertEqual(qos.value, 10) self.assertEqual(consumer.prefetch_count, 10) qos.decrement() self.assertEqual(qos.value, 9) self.assertEqual(consumer.prefetch_count, 9) qos.decrement_eventually() self.assertEqual(qos.value, 8) self.assertEqual(consumer.prefetch_count, 9) # Does not decrement 0 value qos.value = 0 qos.decrement() self.assertEqual(qos.value, 0) qos.increment() self.assertEqual(qos.value, 0)
def test_set(self): consumer = Mock() qos = QoS(consumer, 10) qos.set(12) self.assertEqual(qos.prev, 12) qos.set(qos.prev)
def test_exceeds_short(self): qos = QoS(Mock(), PREFETCH_COUNT_MAX - 1) qos.update() self.assertEqual(qos.value, PREFETCH_COUNT_MAX - 1) qos.increment_eventually() self.assertEqual(qos.value, PREFETCH_COUNT_MAX) qos.increment_eventually() self.assertEqual(qos.value, PREFETCH_COUNT_MAX + 1) qos.decrement_eventually() self.assertEqual(qos.value, PREFETCH_COUNT_MAX) qos.decrement_eventually() self.assertEqual(qos.value, PREFETCH_COUNT_MAX - 1)
def test_consumer_increment_decrement(self): consumer = Mock() qos = QoS(consumer, 10) qos.update() self.assertEqual(qos.value, 10) consumer.qos.assert_called_with(prefetch_count=10) qos.decrement_eventually() qos.update() self.assertEqual(qos.value, 9) consumer.qos.assert_called_with(prefetch_count=9) qos.decrement_eventually() self.assertEqual(qos.value, 8) consumer.qos.assert_called_with(prefetch_count=9) self.assertIn({'prefetch_count': 9}, consumer.qos.call_args) # Does not decrement 0 value qos.value = 0 qos.decrement_eventually() self.assertEqual(qos.value, 0) qos.increment_eventually() self.assertEqual(qos.value, 0)
def test_exceeds_short(self): qos = QoS(Mock(), PREFETCH_COUNT_MAX - 1, current_app.log.get_default_logger()) qos.update() self.assertEqual(qos.value, PREFETCH_COUNT_MAX - 1) qos.increment() self.assertEqual(qos.value, PREFETCH_COUNT_MAX) qos.increment() self.assertEqual(qos.value, PREFETCH_COUNT_MAX + 1) qos.decrement() self.assertEqual(qos.value, PREFETCH_COUNT_MAX) qos.decrement() self.assertEqual(qos.value, PREFETCH_COUNT_MAX - 1)
def __init__(self, value): self.value = value QoS.__init__(self, None, value)
def test_consumer_increment_decrement(self): consumer = Mock() qos = QoS(consumer, 10, current_app.log.get_default_logger()) qos.update() self.assertEqual(qos.value, 10) self.assertIn({"prefetch_count": 10}, consumer.qos.call_args) qos.decrement() self.assertEqual(qos.value, 9) self.assertIn({"prefetch_count": 9}, consumer.qos.call_args) qos.decrement_eventually() self.assertEqual(qos.value, 8) self.assertIn({"prefetch_count": 9}, consumer.qos.call_args) # Does not decrement 0 value qos.value = 0 qos.decrement() self.assertEqual(qos.value, 0) qos.increment() self.assertEqual(qos.value, 0)
def test_set(self): consumer = Mock() qos = QoS(consumer, 10, current_app.log.get_default_logger()) qos.set(12) self.assertEqual(qos.prev, 12) qos.set(qos.prev)
def test_consumer_increment_decrement(self): consumer = Mock() qos = QoS(consumer, 10) qos.update() self.assertEqual(qos.value, 10) consumer.qos.assert_called_with(prefetch_count=10) qos.decrement() self.assertEqual(qos.value, 9) consumer.qos.assert_called_with(prefetch_count=9) qos.decrement_eventually() self.assertEqual(qos.value, 8) consumer.qos.assert_called_with(prefetch_count=9) self.assertIn({"prefetch_count": 9}, consumer.qos.call_args) # Does not decrement 0 value qos.value = 0 qos.decrement() self.assertEqual(qos.value, 0) qos.increment() self.assertEqual(qos.value, 0)