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_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_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_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_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_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_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_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_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_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_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_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)