diff --git a/awx/main/dispatch/worker/task.py b/awx/main/dispatch/worker/task.py index d1273749a1..e63c7a224a 100644 --- a/awx/main/dispatch/worker/task.py +++ b/awx/main/dispatch/worker/task.py @@ -30,11 +30,18 @@ class TaskWorker(BaseWorker): awx.main.tasks.delete_inventory awx.main.tasks.RunProjectUpdate ''' + if not task.startswith('awx.'): + raise ValueError('{} is not a valid awx task'.format(task)) module, target = task.rsplit('.', 1) module = importlib.import_module(module) _call = None if hasattr(module, target): _call = getattr(module, target, None) + if not ( + hasattr(_call, 'apply_async') and hasattr(_call, 'delay') + ): + raise ValueError('{} is not decorated with @task()'.format(task)) + return _call def run_callable(self, body): @@ -78,6 +85,7 @@ class TaskWorker(BaseWorker): try: result = self.run_callable(body) except Exception as exc: + result = exc try: if getattr(exc, 'is_awx_task_error', False): diff --git a/awx/main/tests/functional/test_dispatch.py b/awx/main/tests/functional/test_dispatch.py index 3567d7d37a..e73f226bbc 100644 --- a/awx/main/tests/functional/test_dispatch.py +++ b/awx/main/tests/functional/test_dispatch.py @@ -14,6 +14,10 @@ from awx.main.dispatch.publish import task from awx.main.dispatch.worker import BaseWorker, TaskWorker +def restricted(a, b): + raise AssertionError("This code should not run because it isn't decorated with @task") + + @task() def add(a, b): return a + b @@ -25,6 +29,11 @@ class BaseTask(object): return add(a, b) +class Restricted(object): + def run(self, a, b): + raise AssertionError("This code should not run because it isn't decorated with @task") + + @task() class Adder(BaseTask): def run(self, a, b): @@ -262,6 +271,14 @@ class TestTaskDispatcher: }) assert result == 4 + def test_function_dispatch_must_be_decorated(self): + result = self.tm.perform_work({ + 'task': 'awx.main.tests.functional.test_dispatch.restricted', + 'args': [2, 2] + }) + assert isinstance(result, ValueError) + assert result.message == 'awx.main.tests.functional.test_dispatch.restricted is not decorated with @task()' # noqa + def test_method_dispatch(self): result = self.tm.perform_work({ 'task': 'awx.main.tests.functional.test_dispatch.Adder', @@ -269,6 +286,29 @@ class TestTaskDispatcher: }) assert result == 4 + def test_method_dispatch_must_be_decorated(self): + result = self.tm.perform_work({ + 'task': 'awx.main.tests.functional.test_dispatch.Restricted', + 'args': [2, 2] + }) + assert isinstance(result, ValueError) + assert result.message == 'awx.main.tests.functional.test_dispatch.Restricted is not decorated with @task()' # noqa + + def test_python_function_cannot_be_imported(self): + result = self.tm.perform_work({ + 'task': 'os.system', + 'args': ['ls'], + }) + assert isinstance(result, ValueError) + assert result.message == 'os.system is not a valid awx task' # noqa + + def test_undefined_function_cannot_be_imported(self): + result = self.tm.perform_work({ + 'task': 'awx.foo.bar' + }) + assert isinstance(result, ImportError) + assert result.message == 'No module named foo' # noqa + class TestTaskPublisher: