def test_cgroup_task_runner_super_calls(self, mock_super_on_finish, mock_super_init):
"""
This test ensures that initiating CgroupTaskRunner object
calls init method of BaseTaskRunner,
and when task finishes, CgroupTaskRunner.on_finish() calls
super.on_finish() to delete the temp cfg file.
"""
local_task_job = mock.Mock()
local_task_job.task_instance = mock.MagicMock()
local_task_job.task_instance.run_as_user = None
local_task_job.task_instance.command_as_list.return_value = ['sleep', '1000']
runner = CgroupTaskRunner(local_task_job)
self.assertTrue(mock_super_init.called)
runner.on_finish()
self.assertTrue(mock_super_on_finish.called)
def get_task_runner(local_task_job):
"""
Get the task runner that can be used to run the given job.
:param local_task_job: The LocalTaskJob associated with the TaskInstance
that needs to be executed.
:type local_task_job: airflow.jobs.LocalTaskJob
:return: The task runner to use to run the task.
:rtype: airflow.task_runner.base_task_runner.BaseTaskRunner
"""
if _TASK_RUNNER == "BashTaskRunner":
return BashTaskRunner(local_task_job)
elif _TASK_RUNNER == "CgroupTaskRunner":
from airflow.contrib.task_runner.cgroup_task_runner import CgroupTaskRunner
return CgroupTaskRunner(local_task_job)
else:
raise AirflowException("Unknown task runner type {}".format(_TASK_RUNNER))