def test_get_extension_tasks(self, app):
main()
extensions = sorted(entrypoint.get_extension_tasks('plugin2'))
self.assertEqual(extensions, [
'girder_worker._test_plugins.tasks.celery_task',
'girder_worker._test_plugins.tasks.function_task'
])
def test_get_extension_tasks():
with mock.patch('girder_worker.__main__.app'):
main()
extensions = sorted(entrypoint.get_extension_tasks('plugin2'))
assert extensions == [
'girder_worker._test_plugins.tasks.celery_task',
'girder_worker._test_plugins.tasks.function_task'
]
def test_register_extension():
@decorators.argument('n', types.Integer)
def echo(n):
return n
@app.task
@decorators.argument('n', types.Integer)
def echo_celery(n):
return n
tasks = {
'%s.echo' % __name__: echo,
'%s.echo_celery' % __name__: echo_celery
}
entrypoint.register_extension('echo_tasks', tasks)
exts = entrypoint.get_extensions()
assert 'echo_tasks' in exts
assert entrypoint.get_extension_tasks('echo_tasks') == tasks
celery_tasks = entrypoint.get_extension_tasks('echo_tasks',
celery_only=True)
assert list(celery_tasks.keys()) == ['%s.echo_celery' % __name__]
def test_register_extension(self):
@decorators.argument('n', types.Integer)
def echo(n):
return n
@app.task
@decorators.argument('n', types.Integer)
def echo_celery(n):
return n
tasks = {
'%s.echo' % __name__: echo,
'%s.echo_celery' % __name__: echo_celery
}
entrypoint.register_extension('echo_tasks', tasks)
exts = entrypoint.get_extensions()
self.assertIn('echo_tasks', exts)
self.assertEqual(entrypoint.get_extension_tasks('echo_tasks'), tasks)
celery_tasks = entrypoint.get_extension_tasks('echo_tasks',
celery_only=True)
self.assertEqual(celery_tasks.keys(), ['%s.echo_celery' % __name__])
def celeryTaskDescriptionForFolder(self, folder, extension, params):
import_all_includes()
try:
tasks = get_extension_tasks(extension, celery_only=True)
except KeyError:
raise RestException('Unknown girder_worker extension')
user = self.getCurrentUser()
itemModel = self.model('item')
items = []
for name, func in six.iteritems(tasks):
desc = decorators.describe_function(func)
item = itemModel.createItem(name=desc['name'],
creator=user,
folder=folder,
description=desc.get('description', ''),
reuseExisting=True)
items.append(describeItemTaskFromFunction(func, item, name))
return items
def celeryTaskDescriptionForFolder(self, folder, extension, params):
import_all_includes()
try:
tasks = get_extension_tasks(extension, celery_only=True)
except KeyError:
raise RestException('Unknown girder_worker extension')
user = self.getCurrentUser()
itemModel = self.model('item')
items = []
for name, func in six.iteritems(tasks):
desc = decorators.describe_function(func)
item = itemModel.createItem(
name=desc['name'],
creator=user,
folder=folder,
description=desc.get('description', ''),
reuseExisting=True
)
items.append(describeItemTaskFromFunction(func, item, name))
return items