def test_job_id_setter(self): cel_exec = CeleryExecutor() k8s_exec = KubernetesExecutor() cel_k8s_exec = CeleryKubernetesExecutor(cel_exec, k8s_exec) job_id = 'this-job-id' cel_k8s_exec.job_id = job_id assert cel_exec.job_id == k8s_exec.job_id == cel_k8s_exec.job_id == job_id
def when_using_celery_executor(): celery_executor_mock = mock.MagicMock() k8s_executor_mock = mock.MagicMock() cke = CeleryKubernetesExecutor(celery_executor_mock, k8s_executor_mock) ti = mock.MagicMock() ti.queue = 'non-kubernetes-queue' cke.queue_task_instance( ti, mark_success, pickle_id, ignore_all_deps, ignore_depends_on_past, ignore_task_deps, ignore_ti_state, pool, cfg_path ) k8s_executor_mock.queue_task_instance.assert_not_called() celery_executor_mock.queue_task_instance.assert_called_once_with( ti, mark_success, pickle_id, ignore_all_deps, ignore_depends_on_past, ignore_task_deps, ignore_ti_state, pool, cfg_path )
def test_queue_task_instance(self, test_queue): celery_executor_mock = mock.MagicMock() k8s_executor_mock = mock.MagicMock() cke = CeleryKubernetesExecutor(celery_executor_mock, k8s_executor_mock) ti = mock.MagicMock() ti.queue = test_queue kwargs = dict( task_instance=ti, mark_success=False, pickle_id=None, ignore_all_deps=False, ignore_depends_on_past=False, ignore_task_deps=False, ignore_ti_state=False, pool=None, cfg_path=None, ) kwarg_values = kwargs.values() cke.queue_task_instance(**kwargs) if test_queue == KUBERNETES_QUEUE: k8s_executor_mock.queue_task_instance.assert_called_once_with(*kwarg_values) celery_executor_mock.queue_task_instance.assert_not_called() else: celery_executor_mock.queue_task_instance.assert_called_once_with(*kwarg_values) k8s_executor_mock.queue_task_instance.assert_not_called()
def test_start(self): celery_executor_mock = mock.MagicMock() k8s_executor_mock = mock.MagicMock() cke = CeleryKubernetesExecutor(celery_executor_mock, k8s_executor_mock) cke.start() celery_executor_mock.start.assert_called() k8s_executor_mock.start.assert_called()
def test_terminate(self): celery_executor_mock = mock.MagicMock() k8s_executor_mock = mock.MagicMock() cke = CeleryKubernetesExecutor(celery_executor_mock, k8s_executor_mock) cke.terminate() celery_executor_mock.terminate.assert_called_once() k8s_executor_mock.terminate.assert_called_once()
def test_end(self): celery_executor_mock = mock.MagicMock() k8s_executor_mock = mock.MagicMock() cke = CeleryKubernetesExecutor(celery_executor_mock, k8s_executor_mock) cke.end() celery_executor_mock.end.assert_called_once() k8s_executor_mock.end.assert_called_once()
def when_ti_in_celery_executor(): celery_executor_mock = mock.MagicMock() k8s_executor_mock = mock.MagicMock() cke = CeleryKubernetesExecutor(celery_executor_mock, k8s_executor_mock) celery_executor_mock.has_task.return_value = True assert cke.has_task(ti) celery_executor_mock.has_task.assert_called_once_with(ti)
def test_has_tasks(self, celery_has, k8s_has, cke_has): celery_executor_mock = mock.MagicMock() k8s_executor_mock = mock.MagicMock() cke = CeleryKubernetesExecutor(celery_executor_mock, k8s_executor_mock) celery_executor_mock.has_task.return_value = celery_has k8s_executor_mock.has_task.return_value = k8s_has ti = mock.MagicMock() assert cke.has_task(ti) == cke_has celery_executor_mock.has_task.assert_called_once_with(ti) if not celery_has: k8s_executor_mock.has_task.assert_called_once_with(ti)
def when_using_celery_executor(): celery_executor_mock = mock.MagicMock() k8s_executor_mock = mock.MagicMock() cke = CeleryKubernetesExecutor(celery_executor_mock, k8s_executor_mock) simple_task_instance = mock.MagicMock() simple_task_instance.queue = 'non-kubernetes-queue' cke.queue_command(simple_task_instance, command, priority, queue) celery_executor_mock.queue_command.assert_called_once_with( simple_task_instance, command, priority, queue) k8s_executor_mock.queue_command.assert_not_called()
def when_ti_in_celery_executor(): celery_executor_mock = mock.MagicMock() k8s_executor_mock = mock.MagicMock() ti.queue = "default" cke = CeleryKubernetesExecutor(celery_executor_mock, k8s_executor_mock) celery_executor_mock.try_adopt_task_instances.return_value = [] k8s_executor_mock.try_adopt_task_instances.return_value = [] cke.try_adopt_task_instances([ti]) celery_executor_mock.try_adopt_task_instances.assert_called_once_with( [ti]) k8s_executor_mock.try_adopt_task_instances.assert_called_once_with( [])
def test_adopt_tasks(self, num_k8s, num_celery): celery_executor_mock = mock.MagicMock() k8s_executor_mock = mock.MagicMock() def mock_ti(queue): ti = mock.MagicMock() ti.queue = queue return ti celery_tis = [mock_ti('default') for _ in range(num_celery)] k8s_tis = [mock_ti(KUBERNETES_QUEUE) for _ in range(num_k8s)] cke = CeleryKubernetesExecutor(celery_executor_mock, k8s_executor_mock) cke.try_adopt_task_instances(celery_tis + k8s_tis) celery_executor_mock.try_adopt_task_instances.assert_called_once_with(celery_tis) k8s_executor_mock.try_adopt_task_instances.assert_called_once_with(k8s_tis)
def test_get_event_buffer(self): celery_executor_mock = mock.MagicMock() k8s_executor_mock = mock.MagicMock() cke = CeleryKubernetesExecutor(celery_executor_mock, k8s_executor_mock) dag_ids = ['dag_ids'] events_in_celery = {('dag_id', 'task_id', '2020-08-30', 1): ('failed', 'failed task')} events_in_k8s = {('dag_id_2', 'task_id_2', '2020-08-30', 1): ('success', None)} celery_executor_mock.get_event_buffer.return_value = events_in_celery k8s_executor_mock.get_event_buffer.return_value = events_in_k8s events = cke.get_event_buffer(dag_ids) assert events == {**events_in_celery, **events_in_k8s} celery_executor_mock.get_event_buffer.assert_called_once_with(dag_ids) k8s_executor_mock.get_event_buffer.assert_called_once_with(dag_ids)
def test_queue_command(self, test_queue, k8s_queue_cmd, celery_queue_cmd): kwargs = dict( command=['airflow', 'run', 'dag'], priority=1, queue='default', ) kwarg_values = kwargs.values() cke = CeleryKubernetesExecutor(CeleryExecutor(), KubernetesExecutor()) simple_task_instance = mock.MagicMock() simple_task_instance.queue = test_queue cke.queue_command(simple_task_instance, **kwargs) if test_queue == KUBERNETES_QUEUE: k8s_queue_cmd.assert_called_once_with(simple_task_instance, *kwarg_values) celery_queue_cmd.assert_not_called() else: celery_queue_cmd.assert_called_once_with(simple_task_instance, *kwarg_values) k8s_queue_cmd.assert_not_called()
def test_running(self): celery_executor_mock = mock.MagicMock() k8s_executor_mock = mock.MagicMock() cke = CeleryKubernetesExecutor(celery_executor_mock, k8s_executor_mock) celery_running_tasks = {('dag_id', 'task_id', '2020-08-30', 1)} k8s_running_tasks = {('dag_id_2', 'task_id_2', '2020-08-30', 2)} celery_executor_mock.running = celery_running_tasks k8s_executor_mock.running = k8s_running_tasks assert cke.running == celery_running_tasks.union(k8s_running_tasks)
def test_queued_tasks(self): celery_executor_mock = mock.MagicMock() k8s_executor_mock = mock.MagicMock() cke = CeleryKubernetesExecutor(celery_executor_mock, k8s_executor_mock) celery_queued_tasks = {('dag_id', 'task_id', '2020-08-30', 1): 'queued_command'} k8s_queued_tasks = {('dag_id_2', 'task_id_2', '2020-08-30', 2): 'queued_command'} celery_executor_mock.queued_tasks = celery_queued_tasks k8s_executor_mock.queued_tasks = k8s_queued_tasks expected_queued_tasks = {**celery_queued_tasks, **k8s_queued_tasks} assert cke.queued_tasks == expected_queued_tasks