def test_babysit_creates_correct_records(self, add_consumer, active_queues):
"""
Test babysit() with a blank database. It should create the correct AvailableQueues.
"""
tasks.babysit()
# babysit() should have called the active_queues() method
active_queues.assert_called_once_with()
# There should be three ActiveQueues, one for each reserved worker in the mock data
aqc = AvailableQueue.get_collection()
self.assertEqual(aqc.count(), 3)
# Let's make sure their names and num_reservations counts are correct
self.assertEqual(aqc.find_one({'_id': RESERVED_WORKER_1})['num_reservations'], 0)
self.assertEqual(aqc.find_one({'_id': RESERVED_WORKER_2})['num_reservations'], 0)
self.assertEqual(aqc.find_one({'_id': RESERVED_WORKER_3})['num_reservations'], 0)
# Reserved worker 3 wasn't assigned to a queue, so babysit() should have assigned it to one
add_consumer.assert_called_once_with(queue=RESERVED_WORKER_3,
destination=(RESERVED_WORKER_3,))
def test__delete_queue(self, logger, cancel, active_queues, mock_add_consumer):
"""
Assert that the correct Tasks get canceled when their queue is deleted, and that the queue
is removed from the database.
"""
# Let's start off by creating the existing queues and what not by calling babysit
tasks.babysit()
# Let's simulate three tasks being assigned to RESERVED_WORKER_2, with two of them being
# in an incomplete state and one in a complete state. We will delete RESERVED_WORKER_2's
# queue, which should cause the two to get canceled. Let's put task_1 in progress
task_1 = TaskStatusManager.create_task_status('task_1', RESERVED_WORKER_2,
state=CALL_RUNNING_STATE)
task_2 = TaskStatusManager.create_task_status('task_2', RESERVED_WORKER_2,
state=CALL_WAITING_STATE)
# This task shouldn't get canceled because it isn't in an incomplete state
task_3 = TaskStatusManager.create_task_status('task_3', RESERVED_WORKER_2,
state=CALL_FINISHED_STATE)
# Let's make a task in a worker that is still present just to make sure it isn't touched.
task_4 = TaskStatusManager.create_task_status('task_4', RESERVED_WORKER_1,
state=CALL_RUNNING_STATE)
# Let's just make sure the babysit() worked and that we have an AvailableQueue with RR2
aqc = AvailableQueue.get_collection()
self.assertEqual(aqc.find({'_id': RESERVED_WORKER_2}).count(), 1)
# Now let's delete the queue named RESERVED_WORKER_2
tasks._delete_queue.apply_async(args=(RESERVED_WORKER_2,),
queue=tasks.RESOURCE_MANAGER_QUEUE)
# cancel() should have been called twice with task_1 and task_2 as parameters
self.assertEqual(cancel.call_count, 2)
# Let's build a set out of the two times that cancel was called. We can't know for sure
# which order the Tasks got canceled in, but we can assert that the correct two tasks were
# canceled (task_3 should not appear in this set).
cancel_param_set = set([c[1] for c in cancel.mock_calls])
self.assertEqual(cancel_param_set, set([('task_1',), ('task_2',)]))
# We should have logged that we are canceling the tasks
self.assertEqual(logger.call_count, 0)
self.assertTrue(RESERVED_WORKER_2 in logger.mock_calls[0][1][0])
self.assertTrue('Canceling the tasks' in logger.mock_calls[0][1][0])
# The queue should have been deleted
self.assertEqual(aqc.find({'_id': RESERVED_WORKER_2}).count(), 0)
# The other queues (1 and 3) should remain
self.assertEqual(aqc.find().count(), 2)
def test_babysit_cancels_correct_tasks(self, logger, add_consumer, cancel, active_queues):
"""
When babysit() discovers that a worker has gone missing, it should cancel all of the tasks
that were in its queue.
"""
# Let's start off by creating the existing queues and what not by calling babysit
tasks.babysit()
# Now, let's add another AvailableQueue that isn't found in our active_queues mock so that
# babysit() can notice that it appears to have gone missing next time it's called. We need
# to also mark it as having been missing for at least 5 minutes.
missing_available_queue = AvailableQueue('%s4' % tasks.RESERVED_WORKER_NAME_PREFIX, 2,
datetime.utcnow() - timedelta(minutes=5))
missing_available_queue.save()
# Let's simulate three tasks being assigned to this AvailableQueue, with two of them being
# in an incomplete state and one in a complete state. The two should get canceled.
# Let's put task_1 in progress
task_1 = TaskStatusManager.create_task_status('task_1', missing_available_queue.name,
state=CALL_RUNNING_STATE)
task_2 = TaskStatusManager.create_task_status('task_2', missing_available_queue.name,
state=CALL_WAITING_STATE)
# This task shouldn't get canceled because it isn't in an incomplete state
task_3 = TaskStatusManager.create_task_status('task_3', missing_available_queue.name,
state=CALL_FINISHED_STATE)
# Let's make a task in a worker that is still present just to make sure it isn't touched.
task_4 = TaskStatusManager.create_task_status('task_4', RESERVED_WORKER_1,
state=CALL_RUNNING_STATE)
# Now, let's call babysit() again. This time, it should delete the AvailableQueue, and it
# should cancel task_1 and task_2. task_3 should be left alone.
tasks.babysit()
# cancel() should have been called twice with task_1 and task_2 as parameters
self.assertEqual(cancel.call_count, 2)
# Let's build a set out of the two times that cancel was called. We can't know for sure
# which order the Tasks got canceled in, but we can assert that the correct two tasks were
# canceled (task_3 should not appear in this set).
cancel_param_set = set([c[1] for c in cancel.mock_calls])
self.assertEqual(cancel_param_set, set([('task_1',), ('task_2',)]))
# We should have logged that we are canceling the tasks
self.assertEqual(logger.call_count, 0)
self.assertTrue(missing_available_queue.name in logger.mock_calls[0][1][0])
self.assertTrue('Canceling the tasks' in logger.mock_calls[0][1][0])
def test_babysit_resets_missing_since_on_reappearing_workers(self, add_consumer, active_queues):
"""
Let's simulate an AvailableQueue having been missing in the past by setting its
missing_since attribute to two minutes ago. It is part of the mocked active_queues() call,
so we expect babysit() to set its missing_since attribute back to None. Note that this one
has been missing for more than five minutes, but it got lucky because it is back just in
time to avoid being deleted.
"""
available_queue_2 = AvailableQueue(name=RESERVED_WORKER_2,
missing_since=datetime.utcnow() - timedelta(minutes=6))
available_queue_2.save()
tasks.babysit()
# babysit() should have called the active_queues() method
active_queues.assert_called_once_with()
# There should be three ActiveQueues, one for each reserved worker in the mock data
aqc = AvailableQueue.get_collection()
self.assertEqual(aqc.count(), 3)
# Make sure it's set back to None
aq_2 = aqc.find_one({'_id': RESERVED_WORKER_2})
self.assertEqual(aq_2['num_reservations'], 0)
self.assertEqual(aq_2['missing_since'], None)
def test_babysit_deletes_correct_records(self, add_consumer, _delete_queue_apply_async,
active_queues):
"""
Test babysit() with pre-existing state. It should create the correct AvailableQueues, and
delete other ones, and leave others in place.
"""
# This AvailableQueue should remain in the DB
available_queue_2 = AvailableQueue(name=RESERVED_WORKER_2)
available_queue_2.save()
# This AvailableQueue doesn't exist anymore since it's not in the mock results, and it's
# been missing for five minutes, so it should get deleted
available_queue_4 = AvailableQueue(name='%s4' % tasks.RESERVED_WORKER_NAME_PREFIX,
missing_since=datetime.utcnow() - timedelta(minutes=5))
available_queue_4.save()
# This AvailableQueue doesn't exist anymore since it's not in the mock results, but it's
# been missing for less than five minutes, so it should not get deleted
available_queue_5 = AvailableQueue(name='%s5' % tasks.RESERVED_WORKER_NAME_PREFIX,
missing_since=datetime.utcnow() - timedelta(minutes=2))
available_queue_5.save()
# This AvailableQueue doesn't exist anymore since it's not in the mock results, but it
# hasn't been missing before (i.e., it's missing_since attribute is None), so it should not
# get deleted. It's missing_since attribute should be set to a datetime, however.
available_queue_6 = AvailableQueue(name='%s6' % tasks.RESERVED_WORKER_NAME_PREFIX,
missing_since=None)
available_queue_6.save()
# This should cause queue 4 to get deleted, and 6 to get marked as missing.
tasks.babysit()
# babysit() should have called the active_queues() method
active_queues.assert_called_once_with()
# There should be five ActiveQueues, one for each reserved worker in the mock data (3), and
# numbers 5 and 6 that we created above should also remain because they have been missing
# for less than five minutes.
aqc = AvailableQueue.get_collection()
self.assertEqual(aqc.count(), 5)
# Let's make sure their names, num_reservations counts, and missing_since attributes are
# correct
aq_1 = aqc.find_one({'_id': RESERVED_WORKER_1})
self.assertEqual(aq_1['num_reservations'], 0)
self.assertEqual(aq_1['missing_since'], None)
aq_2 = aqc.find_one({'_id': RESERVED_WORKER_2})
self.assertEqual(aq_2['num_reservations'], 0)
self.assertEqual(aq_2['missing_since'], None)
aq_3 = aqc.find_one({'_id': RESERVED_WORKER_3})
self.assertEqual(aq_3['num_reservations'], 0)
self.assertEqual(aq_3['missing_since'], None)
# Numbers 5 and 6 should exist, with non-null missing_since attributes
aq_5 = aqc.find_one({'_id': '%s5' % tasks.RESERVED_WORKER_NAME_PREFIX})
self.assertEqual(aq_5['num_reservations'], 0)
self.assertEqual(type(aq_5['missing_since']), datetime)
self.assertTrue(aq_5['missing_since'] < datetime.utcnow() - timedelta(minutes=2))
aq_6 = aqc.find_one({'_id': '%s6' % tasks.RESERVED_WORKER_NAME_PREFIX})
self.assertEqual(aq_6['num_reservations'], 0)
self.assertEqual(type(aq_6['missing_since']), datetime)
self.assertTrue(aq_6['missing_since'] < datetime.utcnow())
# Reserved worker 3 wasn't assigned to a queue, so babysit() should have assigned it to one
add_consumer.assert_called_once_with(queue=RESERVED_WORKER_3,
destination=(RESERVED_WORKER_3,))
# Make sure that _delete_queue was called for #4, and that the delete task was sent to the
# RESOURCE_MANAGER_QUEUE
_delete_queue_apply_async.assert_called_once_with(
args=('%s4' % tasks.RESERVED_WORKER_NAME_PREFIX,), queue=tasks.RESOURCE_MANAGER_QUEUE)
