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 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()