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