def test_pending_pod_timeout_multi_namespace_mode( self, mock_kubescheduler, mock_get_kube_client, mock_kubernetes_job_watcher): mock_delete_pod = mock_kubescheduler.return_value.delete_pod mock_kube_client = mock_get_kube_client.return_value now = timezone.utcnow() pending_pods = [ k8s.V1Pod(metadata=k8s.V1ObjectMeta( name="foo90", labels={"airflow-worker": "123"}, creation_timestamp=now - timedelta(seconds=500), namespace="anothernamespace", )), ] mock_kube_client.list_pod_for_all_namespaces.return_value.items = pending_pods config = { ('kubernetes', 'namespace'): 'mynamespace', ('kubernetes', 'multi_namespace_mode'): 'true', ('kubernetes', 'kube_client_request_args'): '{"sentinel": "foo"}', } with conf_vars(config): executor = KubernetesExecutor() executor.job_id = "123" executor.start() executor._check_worker_pods_pending_timeout() mock_kube_client.list_pod_for_all_namespaces.assert_called_once_with( field_selector='status.phase=Pending', label_selector='airflow-worker=123', limit=100, sentinel='foo', ) mock_delete_pod.assert_called_once_with('foo90', 'anothernamespace')
def test_change_state_running(self, mock_get_kube_client, mock_kubernetes_job_watcher): executor = KubernetesExecutor() executor.start() key = ('dag_id', 'task_id', 'ex_time', 'try_number1') executor._change_state(key, State.RUNNING, 'pod_id') self.assertTrue(executor.event_buffer[key] == State.RUNNING)
def test_run_next_exception(self, mock_get_kube_client, mock_kubernetes_job_watcher): # When a quota is exceeded this is the ApiException we get r = HTTPResponse() r.body = { "kind": "Status", "apiVersion": "v1", "metadata": {}, "status": "Failure", "message": "pods \"podname\" is forbidden: " + "exceeded quota: compute-resources, " + "requested: limits.memory=4Gi, " + "used: limits.memory=6508Mi, " + "limited: limits.memory=10Gi", "reason": "Forbidden", "details": { "name": "podname", "kind": "pods" }, "code": 403 }, r.status = 403 r.reason = "Forbidden" # A mock kube_client that throws errors when making a pod mock_kube_client = mock.patch('kubernetes.client.CoreV1Api', autospec=True) mock_kube_client.create_namespaced_pod = mock.MagicMock( side_effect=ApiException(http_resp=r)) mock_get_kube_client.return_value = mock_kube_client kubernetesExecutor = KubernetesExecutor() kubernetesExecutor.start() # Execute a task while the Api Throws errors try_number = 1 kubernetesExecutor.execute_async(key=('dag', 'task', datetime.utcnow(), try_number), command='command', executor_config={}) kubernetesExecutor.sync() kubernetesExecutor.sync() assert mock_kube_client.create_namespaced_pod.called self.assertFalse(kubernetesExecutor.task_queue.empty()) # Disable the ApiException mock_kube_client.create_namespaced_pod.side_effect = None # Execute the task without errors should empty the queue kubernetesExecutor.sync() assert mock_kube_client.create_namespaced_pod.called self.assertTrue(kubernetesExecutor.task_queue.empty())
def test_change_state_failed(self, mock_delete_pod, mock_get_kube_client, mock_kubernetes_job_watcher): executor = KubernetesExecutor() executor.start() key = ('dag_id', 'task_id', 'ex_time', 'try_number3') executor._change_state(key, State.FAILED, 'pod_id') self.assertTrue(executor.event_buffer[key] == State.FAILED) mock_delete_pod.assert_called_once_with('pod_id')
def test_change_state_success(self, mock_delete_pod, mock_get_kube_client, mock_kubernetes_job_watcher): executor = KubernetesExecutor() executor.start() test_time = timezone.utcnow() key = ('dag_id', 'task_id', test_time, 'try_number2') executor._change_state(key, State.SUCCESS, 'pod_id', 'default') self.assertTrue(executor.event_buffer[key] == State.SUCCESS) mock_delete_pod.assert_called_once_with('pod_id', 'default')
def test_change_state_success(self, mock_delete_pod, mock_get_kube_client, mock_kubernetes_job_watcher, mock_kube_config): executor = KubernetesExecutor() executor.start() key = ('dag_id', 'task_id', 'ex_time', 'try_number2') executor._change_state(key, State.SUCCESS, 'pod_id') self.assertTrue(executor.event_buffer[key] == State.SUCCESS) mock_delete_pod.assert_called_with('pod_id')
def test_change_state_skip_pod_deletion(self, mock_delete_pod, mock_get_kube_client, mock_kubernetes_job_watcher): executor = KubernetesExecutor() executor.kube_config.delete_worker_pods = False executor.start() key = ('dag_id', 'task_id', 'ex_time', 'try_number2') executor._change_state(key, State.SUCCESS, 'pod_id') self.assertTrue(executor.event_buffer[key] == State.SUCCESS) mock_delete_pod.assert_not_called()
def test_change_state_failed_pod_deletion(self, mock_delete_pod, mock_get_kube_client, mock_kubernetes_job_watcher): executor = KubernetesExecutor() executor.kube_config.delete_worker_pods_on_failure = True executor.start() key = ('dag_id', 'task_id', 'ex_time', 'try_number2') executor._change_state(key, State.FAILED, 'pod_id', 'test-namespace') self.assertTrue(executor.event_buffer[key] == State.FAILED) mock_delete_pod.assert_called_once_with('pod_id', 'test-namespace')
def test_change_state_failed(self, mock_delete_pod, mock_get_kube_client, mock_kubernetes_job_watcher): executor = KubernetesExecutor() executor.kube_config.delete_worker_pods = False executor.kube_config.delete_worker_pods_on_failure = False executor.start() test_time = timezone.utcnow() key = ('dag_id', 'task_id', test_time, 'try_number3') executor._change_state(key, State.FAILED, 'pod_id', 'default') self.assertTrue(executor.event_buffer[key] == State.FAILED) mock_delete_pod.assert_not_called()
def test_run_next_exception(self, mock_get_kube_client, mock_kubernetes_job_watcher): import sys path = sys.path[ 0] + '/tests/kubernetes/pod_generator_base_with_secrets.yaml' # When a quota is exceeded this is the ApiException we get response = HTTPResponse( body= '{"kind": "Status", "apiVersion": "v1", "metadata": {}, "status": "Failure", ' '"message": "pods \\"podname\\" is forbidden: exceeded quota: compute-resources, ' 'requested: limits.memory=4Gi, used: limits.memory=6508Mi, limited: limits.memory=10Gi", ' '"reason": "Forbidden", "details": {"name": "podname", "kind": "pods"}, "code": 403}' ) response.status = 403 response.reason = "Forbidden" # A mock kube_client that throws errors when making a pod mock_kube_client = mock.patch('kubernetes.client.CoreV1Api', autospec=True) mock_kube_client.create_namespaced_pod = mock.MagicMock( side_effect=ApiException(http_resp=response)) mock_get_kube_client.return_value = mock_kube_client mock_api_client = mock.MagicMock() mock_api_client.sanitize_for_serialization.return_value = {} mock_kube_client.api_client = mock_api_client config = { ('kubernetes', 'pod_template_file'): path, } with conf_vars(config): kubernetes_executor = KubernetesExecutor() kubernetes_executor.start() # Execute a task while the Api Throws errors try_number = 1 kubernetes_executor.execute_async( key=('dag', 'task', datetime.utcnow(), try_number), queue=None, command=['airflow', 'tasks', 'run', 'true', 'some_parameter'], ) kubernetes_executor.sync() kubernetes_executor.sync() assert mock_kube_client.create_namespaced_pod.called self.assertFalse(kubernetes_executor.task_queue.empty()) # Disable the ApiException mock_kube_client.create_namespaced_pod.side_effect = None # Execute the task without errors should empty the queue kubernetes_executor.sync() assert mock_kube_client.create_namespaced_pod.called self.assertTrue(kubernetes_executor.task_queue.empty())
def test_delete_pod_successfully(self, mock_watcher, mock_client, mock_kube_client): # pylint: disable=unused-argument pod_id = "my-pod-1" namespace = "my-namespace-1" mock_delete_namespace = mock.MagicMock() mock_kube_client.return_value.delete_namespaced_pod = mock_delete_namespace kube_executor = KubernetesExecutor() kube_executor.job_id = "test-job-id" kube_executor.start() kube_executor.kube_scheduler.delete_pod(pod_id, namespace) mock_delete_namespace.assert_called_with( pod_id, namespace, body=mock_client.V1DeleteOptions())
def test_delete_pod_404_not_raised(self, mock_watcher, mock_client, mock_kube_client): # pylint: disable=unused-argument pod_id = "my-pod-1" namespace = "my-namespace-3" mock_delete_namespace = mock.MagicMock() mock_kube_client.return_value.delete_namespaced_pod = mock_delete_namespace # ApiException not raised because the status is 404 mock_kube_client.return_value.delete_namespaced_pod.side_effect = ApiException( status=404) kube_executor = KubernetesExecutor() kube_executor.job_id = "test-job-id" kube_executor.start() kube_executor.kube_scheduler.delete_pod(pod_id, namespace) mock_delete_namespace.assert_called_with( pod_id, namespace, body=mock_client.V1DeleteOptions())