Skip to content

Commit fcd607a

Browse files
author
Sameer Mesiah
committed
Add dedicated unit tests for the Kubernetes scheduler pod creation
and pod patching metrics, covering both success and failure paths. Extend the existing task adoption test to verify that the kubernetes_executor.adopt_task_instances.duration metric is emitted with the expected team_name tag when configured.
1 parent 8ef00a0 commit fcd607a

1 file changed

Lines changed: 280 additions & 1 deletion

File tree

‎providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py‎

Lines changed: 280 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -246,6 +246,141 @@ def test_delete_pod_404_not_raised(self, mock_watcher, mock_client, mock_kube_cl
246246
finally:
247247
kube_executor.end()
248248

249+
@pytest.mark.skipif(
250+
AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
251+
)
252+
@pytest.mark.parametrize(
253+
("team_name", "timer_tags", "status_tags"),
254+
[
255+
pytest.param(
256+
None,
257+
{},
258+
{"status": "200"},
259+
id="without_team",
260+
),
261+
pytest.param(
262+
"team_a",
263+
{"team_name": "team_a"},
264+
{"status": "200", "team_name": "team_a"},
265+
id="with_team",
266+
marks=pytest.mark.skipif(
267+
not AIRFLOW_V_3_1_PLUS,
268+
reason="team_name metrics require Airflow 3.1+",
269+
),
270+
),
271+
],
272+
)
273+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.Stats")
274+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
275+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.client")
276+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
277+
def test_run_pod_async_emits_metrics_on_success(
278+
self,
279+
mock_watcher,
280+
mock_client,
281+
mock_kube_client,
282+
mock_stats,
283+
team_name,
284+
timer_tags,
285+
status_tags,
286+
):
287+
pod = mock.MagicMock()
288+
pod.metadata.namespace = "default"
289+
290+
mock_api_client = mock.Mock()
291+
mock_api_client.sanitize_for_serialization.return_value = {}
292+
293+
mock_kube_client.return_value.api_client = mock_api_client
294+
mock_kube_client.return_value.create_namespaced_pod = mock.MagicMock()
295+
296+
kube_executor = KubernetesExecutor()
297+
kube_executor.team_name = team_name
298+
kube_executor.job_id = 1
299+
kube_executor.start()
300+
301+
try:
302+
kube_executor.kube_scheduler.run_pod_async(pod)
303+
304+
mock_stats.timer.assert_any_call(
305+
"kubernetes_executor.pod_creation",
306+
tags=timer_tags,
307+
)
308+
309+
mock_stats.incr.assert_any_call(
310+
"kubernetes_executor.pod_creation_status",
311+
tags=status_tags,
312+
)
313+
finally:
314+
kube_executor.end()
315+
316+
@pytest.mark.skipif(
317+
AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
318+
)
319+
@pytest.mark.parametrize(
320+
("team_name", "timer_tags", "status_tags"),
321+
[
322+
pytest.param(
323+
None,
324+
{},
325+
{"status": "429"},
326+
id="without_team",
327+
),
328+
pytest.param(
329+
"team_a",
330+
{"team_name": "team_a"},
331+
{"status": "429", "team_name": "team_a"},
332+
id="with_team",
333+
marks=pytest.mark.skipif(
334+
not AIRFLOW_V_3_1_PLUS,
335+
reason="team_name metrics require Airflow 3.1+",
336+
),
337+
),
338+
],
339+
)
340+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.Stats")
341+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
342+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.client")
343+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
344+
def test_run_pod_async_emits_metrics_on_failure(
345+
self,
346+
mock_watcher,
347+
mock_client,
348+
mock_kube_client,
349+
mock_stats,
350+
team_name,
351+
timer_tags,
352+
status_tags,
353+
):
354+
pod = mock.MagicMock()
355+
pod.metadata.namespace = "default"
356+
357+
mock_api_client = mock.Mock()
358+
mock_api_client.sanitize_for_serialization.return_value = {}
359+
360+
mock_kube_client.return_value.api_client = mock_api_client
361+
mock_kube_client.return_value.create_namespaced_pod.side_effect = ApiException(status=429)
362+
363+
kube_executor = KubernetesExecutor()
364+
kube_executor.team_name = team_name
365+
kube_executor.job_id = 1
366+
kube_executor.start()
367+
368+
try:
369+
with pytest.raises(ApiException):
370+
kube_executor.kube_scheduler.run_pod_async(pod)
371+
372+
mock_stats.timer.assert_any_call(
373+
"kubernetes_executor.pod_creation",
374+
tags=timer_tags,
375+
)
376+
377+
mock_stats.incr.assert_any_call(
378+
"kubernetes_executor.pod_creation_status",
379+
tags=status_tags,
380+
)
381+
finally:
382+
kube_executor.end()
383+
249384
@pytest.mark.skipif(
250385
AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
251386
)
@@ -370,6 +505,122 @@ def test_delete_pod_emits_metrics_on_failure(
370505
finally:
371506
kube_executor.end()
372507

508+
@pytest.mark.skipif(
509+
AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
510+
)
511+
@pytest.mark.parametrize(
512+
("team_name", "timer_tags", "status_tags"),
513+
[
514+
pytest.param(None, {}, {"status": "200"}, id="without_team"),
515+
pytest.param(
516+
"team_a",
517+
{"team_name": "team_a"},
518+
{"status": "200", "team_name": "team_a"},
519+
id="with_team",
520+
marks=pytest.mark.skipif(
521+
not AIRFLOW_V_3_1_PLUS,
522+
reason="team_name metrics require Airflow 3.1+",
523+
),
524+
),
525+
],
526+
)
527+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.Stats")
528+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
529+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.client")
530+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
531+
def test_patch_pod_emits_metrics_on_success(
532+
self,
533+
mock_watcher,
534+
mock_client,
535+
mock_kube_client,
536+
mock_stats,
537+
team_name,
538+
timer_tags,
539+
status_tags,
540+
):
541+
mock_kube_client.return_value.patch_namespaced_pod = mock.MagicMock()
542+
543+
kube_executor = KubernetesExecutor()
544+
kube_executor.team_name = team_name
545+
kube_executor.job_id = 1
546+
kube_executor.start()
547+
548+
try:
549+
kube_executor.kube_scheduler.patch_pod_executor_done(
550+
pod_name="pod",
551+
namespace="default",
552+
)
553+
554+
mock_stats.timer.assert_any_call(
555+
"kubernetes_executor.pod_patching",
556+
tags=timer_tags,
557+
)
558+
559+
mock_stats.incr.assert_any_call(
560+
"kubernetes_executor.pod_patching_status",
561+
tags=status_tags,
562+
)
563+
finally:
564+
kube_executor.end()
565+
566+
@pytest.mark.skipif(
567+
AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
568+
)
569+
@pytest.mark.parametrize(
570+
("team_name", "timer_tags", "status_tags"),
571+
[
572+
pytest.param(None, {}, {"status": "429"}, id="without_team"),
573+
pytest.param(
574+
"team_a",
575+
{"team_name": "team_a"},
576+
{"status": "429", "team_name": "team_a"},
577+
id="with_team",
578+
marks=pytest.mark.skipif(
579+
not AIRFLOW_V_3_1_PLUS,
580+
reason="team_name metrics require Airflow 3.1+",
581+
),
582+
),
583+
],
584+
)
585+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.Stats")
586+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
587+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.client")
588+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
589+
def test_patch_pod_emits_metrics_on_failure(
590+
self,
591+
mock_watcher,
592+
mock_client,
593+
mock_kube_client,
594+
mock_stats,
595+
team_name,
596+
timer_tags,
597+
status_tags,
598+
):
599+
mock_kube_client.return_value.patch_namespaced_pod.side_effect = ApiException(status=429)
600+
601+
kube_executor = KubernetesExecutor()
602+
kube_executor.team_name = team_name
603+
kube_executor.job_id = 1
604+
kube_executor.start()
605+
606+
try:
607+
kube_executor.kube_scheduler.patch_pod_executor_done(
608+
pod_name="pod",
609+
namespace="default",
610+
)
611+
612+
mock_stats.timer.assert_any_call(
613+
"kubernetes_executor.pod_patching",
614+
tags=timer_tags,
615+
)
616+
617+
mock_stats.incr.assert_any_call(
618+
"kubernetes_executor.pod_patching_status",
619+
tags=status_tags,
620+
)
621+
finally:
622+
kube_executor.end()
623+
373624
def test_running_pod_log_lines(self):
374625
# default behaviour
375626
kube_executor = KubernetesExecutor()
@@ -1205,6 +1456,22 @@ def test_change_state_failed_pod_deletion(
12051456
finally:
12061457
executor.end()
12071458

1459+
@pytest.mark.parametrize(
1460+
("team_name", "timer_tags"),
1461+
[
1462+
pytest.param(None, {}, id="without_team"),
1463+
pytest.param(
1464+
"team_a",
1465+
{"team_name": "team_a"},
1466+
id="with_team",
1467+
marks=pytest.mark.skipif(
1468+
not AIRFLOW_V_3_1_PLUS,
1469+
reason="team_name metrics require Airflow 3.1+",
1470+
),
1471+
),
1472+
],
1473+
)
1474+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor.Stats.timer")
12081475
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor.DynamicClient")
12091476
@mock.patch(
12101477
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor.KubernetesExecutor.adopt_launched_task"
@@ -1213,9 +1480,16 @@ def test_change_state_failed_pod_deletion(
12131480
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor.KubernetesExecutor._adopt_completed_pods"
12141481
)
12151482
def test_try_adopt_task_instances(
1216-
self, mock_adopt_completed_pods, mock_adopt_launched_task, mock_kube_dynamic_client
1483+
self,
1484+
mock_adopt_completed_pods,
1485+
mock_adopt_launched_task,
1486+
mock_kube_dynamic_client,
1487+
mock_stats_timer,
1488+
team_name,
1489+
timer_tags,
12171490
):
12181491
executor = self.kubernetes_executor
1492+
executor.team_name = team_name
12191493
executor.scheduler_job_id = "10"
12201494
ti_key = annotations_to_key(
12211495
{
@@ -1271,6 +1545,11 @@ def test_try_adopt_task_instances(
12711545
mock_adopt_completed_pods.assert_called_once()
12721546
assert reset_tis == [] # This time our return is empty - no TIs to reset
12731547

1548+
mock_stats_timer.assert_any_call(
1549+
"kubernetes_executor.adopt_task_instances.duration",
1550+
tags=timer_tags,
1551+
)
1552+
12741553
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor.DynamicClient")
12751554
@mock.patch(
12761555
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor.KubernetesExecutor._adopt_completed_pods"

0 commit comments

Comments
 (0)