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)