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)