def _test_GET_celery_tasks(self):
"""
Test the GET() method to get all current tasks.
"""
# Populate a couple of task statuses
task_id1 = str(uuid.uuid4())
worker_1 = 'worker_1'
state1 = 'waiting'
task_id2 = str(uuid.uuid4())
worker_2 = 'worker_2'
state2 = 'running'
tags = ['random', 'tags']
TaskStatus(task_id1, worker_1, tags, state1).save()
TaskStatus(task_id2, worker_2, tags, state2).save()
status, body = self.get('/v2/tasks/')
# Validate
self.assertEqual(200, status)
self.assertTrue(len(body) == 2)
for task in body:
if task['task_id'] == task_id1:
self.assertEqual(task['_href'],
serialization.dispatch.task_result_href(task)['_href'])
self.assertEquals(task['state'], state1)
self.assertEqual(task['worker_name'], worker_1)
else:
self.assertEqual(task['_href'],
serialization.dispatch.task_result_href(task)['_href'])
self.assertEquals(task['state'], state2)
self.assertEqual(task['worker_name'], worker_2)
self.assertEquals(task['tags'], tags)
def test_find_by_criteria_with_result(self):
tags = ['test', 'tags']
TaskStatus(task_id='1').save()
TaskStatus(task_id='2', tags=tags).save()
result = 'done'
TaskStatus(task_id='3',
tags=tags,
state=constants.CALL_FINISHED_STATE,
result=result).save()
filters = {'tags': tags, 'task_id': {'$in': ['1', '3']}}
fields = ['task_id', 'tags', 'result']
limit = 1
sort = (('task_id', DESCENDING), )
criteria = Criteria(filters=filters,
fields=fields,
limit=limit,
sort=sort)
query_set = TaskStatus.objects.find_by_criteria(criteria)
self.assertEqual(len(query_set), 1)
self.assertEqual(query_set[0].task_id, '3')
self.assertEqual(query_set[0].result, result)
task_state_default = constants.CALL_WAITING_STATE
self.assertEqual(query_set[0].state, task_state_default)
def test_task_id_validation(self):
# Valid task_id
valid_task_id = str(uuid4())
TaskStatus(valid_task_id).save()
# Invalid task_id
invalid_task_ids = [4, {}, None, uuid4(), ('a', 'b'), object(), []]
for invalid_task_id in invalid_task_ids:
self.assertRaises(ValidationError,
TaskStatus(invalid_task_id).save)
def test_create_task_status_duplicate_task_id(self):
"""
Tests TaskStatus creation with a duplicate task id.
"""
task_id = self.get_random_uuid()
TaskStatus(task_id).save()
try:
TaskStatus(task_id).save()
except Exception, e:
self.assertTrue(task_id in e.message)
def test_task_type_validation(self):
# Valid task_type
task_id = str(uuid4())
valid_task_type = 'task_type'
TaskStatus(task_id=task_id, task_type=valid_task_type).save()
# Invalid task_type
invalid_task_types = [4, {}, uuid4(), ('a', 'b'), object(), []]
for invalid_task_type in invalid_task_types:
self.assertRaises(
ValidationError,
TaskStatus(task_id=task_id, task_type=invalid_task_type).save)
def test_error_validation(self):
# Valid error
task_id = str(uuid4())
valid_error = {'error': 'some error'}
TaskStatus(task_id=task_id, error=valid_error).save()
# Invalid error
invalid_errors = [4, uuid4(), object(), 'tags', [1, 2]]
for invalid_error in invalid_errors:
self.assertRaises(
ValidationError,
TaskStatus(task_id=task_id, error=invalid_error).save)
def test_state_validation(self):
# Valid state
valid_states = constants.CALL_STATES
for valid_state in valid_states:
TaskStatus(task_id=str(uuid4()), state=valid_state).save()
# Invalid state
invalid_states = [4, {}, uuid4(), object(), 'invalid_state', []]
for invalid_state in invalid_states:
self.assertRaises(
ValidationError,
TaskStatus(task_id=str(uuid4()), state=invalid_state).save)
def test_tags_validation(self):
# Valid tags
task_id = str(uuid4())
valid_tags = ['tag1', 'tag2']
TaskStatus(task_id=task_id, tags=valid_tags).save()
# Invalid tags
invalid_tags = [4, {}, uuid4(), object(), 'tags', [1, 2]]
for invalid_tag in invalid_tags:
self.assertRaises(
ValidationError,
TaskStatus(task_id=task_id, tags=invalid_tag).save)
def test_spawned_tasks_validation(self):
# Valid spawned_tasks
task_id = str(uuid4())
valid_spawned_tasks = ['spawned1', 'spawned2']
TaskStatus(task_id=task_id, spawned_tasks=valid_spawned_tasks).save()
# Invalid spawned_tasks
invalid_spawned_tasks = [4, uuid4(), object(), 'tags', [1, 2], {}]
for invalid_spawned_task in invalid_spawned_tasks:
self.assertRaises(
ValidationError,
TaskStatus(task_id=task_id,
spawned_tasks=invalid_spawned_task).save)
def test_worker_name_validation(self):
# Valid worker_name
task_id = str(uuid4())
valid_worker_name = 'worker_name'
TaskStatus(task_id=task_id, worker_name=valid_worker_name).save()
# Invalid worker_name
invalid_worker_names = [4, {}, uuid4(), ('a', 'b'), object(), []]
for invalid_worker_name in invalid_worker_names:
self.assertRaises(
ValidationError,
TaskStatus(task_id=task_id,
worker_name=invalid_worker_name).save)
def test_progress_report_validation(self):
# Valid progress_report
task_id = str(uuid4())
valid_progress_report = {'progress': 'going good'}
TaskStatus(task_id=task_id,
progress_report=valid_progress_report).save()
# Invalid progress_report
invalid_progress_reports = [4, uuid4(), object(), 'tags', [1, 2], ()]
for invalid_progress_report in invalid_progress_reports:
self.assertRaises(
ValidationError,
TaskStatus(task_id=task_id,
progress_report=invalid_progress_report).save)
def test_start_time_validation(self):
# Valid start_time
task_id = str(uuid4())
valid_start_time = dateutils.format_iso8601_datetime(datetime.now())
TaskStatus(task_id=task_id, start_time=valid_start_time).save()
# Invalid start_time
invalid_start_times = [
4, {}, uuid4(), ('a', 'b'),
object(), [],
datetime.now()
]
for invalid_start_time in invalid_start_times:
self.assertRaises(
ValidationError,
TaskStatus(task_id=task_id, task_type=invalid_start_time).save)
def test_task_status_update(self):
"""
Tests the successful operation of task status update.
"""
task_id = self.get_random_uuid()
worker_name = 'special_worker_name'
tags = ['test-tag1', 'test-tag2']
state = 'waiting'
TaskStatus(task_id, worker_name, tags, state).save()
now = datetime.now(dateutils.utc_tz())
start_time = dateutils.format_iso8601_datetime(now)
delta = {
'start_time': start_time,
'state': 'running',
'progress_report': {
'report-id': 'my-progress'
}
}
TaskStatus.objects(task_id=task_id).update_one(
set__start_time=delta['start_time'],
set__state=delta['state'],
set__progress_report=delta['progress_report'])
task_status = TaskStatus.objects(task_id=task_id).first()
self.assertEqual(task_status['start_time'], delta['start_time'])
# Make sure that parse_iso8601_datetime is able to parse the start_time without errors
dateutils.parse_iso8601_datetime(task_status['start_time'])
self.assertEqual(task_status['state'], delta['state'])
self.assertEqual(task_status['progress_report'],
delta['progress_report'])
self.assertEqual(task_status['worker_name'], worker_name)
def test_save_update_defaults(self):
"""
Test the save method with default arguments when the object is already in the database.
"""
task_id = str(uuid4())
worker_name = 'worker_name'
tags = ['tag_1', 'tag_2']
state = constants.CALL_ACCEPTED_STATE
spawned_tasks = ['foo']
error = {'error': 'some_error'}
progress_report = {
'what do we want?': 'progress!',
'when do we want it?': 'now!'
}
task_type = 'some.task'
start_time = datetime.now()
finish_time = start_time + timedelta(minutes=5)
start_time = dateutils.format_iso8601_datetime(start_time)
finish_time = dateutils.format_iso8601_datetime(finish_time)
result = None
ts = TaskStatus(task_id,
worker_name,
tags,
state,
spawned_tasks=spawned_tasks,
error=error,
progress_report=progress_report,
task_type=task_type,
start_time=start_time,
finish_time=finish_time,
result=result)
# Let's go ahead and insert the object
ts.save()
# Now let's alter it a bit, and make sure the alteration makes it to the DB correctly.
new_state = constants.CALL_RUNNING_STATE
ts.state = new_state
# This should update ts in the database
ts.save()
ts = TaskStatus.objects()
# There should only be one TaskStatus in the db
self.assertEqual(len(ts), 1)
ts = ts[0]
# Make sure all the attributes are correct
self.assertEqual(ts['task_id'], task_id)
self.assertEqual(ts['worker_name'], worker_name)
self.assertEqual(ts['tags'], tags)
# The state should have been updated
self.assertEqual(ts['state'], new_state)
self.assertEqual(ts['error'], error)
self.assertEqual(ts['spawned_tasks'], spawned_tasks)
self.assertEqual(ts['progress_report'], progress_report)
self.assertEqual(ts['task_type'], task_type)
self.assertEqual(ts['start_time'], start_time)
self.assertEqual(ts['finish_time'], finish_time)
self.assertEqual(ts['result'], result)
# These are always None
self.assertEqual(ts['traceback'], None)
self.assertEqual(ts['exception'], None)
def test___init__(self):
"""
Test the __init__() method.
"""
task_id = 'a_task_id'
queue = 'some_queue'
tags = ['tag_1', 'tag_2']
state = 'a state'
spawned_tasks = ['foo']
error = 'some_error'
progress_report = {'what do we want?': 'progress!', 'when do we want it?': 'now!'}
task_type = 'some.task'
start_time = datetime.now()
finish_time = start_time + timedelta(minutes=5)
result = None
ts = TaskStatus(
task_id, queue, tags, state, spawned_tasks=spawned_tasks, error=error,
progress_report=progress_report, task_type=task_type, start_time=start_time,
finish_time=finish_time, result=result)
self.assertEqual(ts.task_id, task_id)
self.assertEqual(ts.queue, queue)
self.assertEqual(ts.tags, tags)
self.assertEqual(ts.state, state)
self.assertEqual(ts.error, error)
self.assertEqual(ts.spawned_tasks, spawned_tasks)
self.assertEqual(ts.progress_report, progress_report)
self.assertEqual(ts.task_type, task_type)
self.assertEqual(ts.start_time, start_time)
self.assertEqual(ts.finish_time, finish_time)
self.assertEqual(ts.result, result)
self.assertEqual(ts.traceback, None)
self.assertEqual(ts.exception, None)
def test_updates_task_status_correctly(self, mock_request):
exc = Exception()
task_id = str(uuid.uuid4())
args = [1, 'b', 'iii']
kwargs = {'1': 'for the money', 'tags': ['test_tags']}
class EInfo(object):
"""
on_failure handler expects an instance of celery's ExceptionInfo class
as one of the attributes. It stores string representation of traceback
in it's traceback instance variable. This is a stub to imitate that behavior.
"""
def __init__(self):
self.traceback = "string_repr_of_traceback"
einfo = EInfo()
mock_request.called_directly = False
task_status = TaskStatus(task_id).save()
self.assertEqual(task_status['state'], 'waiting')
self.assertEqual(task_status['finish_time'], None)
self.assertEqual(task_status['traceback'], None)
task = tasks.Task()
task.on_failure(exc, task_id, args, kwargs, einfo)
new_task_status = TaskStatus.objects(task_id=task_id).first()
self.assertEqual(new_task_status['state'], 'error')
self.assertFalse(new_task_status['finish_time'] is None)
# Make sure that parse_iso8601_datetime is able to parse the finish_time without errors
dateutils.parse_iso8601_datetime(new_task_status['finish_time'])
self.assertEqual(new_task_status['traceback'], einfo.traceback)
def apply_async(self, *args, **kwargs):
"""
A wrapper around the Celery apply_async method. It allows us to accept a few more
parameters than Celery does for our own purposes, listed below. It also allows us
to create and update task status which can be used to track status of this task
during it's lifetime.
:param queue: The queue that the task has been placed into (optional, defaults to
the general Celery queue.)
:type queue: basestring
:param tags: A list of tags (strings) to place onto the task, used for searching for
tasks by tag
:type tags: list
:return: An AsyncResult instance as returned by Celery's apply_async
:rtype: celery.result.AsyncResult
"""
routing_key = kwargs.get(
'routing_key',
defaults.NAMESPACES['CELERY']['DEFAULT_ROUTING_KEY'].default)
tags = kwargs.pop('tags', [])
async_result = super(Task, self).apply_async(*args, **kwargs)
async_result.tags = tags
# Create a new task status with the task id and tags.
task_status = TaskStatus(task_id=async_result.id,
task_type=self.name,
state=constants.CALL_WAITING_STATE,
worker_name=routing_key,
tags=tags)
# To avoid the race condition where __call__ method below is called before
# this change is propagated to all db nodes, using an 'upsert' here and setting
# the task state to 'waiting' only on an insert.
task_status.save(fields_to_set_on_insert=['state', 'start_time'])
return async_result
def test_async_result(self, mock_request):
retval = AsyncResult('foo-id')
task_id = str(uuid.uuid4())
args = [1, 'b', 'iii']
kwargs = {
'1': 'for the money',
'tags': ['test_tags'],
'routing_key': WORKER_2
}
mock_request.called_directly = False
task_status = TaskStatus(task_id).save()
self.assertEqual(task_status['state'], 'waiting')
self.assertEqual(task_status['finish_time'], None)
task = tasks.Task()
task.on_success(retval, task_id, args, kwargs)
new_task_status = TaskStatus.objects(task_id=task_id).first()
self.assertEqual(new_task_status['state'], 'finished')
self.assertEqual(new_task_status['result'], None)
self.assertFalse(new_task_status['finish_time'] is None)
# Make sure that parse_iso8601_datetime is able to parse the finish_time without errors
dateutils.parse_iso8601_datetime(new_task_status['finish_time'])
self.assertEqual(new_task_status['spawned_tasks'], ['foo-id'])
def test_spawned_task_dict(self, mock_request):
retval = tasks.TaskResult(spawned_tasks=[{
'task_id': 'foo-id'
}],
result='bar')
task_id = str(uuid.uuid4())
args = [1, 'b', 'iii']
kwargs = {
'1': 'for the money',
'tags': ['test_tags'],
'routing_key': WORKER_2
}
mock_request.called_directly = False
task_status = TaskStatus(task_id).save()
self.assertEqual(task_status['state'], 'waiting')
self.assertEqual(task_status['finish_time'], None)
task = tasks.Task()
task.on_success(retval, task_id, args, kwargs)
new_task_status = TaskStatus.objects(task_id=task_id).first()
self.assertEqual(new_task_status['state'], 'finished')
self.assertEqual(new_task_status['result'], 'bar')
self.assertFalse(new_task_status['finish_time'] is None)
self.assertEqual(new_task_status['spawned_tasks'], ['foo-id'])
def test_create_task_status_invalid_task_id(self):
"""
Test that TaskStatus creation with an invalid task id raises the correct error.
"""
try:
TaskStatus(None).save()
except ValidationError, e:
self.assertTrue('task_id' in e.message)
def test_set_accepted(self):
task_id = self.get_random_uuid()
TaskStatus(task_id, state=constants.CALL_WAITING_STATE).save()
TaskStatus.objects(task_id=task_id, state=constants.CALL_WAITING_STATE).\
update_one(set__state=constants.CALL_ACCEPTED_STATE)
task_status = TaskStatus.objects.get(task_id=task_id)
self.assertTrue(task_status['state'], constants.CALL_ACCEPTED_STATE)
def test_GET_has_correct_worker_name_attribute(self):
task_id = '1234abcd'
TaskStatus(task_id, worker_name='worker1').save()
result = self.task_resource.GET(task_id)
result_json = json.loads(result)
self.assertTrue('worker_name' in result_json)
self.assertTrue(result_json['worker_name'] == 'worker1')
def test_DELETE_doesnt_cancel_spawned_celery_task(self, revoke):
"""
Test the DELETE() which should cause a revoke call to Celery's Controller.
This also tests that the spawned tasks are canceled as well.
"""
task_id = '1234abcd'
spawned_task_id = 'spawned_task'
spawned_by_spawned_task_id = 'spawned_by_spawned_task'
TaskStatus(task_id).save()
TaskStatus(spawned_task_id).save()
TaskStatus(spawned_by_spawned_task_id).save()
TaskStatus.objects(task_id=task_id).update_one(set__spawned_tasks=[spawned_task_id])
TaskStatus.objects(task_id=spawned_task_id).update_one(
set__spawned_tasks=[spawned_by_spawned_task_id])
self.task_resource.DELETE(task_id)
self.assertEqual(revoke.call_count, 1)
revoke.assert_called_once_with(task_id, terminate=True)
def test_DELETE_completed_celery_task(self):
"""
Test the DELETE() method does not change the state of a task that is already complete
"""
task_id = '1234abcd'
TaskStatus(task_id, state=constants.CALL_FINISHED_STATE).save()
self.task_resource.DELETE(task_id)
task_status = TaskStatus.objects(task_id=task_id).first()
self.assertEqual(task_status['state'], constants.CALL_FINISHED_STATE)
def test_save_insert_with_set_on_insert(self):
"""
Test the save method with set on insert arguments when the object is not already in the
database.
"""
task_id = str(uuid4())
worker_name = 'some_worker'
tags = ['tag_1', 'tag_2']
state = constants.CALL_RUNNING_STATE
spawned_tasks = ['foo']
error = {'error': 'some_error'}
progress_report = {
'what do we want?': 'progress!',
'when do we want it?': 'now!'
}
task_type = 'some.task'
start_time = datetime.now()
finish_time = start_time + timedelta(minutes=5)
start_time = dateutils.format_iso8601_datetime(start_time)
finish_time = dateutils.format_iso8601_datetime(finish_time)
result = None
ts = TaskStatus(task_id,
worker_name,
tags,
state,
spawned_tasks=spawned_tasks,
error=error,
progress_report=progress_report,
task_type=task_type,
start_time=start_time,
finish_time=finish_time,
result=result)
# This should cause ts to be in the database
ts.save_with_set_on_insert(
fields_to_set_on_insert=['state', 'start_time'])
ts = TaskStatus.objects()
# There should only be one TaskStatus in the db
self.assertEqual(len(ts), 1)
ts = ts[0]
# Make sure all the attributes are correct
self.assertEqual(ts['task_id'], task_id)
self.assertEqual(ts['worker_name'], worker_name)
self.assertEqual(ts['tags'], tags)
self.assertEqual(ts['state'], state)
self.assertEqual(ts['error'], error)
self.assertEqual(ts['spawned_tasks'], spawned_tasks)
self.assertEqual(ts['progress_report'], progress_report)
self.assertEqual(ts['task_type'], task_type)
self.assertEqual(ts['start_time'], start_time)
self.assertEqual(ts['finish_time'], finish_time)
self.assertEqual(ts['result'], result)
# These are always None
self.assertEqual(ts['traceback'], None)
self.assertEqual(ts['exception'], None)
def _test_GET_celery_tasks_by_tags(self):
"""
Test the GET() method to get all current tasks.
"""
# Populate a few of task statuses
task_id1 = str(uuid.uuid4())
worker_1 = 'worker_1'
state1 = 'waiting'
tags1 = ['random', 'tags']
task_id2 = str(uuid.uuid4())
worker_2 = 'worker_2'
state2 = 'running'
tags2 = ['random', 'tags']
task_id3 = str(uuid.uuid4())
worker_3 = 'worker_3'
state3 = 'running'
tags3 = ['random']
TaskStatus(task_id1, worker_1, tags1, state1).save()
TaskStatus(task_id2, worker_2, tags2, state2).save()
TaskStatus(task_id3, worker_3, tags3, state3).save()
# Validate for tags
status, body = self.get('/v2/tasks/?tag=random&tag=tags')
self.assertEqual(200, status)
self.assertTrue(len(body) == 2)
for task in body:
if task['task_id'] == task_id1:
self.assertEquals(task['state'], state1)
self.assertEqual(task['worker_name'], worker_1)
self.assertEqual(task['tags'], tags1)
else:
self.assertEqual(task['task_id'], task_id2)
self.assertEquals(task['state'], state2)
self.assertEqual(task['worker_name'], worker_2)
self.assertEquals(task['tags'], tags2)
# Negative test
status, body = self.get('/v2/tasks/?tag=non_existent')
self.assertEqual(200, status)
self.assertTrue(len(body) == 0)
def save_update_with_set_on_insert(self):
"""
Test the save method with set on insert arguments when the object is already in the
database.
"""
task_id = 'a_task_id'
queue = 'some_queue'
tags = ['tag_1', 'tag_2']
state = 'a state'
spawned_tasks = ['foo']
error = 'some_error'
progress_report = {'what do we want?': 'progress!', 'when do we want it?': 'now!'}
task_type = 'some.task'
start_time = datetime.now()
finish_time = start_time + timedelta(minutes=5)
result = None
ts = TaskStatus(
task_id, queue, tags, state, spawned_tasks=spawned_tasks, error=error,
progress_report=progress_report, task_type=task_type, start_time=start_time,
finish_time=finish_time, result=result)
# Put the object in the database, and then change some of it settings.
ts.save()
new_queue = 'a_different_queue'
new_state = 'some_other_state'
new_start_time = start_time + timedelta(minutes=10)
ts.queue = new_queue
ts.state = new_state
ts.start_time = new_start_time
# This should update the queue on ts in the database, but should not update the state or
# start_time
ts.save(fields_to_set_on_insert=['state', 'start_time'])
ts = TaskStatus.get_collection().find()
# There should only be one TaskStatus in the db
self.assertEqual(len(ts), 1)
ts = ts[0]
# Make sure all the attributes are correct
self.assertEqual(ts['task_id'], task_id)
# Queue should have been updated
self.assertEqual(ts['queue'], new_queue)
self.assertEqual(ts['tags'], tags)
# state should not have been updated
self.assertEqual(ts['state'], state)
self.assertEqual(ts['error'], error)
self.assertEqual(ts['spawned_tasks'], spawned_tasks)
self.assertEqual(ts['progress_report'], progress_report)
self.assertEqual(ts['task_type'], task_type)
# start_time should not have been updated
self.assertEqual(ts['start_time'], start_time)
self.assertEqual(ts['finish_time'], finish_time)
self.assertEqual(ts['result'], result)
# These are always None
self.assertEqual(ts['traceback'], None)
self.assertEqual(ts['exception'], None)
def test_DELETE_celery_task(self, revoke):
"""
Test the DELETE() method with a UUID that does not correspond to a UUID that the
coordinator is aware of. This should cause a revoke call to Celery's Controller.
"""
task_id = '1234abcd'
TaskStatus(task_id).save()
self.task_resource.DELETE(task_id)
revoke.assert_called_once_with(task_id, terminate=True)
def test_cancel_after_task_canceled(self, *unused_mocks):
"""
Test that canceling a task that was already canceled results in no change
to the task state.
"""
task_id = '1234abcd'
TaskStatus(task_id, 'test_worker', state=CALL_CANCELED_STATE).save()
tasks.cancel(task_id)
task_status = TaskStatus.objects(task_id=task_id).first()
self.assertEqual(task_status['state'], CALL_CANCELED_STATE)
def test_cancel_after_task_finished(self, _logger, revoke):
"""
Test that canceling a task that is already finished results in no change
to the task state.
"""
task_id = '1234abcd'
TaskStatus(task_id, 'test_worker', state=CALL_FINISHED_STATE).save()
tasks.cancel(task_id)
task_status = TaskStatus.objects(task_id=task_id).first()
self.assertEqual(task_status['state'], CALL_FINISHED_STATE)
