XD-DENG commented on code in PR #27611:
URL: https://github.com/apache/airflow/pull/27611#discussion_r1019753709
##########
tests/executors/test_kubernetes_executor.py:
##########
@@ -293,6 +293,41 @@ def test_run_next_exception_requeue(
assert kubernetes_executor.task_queue.empty()
assert kubernetes_executor.event_buffer[task_instance_key][0]
== State.FAILED
+ @pytest.mark.skipif(
+ AirflowKubernetesScheduler is None, reason="kubernetes python package
is not installed"
+ )
+ @mock.patch("airflow.executors.kubernetes_executor.pod_mutation_hook")
+ @mock.patch("airflow.executors.kubernetes_executor.get_kube_client")
+ def test_run_next_pmh_error_requeue(self, mock_get_kube_client, mock_pmh):
+ """
+ Exception during Pod Mutation Hook execution should be handled
gracefully.
+ The job should also be re-queued in case the Pod Mutation Hook failed
due to transient reason.
+
+ """
+ mock_pmh.side_effect = Exception("Purposely generate error for test")
+
+ mock_kube_client = mock.patch("kubernetes.client.CoreV1Api",
autospec=True)
+ mock_kube_client.create_namespaced_pod = mock.MagicMock()
+ mock_get_kube_client.return_value = mock_kube_client
+
+ kubernetes_executor = self.kubernetes_executor
+ kubernetes_executor.start()
+ try_number = 1
+ task_instance_key = TaskInstanceKey("dag", "task", "run_id",
try_number)
+ kubernetes_executor.execute_async(
+ key=task_instance_key,
+ queue=None,
+ command=["airflow", "tasks", "run", "true", "some_parameter"],
+ )
+ kubernetes_executor.sync()
+
+ # The pod_mutation_hook should have been called once.
+ assert mock_pmh.call_count == 1
+ # There should be no pod creation request sent
+ assert mock_kube_client.create_namespaced_pod.call_count == 0
+ # The task is not re-queued
+ assert kubernetes_executor.task_queue.empty()
Review Comment:
Thanks @jedcunningham .
Added more comprehensive test updates at
https://github.com/apache/airflow/pull/27611/commits/0b86fb3693540e3ee3fba434ceaaf9c8934cb85c.
Will merge after the CI passed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
