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())
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")
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")
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)
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()
def setUp(self):
self.ready_queue = Queue()
self.eta_scheduler = Scheduler(self.ready_queue)
self.logger = get_logger()
self.logger.setLevel(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)
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)
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)
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)
