From a3f2fb0477995ff2c360443b91ff20057fc53ad0 Mon Sep 17 00:00:00 2001 From: insomnes Date: Wed, 26 Mar 2025 00:25:01 +0100 Subject: [PATCH] test: k8s system tests speed-up active_deadline_seconds test - Refactor for fail / success scenarios - Lower timer and remove second timer for fail scenario --- .../test_kubernetes_pod_operator.py | 28 +++++++++++++++---- 1 file changed, 22 insertions(+), 6 deletions(-) diff --git a/kubernetes_tests/test_kubernetes_pod_operator.py b/kubernetes_tests/test_kubernetes_pod_operator.py index e5637b5a9db17..c334e858124e7 100644 --- a/kubernetes_tests/test_kubernetes_pod_operator.py +++ b/kubernetes_tests/test_kubernetes_pod_operator.py @@ -20,6 +20,7 @@ import logging import os import shutil +from contextlib import nullcontext from copy import copy from unittest import mock from unittest.mock import ANY, MagicMock @@ -1455,21 +1456,29 @@ def __getattr__(self, name): class TestKubernetesPodOperator(BaseK8STest): - @pytest.mark.parametrize("active_deadline_seconds", [10, 20]) - def test_kubernetes_pod_operator_active_deadline_seconds(self, active_deadline_seconds): + @pytest.mark.parametrize( + "active_deadline_seconds,should_fail", + [(3, True), (60, False)], + ids=["should_fail", "should_not_fail"], + ) + def test_kubernetes_pod_operator_active_deadline_seconds(self, active_deadline_seconds, should_fail): ns = "default" + echo = "echo 'hello world'" + if should_fail: + echo += " && sleep 10" k = KubernetesPodOperator( task_id=f"test_task_{active_deadline_seconds}", active_deadline_seconds=active_deadline_seconds, image="busybox", - cmds=["sh", "-c", "echo 'hello world' && sleep 60"], + cmds=["sh", "-c", echo], namespace=ns, on_finish_action="keep_pod", ) context = create_context(k) - with pytest.raises(AirflowException): + ctx_manager = pytest.raises(AirflowException) if should_fail else nullcontext() + with ctx_manager: k.execute(context) pod = k.find_pod(ns, context, exclude_checked=False) @@ -1480,5 +1489,12 @@ def test_kubernetes_pod_operator_active_deadline_seconds(self, active_deadline_s phase = pod_status.status.phase reason = pod_status.status.reason - assert phase == "Failed" - assert reason == "DeadlineExceeded" + if should_fail: + expected_phase = "Failed" + expected_reason = "DeadlineExceeded" + else: + expected_phase = "Succeeded" + expected_reason = None + + assert phase == expected_phase + assert reason == expected_reason