示例#1
0
 def test_sched_info(self):
     ready_queue = Queue()
     sched = Scheduler(ready_queue)
     item = MockItem("foo")
     sched.enter(item, eta=10, callback=None)
     self.assertDictEqual({"eta": 10, "priority": 0,
                           "item": item}, sched.info().next())
 def test_sched_info(self):
     ready_queue = Queue()
     sched = Scheduler(ready_queue)
     item = MockItem("foo")
     sched.enter(item, eta=10, callback=None)
     self.assertDictEqual({
         "eta": 10,
         "priority": 0,
         "item": item
     },
                          sched.info().next())
 def test_sched_clear(self):
     ready_queue = Queue()
     sched = Scheduler(ready_queue)
     sched.enter(MockItem("foo"), eta=datetime.now(), callback=None)
     self.assertFalse(sched.empty())
     sched.clear()
     self.assertTrue(sched.empty())
示例#4
0
    def test_sched_and_run_now(self):
        ready_queue = Queue()
        sched = Scheduler(ready_queue)
        now = datetime.now()

        callback_called = [False]
        def callback():
            callback_called[0] = True

        sched.enter("foo", eta=now, callback=callback)

        remaining = iter(sched).next()
        self.assertEquals(remaining, 0)
        self.assertTrue(callback_called[0])
        self.assertEquals(ready_queue.get_nowait(), "foo")
示例#5
0
    def test_sched_revoked(self):
        ready_queue = Queue()
        sched = Scheduler(ready_queue)
        now = datetime.now()

        callback_called = [False]
        def callback():
            callback_called[0] = True

        sched.enter(MockItem("foo", revoked=True), eta=now, callback=callback)
        iter(sched).next()
        self.assertFalse(callback_called[0])
        self.assertRaises(Empty, ready_queue.get_nowait)
        self.assertFalse(sched.queue)
        sched.clear()
    def test_sched_and_run_now(self):
        ready_queue = Queue()
        sched = Scheduler(ready_queue)
        now = datetime.now()

        callback_called = [False]

        def callback():
            callback_called[0] = True

        sched.enter(MockItem("foo"), eta=now, callback=callback)

        remaining = iter(sched).next()
        self.assertIsNone(remaining)
        self.assertTrue(callback_called[0])
        self.assertEqual(ready_queue.get_nowait().task_id, "foo")
示例#7
0
    def test_sched_run_later(self):
        ready_queue = Queue()
        sched = Scheduler(ready_queue)
        now = datetime.now()

        callback_called = [False]
        def callback():
            callback_called[0] = True

        eta = now + timedelta(seconds=10)
        sched.enter("foo", eta=eta, callback=callback)

        remaining = iter(sched).next()
        self.assertTrue(remaining > 7)
        self.assertFalse(callback_called[0])
        self.assertRaises(Empty, ready_queue.get_nowait)
    def test_sched_run_later(self):
        ready_queue = Queue()
        sched = Scheduler(ready_queue)
        now = datetime.now()

        callback_called = [False]

        def callback():
            callback_called[0] = True

        eta = now + timedelta(seconds=10)
        sched.enter(MockItem("foo"), eta=eta, callback=callback)

        remaining = iter(sched).next()
        self.assertTrue(remaining > 7 or remaining == sched.max_interval)
        self.assertFalse(callback_called[0])
        self.assertRaises(Empty, ready_queue.get_nowait)
示例#9
0
 def test_sched_clear(self):
     ready_queue = Queue()
     sched = Scheduler(ready_queue)
     sched.enter(MockItem("foo"), eta=datetime.now(), callback=None)
     self.assertFalse(sched.empty())
     sched.clear()
     self.assertTrue(sched.empty())
 def __init__(self):
     self.ready_queue = FastQueue()
     self.ready_queue.put(
         TaskRequest(task_name=mytask.name,
                     task_id=gen_unique_id(),
                     args=(2, 2),
                     kwargs={}))
     self.eta_schedule = Scheduler(self.ready_queue)
     self.event_dispatcher = Dispatcher()
    def test_sched_revoked(self):
        ready_queue = Queue()
        sched = Scheduler(ready_queue)
        now = datetime.now()

        callback_called = [False]

        def callback():
            callback_called[0] = True

        sched.enter(MockItem("foo", revoked=True), eta=now, callback=callback)
        iter(sched).next()
        self.assertFalse(callback_called[0])
        self.assertRaises(Empty, ready_queue.get_nowait)
        self.assertFalse(sched.queue)
        sched.clear()
示例#12
0
 def setUp(self):
     self.ready_queue = Queue()
     self.eta_scheduler = Scheduler(self.ready_queue)
     self.logger = get_logger()
     self.logger.setLevel(0)
示例#13
0
class TestCarrotListener(unittest.TestCase):

    def setUp(self):
        self.ready_queue = Queue()
        self.eta_scheduler = Scheduler(self.ready_queue)
        self.logger = get_logger()
        self.logger.setLevel(0)

    def test_connection(self):
        l = CarrotListener(self.ready_queue, self.eta_scheduler, self.logger,
                           send_events=False)

        c = l.reset_connection()
        self.assertTrue(isinstance(l.amqp_connection, BrokerConnection))

        l.close_connection()
        self.assertTrue(l.amqp_connection is None)
        self.assertTrue(l.task_consumer is None)

        c = l.reset_connection()
        self.assertTrue(isinstance(l.amqp_connection, BrokerConnection))

        l.stop()
        self.assertTrue(l.amqp_connection is None)
        self.assertTrue(l.task_consumer is None)

    def test_receieve_message(self):
        l = CarrotListener(self.ready_queue, self.eta_scheduler, self.logger,
                           send_events=False)
        backend = MockBackend()
        m = create_message(backend, task=foo_task.name,
                           args=[2, 4, 8], kwargs={})

        l.event_dispatcher = MockEventDispatcher()
        l.receive_message(m.decode(), m)

        in_bucket = self.ready_queue.get_nowait()
        self.assertTrue(isinstance(in_bucket, TaskWrapper))
        self.assertEquals(in_bucket.task_name, foo_task.name)
        self.assertEquals(in_bucket.execute(), 2 * 4 * 8)
        self.assertTrue(self.eta_scheduler.empty())

    def test_revoke(self):
        ready_queue = Queue()
        l = CarrotListener(ready_queue, self.eta_scheduler, self.logger,
                           send_events=False)
        backend = MockBackend()
        id = gen_unique_id()
        c = create_message(backend, control={"command": "revoke",
                                             "task_id": id})
        t = create_message(backend, task=foo_task.name, args=[2, 4, 8],
                           kwargs={}, id=id)
        l.event_dispatcher = MockEventDispatcher()
        l.receive_message(c.decode(), c)
        from celery.worker.revoke import revoked
        self.assertTrue(id in revoked)

        l.receive_message(t.decode(), t)
        self.assertTrue(ready_queue.empty())

    def test_receieve_message_not_registered(self):
        l = CarrotListener(self.ready_queue, self.eta_scheduler, self.logger,
                          send_events=False)
        backend = MockBackend()
        m = create_message(backend, task="x.X.31x", args=[2, 4, 8], kwargs={})

        l.event_dispatcher = MockEventDispatcher()
        self.assertFalse(l.receive_message(m.decode(), m))
        self.assertRaises(Empty, self.ready_queue.get_nowait)
        self.assertTrue(self.eta_scheduler.empty())

    def test_receieve_message_eta(self):
        l = CarrotListener(self.ready_queue, self.eta_scheduler, self.logger,
                          send_events=False)
        backend = MockBackend()
        m = create_message(backend, task=foo_task.name,
                           args=[2, 4, 8], kwargs={},
                           eta=(datetime.now() +
                               timedelta(days=1)).isoformat())

        l.reset_connection()
        l.receive_message(m.decode(), m)

        in_hold = self.eta_scheduler.queue[0]
        self.assertEquals(len(in_hold), 4)
        eta, priority, task, on_accept = in_hold
        self.assertTrue(isinstance(task, TaskWrapper))
        self.assertTrue(callable(on_accept))
        self.assertEquals(task.task_name, foo_task.name)
        self.assertEquals(task.execute(), 2 * 4 * 8)
        self.assertRaises(Empty, self.ready_queue.get_nowait)
示例#14
0
class TestCarrotListener(unittest.TestCase):

    def setUp(self):
        self.ready_queue = Queue()
        self.eta_schedule = Scheduler(self.ready_queue)
        self.logger = get_logger()
        self.logger.setLevel(0)

    def test_mainloop(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                           send_events=False)

        class MockConnection(object):

            def drain_events(self):
                return "draining"

        l.connection = PlaceHolder()
        l.connection.connection = MockConnection()

        it = l._mainloop()
        self.assertTrue(it.next(), "draining")

        records = {}
        def create_recorder(key):
            def _recorder(*args, **kwargs):
                records[key] = True
            return _recorder

        l.task_consumer = PlaceHolder()
        l.task_consumer.iterconsume = create_recorder("consume_tasks")
        l.broadcast_consumer = PlaceHolder()
        l.broadcast_consumer.register_callback = create_recorder(
                                                    "broadcast_callback")
        l.broadcast_consumer.iterconsume = create_recorder(
                                             "consume_broadcast")
        l.task_consumer.add_consumer = create_recorder("consumer_add")

        records.clear()
        self.assertEquals(l._detect_wait_method(), l._mainloop)
        self.assertTrue(records.get("broadcast_callback"))
        self.assertTrue(records.get("consume_broadcast"))
        self.assertTrue(records.get("consume_tasks"))

        records.clear()
        l.connection.connection = PlaceHolder()
        self.assertTrue(l._detect_wait_method() is l.task_consumer.iterconsume)
        self.assertTrue(records.get("consumer_add"))

    def test_connection(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                           send_events=False)

        l.reset_connection()
        self.assertTrue(isinstance(l.connection, BrokerConnection))

        l.close_connection()
        self.assertTrue(l.connection is None)
        self.assertTrue(l.task_consumer is None)

        l.reset_connection()
        self.assertTrue(isinstance(l.connection, BrokerConnection))

        l.stop()
        self.assertTrue(l.connection is None)
        self.assertTrue(l.task_consumer is None)

    def test_receive_message_control_command(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                           send_events=False)
        backend = MockBackend()
        m = create_message(backend, control={"command": "shutdown"})
        l.event_dispatcher = MockEventDispatcher()
        l.control_dispatch = MockControlDispatch()
        l.receive_message(m.decode(), m)
        self.assertTrue("shutdown" in l.control_dispatch.commands)

    def test_close_connection(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                           send_events=False)
        l._state = RUN
        l.close_connection()

        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                           send_events=False)
        eventer = l.event_dispatcher = MockEventDispatcher()
        heart = l.heart = MockHeart()
        l._state = RUN
        l.close_connection()
        self.assertTrue(eventer.closed)
        self.assertTrue(heart.closed)

    def test_receive_message_unknown(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                           send_events=False)
        backend = MockBackend()
        m = create_message(backend, unknown={"baz": "!!!"})
        l.event_dispatcher = MockEventDispatcher()
        l.control_dispatch = MockControlDispatch()
        import warnings
        with warnings.catch_warnings(record=True) as log:
                l.receive_message(m.decode(), m)
                self.assertTrue(log)
                self.assertTrue("unknown message" in log[0].message.args[0])

    def test_receieve_message(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                           send_events=False)
        backend = MockBackend()
        m = create_message(backend, task=foo_task.name,
                           args=[2, 4, 8], kwargs={})

        l.event_dispatcher = MockEventDispatcher()
        l.receive_message(m.decode(), m)

        in_bucket = self.ready_queue.get_nowait()
        self.assertTrue(isinstance(in_bucket, TaskWrapper))
        self.assertEquals(in_bucket.task_name, foo_task.name)
        self.assertEquals(in_bucket.execute(), 2 * 4 * 8)
        self.assertTrue(self.eta_schedule.empty())

    def test_receieve_message_eta_isoformat(self):
        l = CarrotListener(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.event_dispatcher = MockEventDispatcher()
        l.receive_message(m.decode(), m)

        items = [entry[2] for entry in self.eta_schedule.queue]
        found = 0
        for item in items:
            if item.task_name == foo_task.name:
                found = True
        self.assertTrue(found)

    def test_revoke(self):
        ready_queue = Queue()
        l = CarrotListener(ready_queue, self.eta_schedule, self.logger,
                           send_events=False)
        backend = MockBackend()
        id = gen_unique_id()
        c = create_message(backend, control={"command": "revoke",
                                             "task_id": id})
        t = create_message(backend, task=foo_task.name, args=[2, 4, 8],
                           kwargs={}, id=id)
        l.event_dispatcher = MockEventDispatcher()
        l.receive_message(c.decode(), c)
        from celery.worker.revoke import revoked
        self.assertTrue(id in revoked)

        l.receive_message(t.decode(), t)
        self.assertTrue(ready_queue.empty())

    def test_receieve_message_not_registered(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                          send_events=False)
        backend = MockBackend()
        m = create_message(backend, task="x.X.31x", args=[2, 4, 8], kwargs={})

        l.event_dispatcher = MockEventDispatcher()
        self.assertFalse(l.receive_message(m.decode(), m))
        self.assertRaises(Empty, self.ready_queue.get_nowait)
        self.assertTrue(self.eta_schedule.empty())

    def test_receieve_message_eta(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                          send_events=False)
        backend = MockBackend()
        m = create_message(backend, task=foo_task.name,
                           args=[2, 4, 8], kwargs={},
                           eta=(datetime.now() +
                               timedelta(days=1)).isoformat())

        l.reset_connection()
        p, conf.BROKER_CONNECTION_RETRY = conf.BROKER_CONNECTION_RETRY, False
        try:
            l.reset_connection()
        finally:
            conf.BROKER_CONNECTION_RETRY = p
        l.receive_message(m.decode(), m)

        in_hold = self.eta_schedule.queue[0]
        self.assertEquals(len(in_hold), 4)
        eta, priority, task, on_accept = in_hold
        self.assertTrue(isinstance(task, TaskWrapper))
        self.assertTrue(callable(on_accept))
        self.assertEquals(task.task_name, foo_task.name)
        self.assertEquals(task.execute(), 2 * 4 * 8)
        self.assertRaises(Empty, self.ready_queue.get_nowait)
示例#15
0
class TestCarrotListener(unittest.TestCase):

    def setUp(self):
        self.ready_queue = Queue()
        self.eta_scheduler = Scheduler(self.ready_queue)
        self.logger = get_logger()
        self.logger.setLevel(0)

    def test_connection(self):
        l = CarrotListener(self.ready_queue, self.eta_scheduler, self.logger)

        c = l.reset_connection()
        self.assertTrue(isinstance(l.amqp_connection, BrokerConnection))

        l.close_connection()
        self.assertTrue(l.amqp_connection is None)
        self.assertTrue(l.task_consumer is None)

        c = l.reset_connection()
        self.assertTrue(isinstance(l.amqp_connection, BrokerConnection))

        l.stop()
        self.assertTrue(l.amqp_connection is None)
        self.assertTrue(l.task_consumer is None)

    def test_receieve_message(self):
        l = CarrotListener(self.ready_queue, self.eta_scheduler, self.logger)
        backend = MockBackend()
        m = create_message(backend, task=foo_task.name,
                           args=[2, 4, 8], kwargs={})

        l.receive_message(m.decode(), m)

        in_bucket = self.ready_queue.get_nowait()
        self.assertTrue(isinstance(in_bucket, TaskWrapper))
        self.assertEquals(in_bucket.task_name, foo_task.name)
        self.assertEquals(in_bucket.execute(), 2 * 4 * 8)
        self.assertTrue(self.eta_scheduler.empty())

    def test_receieve_message_not_registered(self):
        l = CarrotListener(self.ready_queue, self.eta_scheduler, self.logger)
        backend = MockBackend()
        m = create_message(backend, task="x.X.31x", args=[2, 4, 8], kwargs={})

        self.assertFalse(l.receive_message(m.decode(), m))
        self.assertRaises(Empty, self.ready_queue.get_nowait)
        self.assertTrue(self.eta_scheduler.empty())

    def test_receieve_message_eta(self):
        l = CarrotListener(self.ready_queue, self.eta_scheduler, self.logger)
        backend = MockBackend()
        m = create_message(backend, task=foo_task.name,
                           args=[2, 4, 8], kwargs={},
                           eta=(datetime.now() +
                               timedelta(days=1)).isoformat())

        l.receive_message(m.decode(), m)

        in_hold = self.eta_scheduler.queue[0]
        self.assertEquals(len(in_hold), 4)
        eta, priority, task, on_accept = in_hold
        self.assertTrue(isinstance(task, TaskWrapper))
        self.assertTrue(callable(on_accept))
        self.assertEquals(task.task_name, foo_task.name)
        self.assertEquals(task.execute(), 2 * 4 * 8)
        self.assertRaises(Empty, self.ready_queue.get_nowait)
示例#16
0
class test_CarrotListener(unittest.TestCase):

    def setUp(self):
        self.ready_queue = FastQueue()
        self.eta_schedule = Scheduler(self.ready_queue)
        self.logger = get_logger()
        self.logger.setLevel(0)

    def test_mainloop(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                           send_events=False)

        class MockConnection(object):

            def drain_events(self):
                return "draining"

        l.connection = MockConnection()
        l.connection.connection = MockConnection()

        it = l._mainloop()
        self.assertTrue(it.next(), "draining")

        records = {}
        def create_recorder(key):
            def _recorder(*args, **kwargs):
                records[key] = True
            return _recorder

        l.task_consumer = PlaceHolder()
        l.task_consumer.iterconsume = create_recorder("consume_tasks")
        l.broadcast_consumer = PlaceHolder()
        l.broadcast_consumer.register_callback = create_recorder(
                                                    "broadcast_callback")
        l.broadcast_consumer.iterconsume = create_recorder(
                                             "consume_broadcast")
        l.task_consumer.add_consumer = create_recorder("consumer_add")

        records.clear()
        self.assertEqual(l._detect_wait_method(), l._mainloop)
        for record in ("broadcast_callback", "consume_broadcast",
                "consume_tasks"):
            self.assertTrue(records.get(record))

        records.clear()
        l.connection.connection = PlaceHolder()
        self.assertIs(l._detect_wait_method(), l.task_consumer.iterconsume)
        self.assertTrue(records.get("consumer_add"))

    def test_connection(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                           send_events=False)

        l.reset_connection()
        self.assertIsInstance(l.connection, BrokerConnection)

        l.stop_consumers()
        self.assertIsNone(l.connection)
        self.assertIsNone(l.task_consumer)

        l.reset_connection()
        self.assertIsInstance(l.connection, BrokerConnection)

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

    def test_receive_message_control_command(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                           send_events=False)
        backend = MockBackend()
        m = create_message(backend, control={"command": "shutdown"})
        l.event_dispatcher = MockEventDispatcher()
        l.control_dispatch = MockControlDispatch()
        l.receive_message(m.decode(), m)
        self.assertIn("shutdown", l.control_dispatch.commands)

    def test_close_connection(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                           send_events=False)
        l._state = RUN
        l.close_connection()

        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                           send_events=False)
        eventer = l.event_dispatcher = MockEventDispatcher()
        heart = l.heart = MockHeart()
        l._state = RUN
        l.stop_consumers()
        self.assertTrue(eventer.closed)
        self.assertTrue(heart.closed)

    def test_receive_message_unknown(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                           send_events=False)
        backend = MockBackend()
        m = create_message(backend, unknown={"baz": "!!!"})
        l.event_dispatcher = MockEventDispatcher()
        l.control_dispatch = MockControlDispatch()

        def with_catch_warnings(log):
            l.receive_message(m.decode(), m)
            self.assertTrue(log)
            self.assertIn("unknown message", log[0].message.args[0])

        context = catch_warnings(record=True)
        execute_context(context, with_catch_warnings)

    def test_receive_message_InvalidTaskError(self):
        logger = MockLogger()
        l = CarrotListener(self.ready_queue, self.eta_schedule, logger,
                           send_events=False)
        backend = MockBackend()
        m = create_message(backend, task=foo_task.name,
            args=(1, 2), kwargs="foobarbaz", id=1)
        l.event_dispatcher = MockEventDispatcher()
        l.control_dispatch = MockControlDispatch()

        l.receive_message(m.decode(), m)
        self.assertIn("Invalid task ignored", logger.logged[0])

    def test_on_decode_error(self):
        logger = MockLogger()
        l = CarrotListener(self.ready_queue, self.eta_schedule, logger,
                           send_events=False)

        class MockMessage(object):
            content_type = "application/x-msgpack"
            content_encoding = "binary"
            body = "foobarbaz"
            acked = False

            def ack(self):
                self.acked = True

        message = MockMessage()
        l.on_decode_error(message, KeyError("foo"))
        self.assertTrue(message.acked)
        self.assertIn("Message decoding error", logger.logged[0])

    def test_receieve_message(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                           send_events=False)
        backend = MockBackend()
        m = create_message(backend, task=foo_task.name,
                           args=[2, 4, 8], kwargs={})

        l.event_dispatcher = MockEventDispatcher()
        l.receive_message(m.decode(), m)

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

    def test_receieve_message_eta_isoformat(self):

        class MockConsumer(object):
            prefetch_count_incremented = False

            def qos(self, **kwargs):
                self.prefetch_count_incremented = True

        l = CarrotListener(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.event_dispatcher = MockEventDispatcher()
        l.task_consumer = MockConsumer()
        l.qos = QoS(l.task_consumer, l.initial_prefetch_count, l.logger)
        l.receive_message(m.decode(), m)

        items = [entry[2] for entry in self.eta_schedule.queue]
        found = 0
        for item in items:
            if item.task_name == foo_task.name:
                found = True
        self.assertTrue(found)
        self.assertTrue(l.task_consumer.prefetch_count_incremented)

    def test_revoke(self):
        ready_queue = FastQueue()
        l = CarrotListener(ready_queue, self.eta_schedule, self.logger,
                           send_events=False)
        backend = MockBackend()
        id = gen_unique_id()
        c = create_message(backend, control={"command": "revoke",
                                             "task_id": id})
        t = create_message(backend, task=foo_task.name, args=[2, 4, 8],
                           kwargs={}, id=id)
        l.event_dispatcher = MockEventDispatcher()
        l.receive_message(c.decode(), c)
        from celery.worker.revoke import revoked
        self.assertIn(id, revoked)

        l.receive_message(t.decode(), t)
        self.assertTrue(ready_queue.empty())

    def test_receieve_message_not_registered(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                          send_events=False)
        backend = MockBackend()
        m = create_message(backend, task="x.X.31x", args=[2, 4, 8], kwargs={})

        l.event_dispatcher = MockEventDispatcher()
        self.assertFalse(l.receive_message(m.decode(), m))
        self.assertRaises(Empty, self.ready_queue.get_nowait)
        self.assertTrue(self.eta_schedule.empty())

    def test_receieve_message_eta(self):
        l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger,
                          send_events=False)
        backend = MockBackend()
        m = create_message(backend, task=foo_task.name,
                           args=[2, 4, 8], kwargs={},
                           eta=(datetime.now() +
                               timedelta(days=1)).isoformat())

        l.reset_connection()
        p, conf.BROKER_CONNECTION_RETRY = conf.BROKER_CONNECTION_RETRY, False
        try:
            l.reset_connection()
        finally:
            conf.BROKER_CONNECTION_RETRY = p
        l.receive_message(m.decode(), m)

        in_hold = self.eta_schedule.queue[0]
        self.assertEqual(len(in_hold), 4)
        eta, priority, task, on_accept = in_hold
        self.assertIsInstance(task, TaskRequest)
        self.assertTrue(callable(on_accept))
        self.assertEqual(task.task_name, foo_task.name)
        self.assertEqual(task.execute(), 2 * 4 * 8)
        self.assertRaises(Empty, self.ready_queue.get_nowait)

    def test_start__consume_messages(self):

        class _QoS(object):
            prev = 3
            next = 4

            def update(self):
                self.prev = self.next

        class _Listener(CarrotListener):
            iterations = 0
            wait_method = None

            def reset_connection(self):
                if self.iterations >= 1:
                    raise KeyError("foo")

            def _detect_wait_method(self):
                return self.wait_method

        called_back = [False]
        def init_callback(listener):
            called_back[0] = True


        l = _Listener(self.ready_queue, self.eta_schedule, self.logger,
                      send_events=False, init_callback=init_callback)
        l.qos = _QoS()

        def raises_KeyError(limit=None):
            yield True
            l.iterations = 1
            raise KeyError("foo")

        l.wait_method = raises_KeyError
        self.assertRaises(KeyError, l.start)
        self.assertTrue(called_back[0])
        self.assertEqual(l.iterations, 1)
        self.assertEqual(l.qos.prev, l.qos.next)

        l = _Listener(self.ready_queue, self.eta_schedule, self.logger,
                      send_events=False, init_callback=init_callback)
        l.qos = _QoS()
        def raises_socket_error(limit=None):
            yield True
            l.iterations = 1
            raise socket.error("foo")

        l.wait_method = raises_socket_error
        self.assertRaises(KeyError, l.start)
        self.assertTrue(called_back[0])
        self.assertEqual(l.iterations, 1)
 def test_empty_queue_yields_None(self):
     ready_queue = Queue()
     sched = Scheduler(ready_queue)
     self.assertIsNone(iter(sched).next())
    def __init__(self,
                 concurrency=None,
                 logfile=None,
                 loglevel=None,
                 send_events=conf.SEND_EVENTS,
                 hostname=None,
                 ready_callback=noop,
                 embed_clockservice=False,
                 pool_cls=conf.CELERYD_POOL,
                 listener_cls=conf.CELERYD_LISTENER,
                 mediator_cls=conf.CELERYD_MEDIATOR,
                 eta_scheduler_cls=conf.CELERYD_ETA_SCHEDULER,
                 schedule_filename=conf.CELERYBEAT_SCHEDULE_FILENAME,
                 task_time_limit=conf.CELERYD_TASK_TIME_LIMIT,
                 task_soft_time_limit=conf.CELERYD_TASK_SOFT_TIME_LIMIT,
                 max_tasks_per_child=conf.CELERYD_MAX_TASKS_PER_CHILD,
                 pool_putlocks=conf.CELERYD_POOL_PUTLOCKS,
                 db=conf.CELERYD_STATE_DB):

        # Options
        self.loglevel = loglevel or self.loglevel
        self.concurrency = concurrency or self.concurrency
        self.logfile = logfile or self.logfile
        self.logger = setup_logger(loglevel, logfile)
        self.hostname = hostname or socket.gethostname()
        self.embed_clockservice = embed_clockservice
        self.ready_callback = ready_callback
        self.send_events = send_events
        self.task_time_limit = task_time_limit
        self.task_soft_time_limit = task_soft_time_limit
        self.max_tasks_per_child = max_tasks_per_child
        self.pool_putlocks = pool_putlocks
        self.db = db
        self._finalize = Finalize(self, self.stop, exitpriority=1)

        if self.db:
            persistence = state.Persistent(self.db)
            Finalize(persistence, persistence.save, exitpriority=5)

        # Queues
        if conf.DISABLE_RATE_LIMITS:
            self.ready_queue = FastQueue()
        else:
            self.ready_queue = TaskBucket(task_registry=registry.tasks)
        self.eta_schedule = Scheduler(self.ready_queue, logger=self.logger)

        self.logger.debug("Instantiating thread components...")

        # Threads + Pool + Consumer
        self.pool = instantiate(pool_cls,
                                self.concurrency,
                                logger=self.logger,
                                initializer=process_initializer,
                                maxtasksperchild=self.max_tasks_per_child,
                                timeout=self.task_time_limit,
                                soft_timeout=self.task_soft_time_limit,
                                putlocks=self.pool_putlocks)
        self.mediator = instantiate(mediator_cls,
                                    self.ready_queue,
                                    callback=self.process_task,
                                    logger=self.logger)
        self.scheduler = instantiate(eta_scheduler_cls,
                                     self.eta_schedule,
                                     logger=self.logger)

        self.clockservice = None
        if self.embed_clockservice:
            self.clockservice = EmbeddedClockService(
                logger=self.logger, schedule_filename=schedule_filename)

        prefetch_count = self.concurrency * conf.CELERYD_PREFETCH_MULTIPLIER
        self.listener = instantiate(listener_cls,
                                    self.ready_queue,
                                    self.eta_schedule,
                                    logger=self.logger,
                                    hostname=self.hostname,
                                    send_events=self.send_events,
                                    init_callback=self.ready_callback,
                                    initial_prefetch_count=prefetch_count,
                                    pool=self.pool)

        # The order is important here;
        #   the first in the list is the first to start,
        # and they must be stopped in reverse order.
        self.components = filter(None,
                                 (self.pool, self.mediator, self.scheduler,
                                  self.clockservice, self.listener))
 def setUp(self):
     self.ready_queue = FastQueue()
     self.eta_schedule = Scheduler(self.ready_queue)
     self.logger = get_logger()
     self.logger.setLevel(0)
class test_CarrotListener(unittest.TestCase):
    def setUp(self):
        self.ready_queue = FastQueue()
        self.eta_schedule = Scheduler(self.ready_queue)
        self.logger = get_logger()
        self.logger.setLevel(0)

    def test_mainloop(self):
        l = CarrotListener(self.ready_queue,
                           self.eta_schedule,
                           self.logger,
                           send_events=False)

        class MockConnection(object):
            def drain_events(self):
                return "draining"

        l.connection = MockConnection()
        l.connection.connection = MockConnection()

        it = l._mainloop()
        self.assertTrue(it.next(), "draining")

        records = {}

        def create_recorder(key):
            def _recorder(*args, **kwargs):
                records[key] = True

            return _recorder

        l.task_consumer = PlaceHolder()
        l.task_consumer.iterconsume = create_recorder("consume_tasks")
        l.broadcast_consumer = PlaceHolder()
        l.broadcast_consumer.register_callback = create_recorder(
            "broadcast_callback")
        l.broadcast_consumer.iterconsume = create_recorder("consume_broadcast")
        l.task_consumer.add_consumer = create_recorder("consumer_add")

        records.clear()
        self.assertEqual(l._detect_wait_method(), l._mainloop)
        for record in ("broadcast_callback", "consume_broadcast",
                       "consume_tasks"):
            self.assertTrue(records.get(record))

        records.clear()
        l.connection.connection = PlaceHolder()
        self.assertIs(l._detect_wait_method(), l.task_consumer.iterconsume)
        self.assertTrue(records.get("consumer_add"))

    def test_connection(self):
        l = CarrotListener(self.ready_queue,
                           self.eta_schedule,
                           self.logger,
                           send_events=False)

        l.reset_connection()
        self.assertIsInstance(l.connection, BrokerConnection)

        l.stop_consumers()
        self.assertIsNone(l.connection)
        self.assertIsNone(l.task_consumer)

        l.reset_connection()
        self.assertIsInstance(l.connection, BrokerConnection)

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

    def test_receive_message_control_command(self):
        l = CarrotListener(self.ready_queue,
                           self.eta_schedule,
                           self.logger,
                           send_events=False)
        backend = MockBackend()
        m = create_message(backend, control={"command": "shutdown"})
        l.event_dispatcher = MockEventDispatcher()
        l.control_dispatch = MockControlDispatch()
        l.receive_message(m.decode(), m)
        self.assertIn("shutdown", l.control_dispatch.commands)

    def test_close_connection(self):
        l = CarrotListener(self.ready_queue,
                           self.eta_schedule,
                           self.logger,
                           send_events=False)
        l._state = RUN
        l.close_connection()

        l = CarrotListener(self.ready_queue,
                           self.eta_schedule,
                           self.logger,
                           send_events=False)
        eventer = l.event_dispatcher = MockEventDispatcher()
        heart = l.heart = MockHeart()
        l._state = RUN
        l.stop_consumers()
        self.assertTrue(eventer.closed)
        self.assertTrue(heart.closed)

    def test_receive_message_unknown(self):
        l = CarrotListener(self.ready_queue,
                           self.eta_schedule,
                           self.logger,
                           send_events=False)
        backend = MockBackend()
        m = create_message(backend, unknown={"baz": "!!!"})
        l.event_dispatcher = MockEventDispatcher()
        l.control_dispatch = MockControlDispatch()

        def with_catch_warnings(log):
            l.receive_message(m.decode(), m)
            self.assertTrue(log)
            self.assertIn("unknown message", log[0].message.args[0])

        context = catch_warnings(record=True)
        execute_context(context, with_catch_warnings)

    def test_receive_message_InvalidTaskError(self):
        logger = MockLogger()
        l = CarrotListener(self.ready_queue,
                           self.eta_schedule,
                           logger,
                           send_events=False)
        backend = MockBackend()
        m = create_message(backend,
                           task=foo_task.name,
                           args=(1, 2),
                           kwargs="foobarbaz",
                           id=1)
        l.event_dispatcher = MockEventDispatcher()
        l.control_dispatch = MockControlDispatch()

        l.receive_message(m.decode(), m)
        self.assertIn("Invalid task ignored", logger.logged[0])

    def test_on_decode_error(self):
        logger = MockLogger()
        l = CarrotListener(self.ready_queue,
                           self.eta_schedule,
                           logger,
                           send_events=False)

        class MockMessage(object):
            content_type = "application/x-msgpack"
            content_encoding = "binary"
            body = "foobarbaz"
            acked = False

            def ack(self):
                self.acked = True

        message = MockMessage()
        l.on_decode_error(message, KeyError("foo"))
        self.assertTrue(message.acked)
        self.assertIn("Message decoding error", logger.logged[0])

    def test_receieve_message(self):
        l = CarrotListener(self.ready_queue,
                           self.eta_schedule,
                           self.logger,
                           send_events=False)
        backend = MockBackend()
        m = create_message(backend,
                           task=foo_task.name,
                           args=[2, 4, 8],
                           kwargs={})

        l.event_dispatcher = MockEventDispatcher()
        l.receive_message(m.decode(), m)

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

    def test_receieve_message_eta_isoformat(self):
        class MockConsumer(object):
            prefetch_count_incremented = False

            def qos(self, **kwargs):
                self.prefetch_count_incremented = True

        l = CarrotListener(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.event_dispatcher = MockEventDispatcher()
        l.task_consumer = MockConsumer()
        l.qos = QoS(l.task_consumer, l.initial_prefetch_count, l.logger)
        l.receive_message(m.decode(), m)

        items = [entry[2] for entry in self.eta_schedule.queue]
        found = 0
        for item in items:
            if item.task_name == foo_task.name:
                found = True
        self.assertTrue(found)
        self.assertTrue(l.task_consumer.prefetch_count_incremented)

    def test_revoke(self):
        ready_queue = FastQueue()
        l = CarrotListener(ready_queue,
                           self.eta_schedule,
                           self.logger,
                           send_events=False)
        backend = MockBackend()
        id = gen_unique_id()
        c = create_message(backend,
                           control={
                               "command": "revoke",
                               "task_id": id
                           })
        t = create_message(backend,
                           task=foo_task.name,
                           args=[2, 4, 8],
                           kwargs={},
                           id=id)
        l.event_dispatcher = MockEventDispatcher()
        l.receive_message(c.decode(), c)
        from celery.worker.state import revoked
        self.assertIn(id, revoked)

        l.receive_message(t.decode(), t)
        self.assertTrue(ready_queue.empty())

    def test_receieve_message_not_registered(self):
        l = CarrotListener(self.ready_queue,
                           self.eta_schedule,
                           self.logger,
                           send_events=False)
        backend = MockBackend()
        m = create_message(backend, task="x.X.31x", args=[2, 4, 8], kwargs={})

        l.event_dispatcher = MockEventDispatcher()
        self.assertFalse(l.receive_message(m.decode(), m))
        self.assertRaises(Empty, self.ready_queue.get_nowait)
        self.assertTrue(self.eta_schedule.empty())

    def test_receieve_message_eta(self):
        l = CarrotListener(self.ready_queue,
                           self.eta_schedule,
                           self.logger,
                           send_events=False)
        backend = MockBackend()
        m = create_message(backend,
                           task=foo_task.name,
                           args=[2, 4, 8],
                           kwargs={},
                           eta=(datetime.now() +
                                timedelta(days=1)).isoformat())

        l.reset_connection()
        p, conf.BROKER_CONNECTION_RETRY = conf.BROKER_CONNECTION_RETRY, False
        try:
            l.reset_connection()
        finally:
            conf.BROKER_CONNECTION_RETRY = p
        l.receive_message(m.decode(), m)

        in_hold = self.eta_schedule.queue[0]
        self.assertEqual(len(in_hold), 4)
        eta, priority, task, on_accept = in_hold
        self.assertIsInstance(task, TaskRequest)
        self.assertTrue(callable(on_accept))
        self.assertEqual(task.task_name, foo_task.name)
        self.assertEqual(task.execute(), 2 * 4 * 8)
        self.assertRaises(Empty, self.ready_queue.get_nowait)

    def test_start__consume_messages(self):
        class _QoS(object):
            prev = 3
            next = 4

            def update(self):
                self.prev = self.next

        class _Listener(CarrotListener):
            iterations = 0
            wait_method = None

            def reset_connection(self):
                if self.iterations >= 1:
                    raise KeyError("foo")

            def _detect_wait_method(self):
                return self.wait_method

        called_back = [False]

        def init_callback(listener):
            called_back[0] = True

        l = _Listener(self.ready_queue,
                      self.eta_schedule,
                      self.logger,
                      send_events=False,
                      init_callback=init_callback)
        l.qos = _QoS()

        def raises_KeyError(limit=None):
            yield True
            l.iterations = 1
            raise KeyError("foo")

        l.wait_method = raises_KeyError
        self.assertRaises(KeyError, l.start)
        self.assertTrue(called_back[0])
        self.assertEqual(l.iterations, 1)
        self.assertEqual(l.qos.prev, l.qos.next)

        l = _Listener(self.ready_queue,
                      self.eta_schedule,
                      self.logger,
                      send_events=False,
                      init_callback=init_callback)
        l.qos = _QoS()

        def raises_socket_error(limit=None):
            yield True
            l.iterations = 1
            raise socket.error("foo")

        l.wait_method = raises_socket_error
        self.assertRaises(KeyError, l.start)
        self.assertTrue(called_back[0])
        self.assertEqual(l.iterations, 1)