def test__release_resource_task_count_two(self):
"""
Test _release_resource() with a resource that has a task count of two. This should simply
decrement the task_count for the resource, but should not remove it from the database.
"""
# Set up two workers
now = datetime.utcnow()
worker_1 = Worker(WORKER_1, now)
worker_1.save()
worker_2 = Worker(WORKER_2, now)
worker_2.save()
# Set up two resource reservations, using our workers from above
reserved_resource_1 = ReservedResource('resource_1', worker_1.name, 7)
reserved_resource_1.save()
reserved_resource_2 = ReservedResource('resource_2', worker_2.name, 2)
reserved_resource_2.save()
# This should reduce the reserved_resource_2 num_reservations to 1.
tasks._release_resource('resource_2')
# Make sure the ReservedResources are also correct
rrc = ReservedResource.get_collection()
self.assertEqual(rrc.count(), 2)
rr_1 = rrc.find_one({'_id': reserved_resource_1.name})
self.assertEqual(rr_1['assigned_queue'],
reserved_resource_1.assigned_queue)
self.assertEqual(rr_1['num_reservations'], 7)
rr_2 = rrc.find_one({'_id': reserved_resource_2.name})
self.assertEqual(rr_2['assigned_queue'],
reserved_resource_2.assigned_queue)
self.assertEqual(rr_2['num_reservations'], 1)
def test__release_resource_task_count_one(self):
"""
Test _release_resource() with a resource that has a task count of one. This should remove
the resource from the database.
"""
# Set up two workers
now = datetime.utcnow()
worker_1 = Worker(WORKER_1, now)
worker_1.save()
worker_2 = Worker(WORKER_2, now)
worker_2.save()
# Set up two reserved resources
reserved_resource_1 = ReservedResource('resource_1', worker_1.name, 7)
reserved_resource_1.save()
reserved_resource_2 = ReservedResource('resource_2', worker_2.name, 1)
reserved_resource_2.save()
# This should remove resource_2 from the _resource_map.
tasks._release_resource('resource_2')
# resource_2 should have been removed from the database
rrc = ReservedResource.get_collection()
self.assertEqual(rrc.count(), 1)
rr_1 = rrc.find_one({'_id': reserved_resource_1.name})
self.assertEqual(rr_1['assigned_queue'],
reserved_resource_1.assigned_queue)
self.assertEqual(rr_1['num_reservations'], 7)
def test_resource_in_resource_map(self):
"""
Test _release_resource() with a valid resource. This should remove the resource from the
database.
"""
# Set up two workers
now = datetime.utcnow()
worker_1 = Worker(WORKER_1, now)
worker_1.save()
worker_2 = Worker(WORKER_2, now)
worker_2.save()
# Set up two reserved resources
reserved_resource_1 = ReservedResource(uuid.uuid4(), worker_1.name,
'resource_1')
reserved_resource_1.save()
reserved_resource_2 = ReservedResource(uuid.uuid4(), worker_2.name,
'resource_2')
reserved_resource_2.save()
# This should remove resource_2 from the _resource_map.
tasks._release_resource(reserved_resource_2.task_id)
# resource_2 should have been removed from the database
rrc = ReservedResource.get_collection()
self.assertEqual(rrc.count(), 1)
rr_1 = rrc.find_one({'_id': reserved_resource_1.task_id})
self.assertEqual(rr_1['worker_name'], reserved_resource_1.worker_name)
self.assertEqual(rr_1['resource_id'], 'resource_1')
def test_filter(self):
"""
Test a filter operation to make sure the results appear to be correct.
"""
# Make three workers. We'll filter for two of them.
now = datetime.utcnow()
kw_1 = Worker('worker_1', now)
kw_1.save()
kw_2 = Worker('worker_2', now)
kw_2.save()
kw_3 = Worker('worker_3', now)
kw_3.save()
criteria = Criteria(filters={'_id': {
'$gt': 'worker_1'
}},
sort=[('_id', pymongo.ASCENDING)])
workers = resources.filter_workers(criteria)
# Let's assert that workers is a generator, and then let's cast it to a list so it's easier
# to test that we got the correct instances back.
self.assertEqual(type(workers), types.GeneratorType)
workers = list(workers)
self.assertEqual(all([isinstance(w, Worker) for w in workers]), True)
self.assertEqual(workers[0].name, 'worker_2')
self.assertEqual(workers[1].name, 'worker_3')
def test_ignores_queues_that_arent_workers(self):
"""
It is possible for the assigned_queue in a ReservedResource to reference a queue that is not
in the workers collection. This test ensures that this queue is properly ignored, even if it
is the most "enticing" choice.
"""
# Set up three Workers, with the least busy one in the middle so that we can
# demonstrate that it did pick the least busy and not the last or first.
now = datetime.utcnow()
worker_1 = Worker('busy_worker', now)
worker_2 = Worker('less_busy_worker', now)
worker_3 = Worker('most_busy_worker', now)
for worker in (worker_1, worker_2, worker_3):
worker.save()
# Now we need to make some reservations against these Workers' queues. We'll give worker_1
# 8 reservations, putting it in the middle of busyness.
rr_1 = ReservedResource(name='resource_1',
assigned_queue=worker_1.queue_name,
num_reservations=8)
# These next two will give worker_2 a total of 7 reservations, so it should get picked.
rr_2 = ReservedResource(name='resource_2',
assigned_queue=worker_2.queue_name,
num_reservations=3)
rr_3 = ReservedResource(name='resource_3',
assigned_queue=worker_2.queue_name,
num_reservations=4)
# These next three will give worker_3 a total of 9 reservations, so it should be the most
# busy.
rr_4 = ReservedResource(name='resource_4',
assigned_queue=worker_3.queue_name,
num_reservations=2)
rr_5 = ReservedResource(name='resource_5',
assigned_queue=worker_3.queue_name,
num_reservations=3)
rr_6 = ReservedResource(name='resource_6',
assigned_queue=worker_3.queue_name,
num_reservations=4)
# Now we will make a ReservedResource that references a queue that does not correspond to a
# Worker and has the lowest reservation count. This RR should be ignored.
rr_7 = ReservedResource(name='resource_7',
assigned_queue='doesnt_exist',
num_reservations=1)
for rr in (rr_1, rr_2, rr_3, rr_4, rr_5, rr_6, rr_7):
rr.save()
worker = resources.get_least_busy_worker()
self.assertEqual(type(worker), Worker)
self.assertEqual(worker.name, 'less_busy_worker')
def test__reserve_resource_with_existing_reservation(self):
"""
Test _reserve_resource() with a resource that has an existing reservation in the database.
It should return the queue listed in the database, and increment the reservation counter.
"""
# Set up a worker with a reservation count of 1
now = datetime.utcnow()
worker_1 = Worker(WORKER_1, now)
worker_1.save()
# Set up a resource reservation, using our worker from above
reserved_resource_1 = ReservedResource('resource_1',
worker_1.queue_name, 1)
reserved_resource_1.save()
# This should increase the reserved_resource_1 num_reservations to 2. worker_1's name should
# be returned
queue = tasks._reserve_resource('resource_1')
self.assertEqual(queue, WORKER_1_QUEUE)
# Make sure the ReservedResource is correct
rrc = ReservedResource.get_collection()
self.assertEqual(rrc.count(), 1)
rr_1 = rrc.find_one({'_id': reserved_resource_1.name})
self.assertEqual(rr_1['assigned_queue'], WORKER_1_QUEUE)
self.assertEqual(rr_1['num_reservations'], 2)
def test_post(self, mock_publish, mock_get_worker_for_reservation):
"""
Test that publish repo group creates a task for a worker.
"""
# Setup
group_id = 'group-1'
distributor_id = 'dist-1'
mock_publish.apply_async_with_reservation.return_value = AsyncResult(
'1234')
self.manager.create_repo_group(group_id)
self.distributor_manager.add_distributor(group_id,
'dummy-group-distributor', {},
distributor_id=distributor_id)
mock_get_worker_for_reservation.return_value = Worker(
'some_queue', datetime.datetime.now())
# Test
data = {'id': distributor_id}
status, body = self.post(
'/v2/repo_groups/%s/actions/publish/' % group_id, data)
# Verify
self.assertEqual(202, status)
self.assertEqual(mock_publish.apply_async_with_reservation.call_count,
1)
def test_post_with_override_config(self, mock_get_worker_for_reservation, mock_uuid,
mock_apply_async):
# Setup
uuid_list = [uuid.uuid4() for i in range(10)]
mock_uuid.uuid4.side_effect = copy.deepcopy(uuid_list)
expected_async_result = AsyncResult(str(uuid_list[0]))
mock_get_worker_for_reservation.return_value = Worker('some_queue', datetime.datetime.now())
upload_id = self.upload_manager.initialize_upload()
self.upload_manager.save_data(upload_id, 0, 'string data')
repo_manager = manager_factory.repo_manager()
repo_manager.create_repo('repo-upload')
importer_manager = manager_factory.repo_importer_manager()
importer_manager.set_importer('repo-upload', 'dummy-importer', {})
# Test
test_override_config = {'key1': 'value1', 'key2': 'value2'}
body = {
'upload_id' : upload_id,
'unit_type_id' : 'dummy-type',
'unit_key' : {'name' : 'foo'},
'unit_metadata' : {'stuff' : 'bar'},
'override_config': test_override_config,
}
status, body = self.post('/v2/repositories/repo-upload/actions/import_upload/', body)
# Verify
self.assertEqual(202, status)
assert_body_matches_async_task(body, expected_async_result)
exepcted_call_args = ['repo-upload', 'dummy-type',
{'name': 'foo'}, {'stuff': 'bar'},
upload_id, test_override_config]
self.assertEqual(exepcted_call_args, mock_apply_async.call_args[0][0])
def handle_worker_heartbeat(event):
"""
Celery event handler for 'worker-heartbeat' events.
The event is first parsed and logged. If this event is from the resource manager, there is
no further processing to be done. Then the existing Worker objects are searched
for one to update. If an existing one is found, it is updated. Otherwise a new
Worker entry is created. Logging at the info and debug level is also done.
:param event: A celery event to handle.
:type event: dict
"""
event_info = _parse_and_log_event(event)
# if this is the resource_manager do nothing
if _is_resource_manager(event):
return
find_worker_criteria = Criteria(filters={'_id': event_info['worker_name']},
fields=('_id', 'last_heartbeat',
'num_reservations'))
find_worker_list = list(resources.filter_workers(find_worker_criteria))
if find_worker_list:
Worker.get_collection().find_and_modify(
query={'_id': event_info['worker_name']},
update={'$set': {
'last_heartbeat': event_info['timestamp']
}})
else:
new_worker = Worker(event_info['worker_name'], event_info['timestamp'])
msg = _("New worker '%(worker_name)s' discovered") % event_info
_logger.info(msg)
new_worker.save()
def test_dispatches__release_resource(self):
self.mock_get_worker_for_reservation.return_value = Worker(
'worker1', datetime.utcnow())
tasks._queue_reserved_task('task_name', 'my_task_id', 'my_resource_id',
[1, 2], {'a': 2})
self.mock__release_resource.apply_async.assert_called_once_with(
('my_task_id', ), routing_key='worker1', exchange='C.dq')
def test_get_unreserved_worker_breaks_out_of_loop(self):
self.mock_get_worker_for_reservation.side_effect = NoWorkers()
self.mock_get_unreserved_worker.return_value = Worker(
'worker1', datetime.utcnow())
tasks._queue_reserved_task('task_name', 'my_task_id', 'my_resource_id',
[1, 2], {'a': 2})
self.assertTrue(not self.mock_time.sleep.called)
def test_update_repo_and_plugins(self, distributor_update,
mock_get_worker_for_reservation):
"""
Tests the aggregate call to update a repo and its plugins.
"""
mock_get_worker_for_reservation.return_value = Worker(
'some_queue', datetime.datetime.now())
self.manager.create_repo('repo-1', 'Original', 'Original Description')
importer_manager = manager_factory.repo_importer_manager()
distributor_manager = manager_factory.repo_distributor_manager()
importer_manager.set_importer('repo-1', 'mock-importer',
{'key-i1': 'orig-1'})
distributor_manager.add_distributor('repo-1',
'mock-distributor',
{'key-d1': 'orig-1'},
True,
distributor_id='dist-1')
distributor_manager.add_distributor('repo-1',
'mock-distributor',
{'key-d2': 'orig-2'},
True,
distributor_id='dist-2')
# Test
repo_delta = {'display_name': 'Updated'}
new_importer_config = {'key-i1': 'updated-1', 'key-i2': 'new-1'}
new_distributor_configs = {
'dist-1': {
'key-d1': 'updated-1'
},
} # only update one of the two distributors
result = self.manager.update_repo_and_plugins('repo-1', repo_delta,
new_importer_config,
new_distributor_configs)
self.assertTrue(isinstance(result, TaskResult))
self.assertEquals(None, result.error)
repo = result.return_value
# Verify
self.assertEqual(repo['id'], 'repo-1')
self.assertEqual(repo['display_name'], 'Updated')
self.assertEqual(repo['description'], 'Original Description')
importer = importer_manager.get_importer('repo-1')
self.assertEqual(importer['config'], new_importer_config)
dist_1 = distributor_manager.get_distributor('repo-1', 'dist-1')
self.assertEqual(dist_1['config'], new_distributor_configs['dist-1'])
dist_2 = distributor_manager.get_distributor('repo-1', 'dist-2')
self.assertEqual(dist_2['config'], {'key-d2': 'orig-2'})
# There should have been a spawned task for the new distributor config
expected_task_id = dispatch.TaskStatus.get_collection().find_one(
{'tags': 'pulp:repository_distributor:dist-1'})['task_id']
self.assertEqual(result.spawned_tasks, [{'task_id': expected_task_id}])
def test_creates_and_saves_reserved_resource(self):
self.mock_get_worker_for_reservation.return_value = Worker(
'worker1', datetime.utcnow())
tasks._queue_reserved_task('task_name', 'my_task_id', 'my_resource_id',
[1, 2], {'a': 2})
self.mock_reserved_resource.assert_called_once_with(
'my_task_id', 'worker1', 'my_resource_id')
self.mock_reserved_resource.return_value.save.assert_called_once_with()
def test_picks_least_busy_worker(self):
"""
Test that the function picks the least busy worker.
"""
# Set up three Workers, with the least busy one in the middle so that we can
# demonstrate that it did pick the least busy and not the last or first.
now = datetime.utcnow()
worker_1 = Worker('busy_worker', now)
worker_2 = Worker('less_busy_worker', now)
worker_3 = Worker('most_busy_worker', now)
for worker in (worker_1, worker_2, worker_3):
worker.save()
# Now we need to make some reservations against these Workers' queues. We'll give worker_1
# 8 reservations, putting it in the middle of busyness.
rr_1 = ReservedResource(name='resource_1',
assigned_queue=worker_1.queue_name,
num_reservations=8)
# These next two will give worker_2 a total of 7 reservations, so it should get picked.
rr_2 = ReservedResource(name='resource_2',
assigned_queue=worker_2.queue_name,
num_reservations=3)
rr_3 = ReservedResource(name='resource_3',
assigned_queue=worker_2.queue_name,
num_reservations=4)
# These next three will give worker_3 a total of 9 reservations, so it should be the most
# busy.
rr_4 = ReservedResource(name='resource_4',
assigned_queue=worker_3.queue_name,
num_reservations=2)
rr_5 = ReservedResource(name='resource_5',
assigned_queue=worker_3.queue_name,
num_reservations=3)
rr_6 = ReservedResource(name='resource_6',
assigned_queue=worker_3.queue_name,
num_reservations=4)
for rr in (rr_1, rr_2, rr_3, rr_4, rr_5, rr_6):
rr.save()
worker = resources.get_least_busy_worker()
self.assertEqual(type(worker), Worker)
self.assertEqual(worker.name, 'less_busy_worker')
def GET(self, task_id):
try:
task = TaskStatus.objects.get(task_id=task_id)
except DoesNotExist:
raise MissingResource(task_id)
task_dict = task_serializer(task)
if 'worker_name' in task_dict:
queue_name = Worker(task_dict['worker_name'], datetime.now()).queue_name
task_dict.update({'queue': queue_name})
return self.ok(task_dict)
def test_post(self, mock_get_worker_for_reservation):
# Setup
group_id = 'group-1'
distributor_id = 'dist-1'
self.manager.create_repo_group(group_id)
self.distributor_manager.add_distributor(group_id, 'dummy-group-distributor', {}, distributor_id=distributor_id)
mock_get_worker_for_reservation.return_value = Worker('some_queue', datetime.datetime.now())
# Test
data = {'id' : distributor_id}
status, body = self.post('/v2/repo_groups/%s/actions/publish/' % group_id, data)
def test_DELETE_completed_celery_task(self):
"""
Test the DELETE() method raises a TaskComplete exception if the task is already complete.
"""
task_id = '1234abcd'
now = datetime.utcnow()
test_worker = Worker('test_worker', now)
TaskStatusManager.create_task_status(
task_id, test_worker.name, state=constants.CALL_FINISHED_STATE)
self.assertRaises(PulpCodedException, self.task_resource.DELETE,
task_id)
def test_cancel_after_task_finished(self, logger, revoke):
task_id = '1234abcd'
now = datetime.utcnow()
test_worker = Worker('test_worker', now)
TaskStatusManager.create_task_status(task_id,
test_worker.name,
state=CALL_FINISHED_STATE)
self.assertRaises(PulpCodedException, tasks.cancel, task_id)
task_status = TaskStatusManager.find_by_task_id(task_id)
self.assertEqual(task_status['state'], CALL_FINISHED_STATE)
def GET(self, task_id):
task = TaskStatusManager.find_by_task_id(task_id)
if task is None:
raise MissingResource(task_id)
else:
link = serialization.link.link_obj('/pulp/api/v2/tasks/%s/' % task_id)
task.update(link)
task.update(serialization.dispatch.spawned_tasks(task))
if 'worker_name' in task:
queue_name = Worker(task['worker_name'], datetime.now()).queue_name
task.update({'queue': queue_name})
return self.ok(task)
def test_dispatches_inner_task(self):
self.mock_get_worker_for_reservation.return_value = Worker(
'worker1', datetime.utcnow())
tasks._queue_reserved_task('task_name', 'my_task_id', 'my_resource_id',
[1, 2], {'a': 2})
apply_async = self.mock_celery.tasks['task_name'].apply_async
apply_async.assert_called_once_with(1,
2,
a=2,
routing_key='worker1',
task_id='my_task_id',
exchange='C.dq')
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'
now = datetime.utcnow()
test_worker = Worker('test_worker', now)
TaskStatusManager.create_task_status(task_id, test_worker.name)
self.task_resource.DELETE(task_id)
revoke.assert_called_once_with(task_id, terminate=True)
def test__release_resource_not_in__resource_map(self):
"""
Test _release_resource() with a resource that is not in the database. This should be
gracefully handled, and result in no changes to the database.
"""
# Set up two workers
worker_1 = Worker(WORKER_1, datetime.utcnow())
worker_1.save()
worker_2 = Worker(WORKER_2, datetime.utcnow())
worker_2.save()
# Set up two resource reservations, using our workers from above
reserved_resource_1 = ReservedResource('resource_1', worker_1.name, 7)
reserved_resource_1.save()
reserved_resource_2 = ReservedResource('resource_2', worker_2.name, 3)
reserved_resource_2.save()
# This should not raise any Exception, but should also not alter either the Worker
# collection or the ReservedResource collection
tasks._release_resource('made_up_resource_id')
# Make sure that the workers collection has not been altered
worker_collection = Worker.get_collection()
self.assertEqual(worker_collection.count(), 2)
worker_1 = worker_collection.find_one({'_id': worker_1.name})
self.assertTrue(worker_1)
worker_2 = worker_collection.find_one({'_id': worker_2.name})
self.assertTrue(worker_2)
# Make sure that the reserved resources collection has not been altered
rrc = ReservedResource.get_collection()
self.assertEqual(rrc.count(), 2)
rr_1 = rrc.find_one({'_id': reserved_resource_1.name})
self.assertEqual(rr_1['assigned_queue'],
reserved_resource_1.assigned_queue)
self.assertEqual(rr_1['num_reservations'], 7)
rr_2 = rrc.find_one({'_id': reserved_resource_2.name})
self.assertEqual(rr_2['assigned_queue'],
reserved_resource_2.assigned_queue)
self.assertEqual(rr_2['num_reservations'], 3)
def test_cancel_successful(self, logger, revoke):
task_id = '1234abcd'
now = datetime.utcnow()
test_worker = Worker('test_worker', now)
TaskStatusManager.create_task_status(task_id, test_worker.name)
tasks.cancel(task_id)
revoke.assert_called_once_with(task_id, terminate=True)
self.assertEqual(logger.info.call_count, 1)
log_msg = logger.info.mock_calls[0][1][0]
self.assertTrue(task_id in log_msg)
self.assertTrue('Task canceled' in log_msg)
task_status = TaskStatusManager.find_by_task_id(task_id)
self.assertEqual(task_status['state'], CALL_CANCELED_STATE)
def test_apply_async_task_canceled(self, apply_async):
"""
Assert that apply_async() honors 'canceled' task status.
"""
args = [1, 'b', 'iii']
kwargs = {'1': 'for the money', 'tags': ['test_tags']}
task_id = 'test_task_id'
now = datetime.utcnow()
TaskStatusManager.create_task_status(task_id,
Worker('test-worker', now),
state=CALL_CANCELED_STATE)
apply_async.return_value = celery.result.AsyncResult(task_id)
task = tasks.Task()
task.apply_async(*args, **kwargs)
task_status = TaskStatusManager.find_by_task_id(task_id)
self.assertEqual(task_status['state'], 'canceled')
self.assertEqual(task_status['start_time'], None)
def test__reserve_resource_without_existing_reservation(self):
"""
Test _reserve_resource() with a resource that does not have an existing reservation in the
database. It should find the least busy worker, add a reservation to the database with that
worker's queue, and then return the queue name.
"""
# Set up a worker
worker_1 = Worker(WORKER_1, datetime.utcnow())
worker_1.save()
queue = tasks._reserve_resource('resource_1')
self.assertEqual(queue, WORKER_1_QUEUE)
# Make sure the ReservedResource is correct
rrc = ReservedResource.get_collection()
self.assertEqual(rrc.count(), 1)
rr_1 = rrc.find_one({'_id': 'resource_1'})
self.assertEqual(rr_1['assigned_queue'], WORKER_1_QUEUE)
self.assertEqual(rr_1['num_reservations'], 1)
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'
now = datetime.utcnow()
test_worker = Worker('test_worker', now)
TaskStatusManager.create_task_status(task_id, test_worker.queue_name)
TaskStatusManager.create_task_status(spawned_task_id,
test_worker.queue_name)
TaskStatusManager.create_task_status(spawned_by_spawned_task_id,
test_worker.queue_name)
TaskStatusManager.update_task_status(
task_id, delta={'spawned_tasks': [spawned_task_id]})
TaskStatusManager.update_task_status(
spawned_task_id,
delta={'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)
