fix(cncf): align how KubernetesPodOperatorAsync handle event and WaitContainerTrigger

Author: Lee-W

astronomer/providers/cncf/kubernetes/operators/kubernetes_pod.py

@@ -5,9 +5,9 @@
 from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import (
     KubernetesPodOperator,
 )
-from airflow.providers.cncf.kubernetes.triggers.kubernetes_pod import (
-    ContainerState,
-    KubernetesPodTrigger,
+from airflow.providers.cncf.kubernetes.utils.pod_manager import (
+    PodPhase,
+    container_is_running,
 )
 from kubernetes.client import models as k8s
 from pendulum import DateTime
@@ -54,21 +54,6 @@ def raise_for_trigger_status(event: Dict[str, Any]) -> None:
             else:
                 raise AirflowException(description)
 
-    def define_container_state(self) -> ContainerState:
-        """Define the state of container"""
-        pod_containers = self.pod.status.container_statuses
-        if not pod_containers:
-            return ContainerState.UNDEFINED
-        container = [c for c in pod_containers if c.name == self.BASE_CONTAINER_NAME][0]
-        for state in (ContainerState.RUNNING, ContainerState.WAITING, ContainerState.TERMINATED):
-            state_obj = getattr(container.state, state)
-            if state_obj:
-                if state != ContainerState.TERMINATED:
-                    return state
-                else:
-                    return ContainerState.TERMINATED if state_obj.exit_code == 0 else ContainerState.FAILED
-        return ContainerState.UNDEFINED
-
     def defer(self, last_log_time: Optional[DateTime] = None, **kwargs: Any) -> None:
         """Defers to ``WaitContainerTrigger`` optionally with last log time."""
         if kwargs:
@@ -98,27 +83,18 @@ def defer(self, last_log_time: Optional[DateTime] = None, **kwargs: Any) -> None
     def execute(self, context: Context) -> Any:  # noqa: D102
         self.pod_request_obj = self.build_pod_request_obj(context)
         self.pod: k8s.V1Pod = self.get_or_create_pod(self.pod_request_obj, context)
-        container_state = self.define_container_state()
         pod_status = self.pod.status.phase
-        if KubernetesPodTrigger.should_wait(pod_status, container_state):
-            self.defer()
-            return
-
-        if container_state == ContainerState.TERMINATED:
-            event = {
-                "name": self.pod.metadata.name,
-                "namespace": self.pod.metadata.namespace,
-                "status": "success",
-                "message": "All containers inside pod have started successfully.",
-            }
-        else:
+        if pod_status in PodPhase.terminal_states or not container_is_running(
+            pod=self.pod, container_name=self.BASE_CONTAINER_NAME
+        ):
             event = {
-                "name": self.pod.metadata.name,
+                "status": "done",
                 "namespace": self.pod.metadata.namespace,
-                "status": "failed",
-                "message": self.pod.status.message,
+                "pod_name": self.pod.metadata.name,
             }
-        return self.trigger_reentry(context=context, event=event)
+            return self.trigger_reentry(context=context, event=event)
+
+        self.defer()
 
     def execute_complete(self, context: Context, event: Dict[str, Any]) -> Any:  # type: ignore[override]
         """Deprecated; replaced by trigger_reentry."""

tests/cncf/kubernetes/operators/test_kubernetes_pod.py

@@ -3,8 +3,10 @@
 
 import pytest
 from airflow.exceptions import TaskDeferred
-from airflow.providers.cncf.kubernetes.triggers.kubernetes_pod import ContainerState
-from airflow.providers.cncf.kubernetes.utils.pod_manager import PodLoggingStatus
+from airflow.providers.cncf.kubernetes.utils.pod_manager import (
+    PodLoggingStatus,
+    PodPhase,
+)
 from kubernetes.client import models as k8s
 
 from astronomer.providers.cncf.kubernetes.operators.kubernetes_pod import (
@@ -172,63 +174,22 @@ def test_defer_with_kwargs(self):
         with pytest.raises(ValueError):
             op.defer(kwargs={"timeout": 10})
 
+    @pytest.mark.parametrize("pod_phase", PodPhase.terminal_states)
     @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.trigger_reentry")
-    @mock.patch(
-        f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.define_container_state",
-        return_value=ContainerState.FAILED,
-    )
     @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.build_pod_request_obj")
     @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.get_or_create_pod")
     @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.defer")
-    def test_execute_failed_before_defer(
-        self,
-        mock_defer,
-        mock_get_or_create_pod,
-        mock_build_pod_request_obj,
-        mock_define_container_state,
-        mock_trigger_reentry,
+    def test_execute_done_before_defer(
+        self, mock_defer, mock_get_or_create_pod, mock_build_pod_request_obj, mock_trigger_reentry, pod_phase
     ):
-        mock_get_or_create_pod.return_value = _build_mock_pod(
-            k8s.V1ContainerState({"running": k8s.V1ContainerStateTerminated(exit_code=1), "waiting": None})
-        )
-        mock_build_pod_request_obj.return_value = {}
-        mock_defer.return_value = {}
-        op = KubernetesPodOperatorAsync(task_id="test_task", name="test-pod", get_logs=True)
-
-        op.execute(context=create_context(op))
-        assert mock_trigger_reentry.called
-        assert not mock_defer.called
-
-    @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.trigger_reentry")
-    @mock.patch(
-        f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.define_container_state",
-        return_value=ContainerState.TERMINATED,
-    )
-    @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.build_pod_request_obj")
-    @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.get_or_create_pod")
-    @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.defer")
-    def test_execute_succeeded_before_defer(
-        self,
-        mock_defer,
-        mock_get_or_create_pod,
-        mock_build_pod_request_obj,
-        mock_define_container_state,
-        mock_trigger_reentry,
-    ):
-        mock_get_or_create_pod.return_value = _build_mock_pod(
-            k8s.V1ContainerState({"running": k8s.V1ContainerStateTerminated(exit_code=0), "waiting": None})
-        )
+        mock_get_or_create_pod.return_value.status.phase = pod_phase
         mock_build_pod_request_obj.return_value = {}
         mock_defer.return_value = {}
         op = KubernetesPodOperatorAsync(task_id="test_task", name="test-pod", get_logs=True)
         assert op.execute(context=create_context(op))
         assert mock_trigger_reentry.called
         assert not mock_defer.called
 
-    @mock.patch(
-        "astronomer.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperatorAsync.define_container_state",
-        return_value=ContainerState.RUNNING,
-    )
     @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.build_pod_request_obj")
     @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.get_or_create_pod")
     @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.defer")
@@ -237,7 +198,6 @@ def test_execute(
         mock_defer,
         mock_get_or_create_pod,
         mock_build_pod_request_obj,
-        mock_define_container_state,
     ):
         """Assert that execute succeeded"""
         mock_get_or_create_pod.return_value = _build_mock_pod(
@@ -256,40 +216,3 @@ def test_execute_complete(self, mock_trigger_reentry):
         mock_trigger_reentry.return_value = {}
         op = KubernetesPodOperatorAsync(task_id="test_task", name="test-pod", get_logs=True)
         assert op.execute_complete(context=create_context(op), event={}) is None
-
-    @pytest.mark.parametrize(
-        "container_state, expected_state",
-        [
-            (
-                {"running": k8s.V1ContainerStateRunning(), "terminated": None, "waiting": None},
-                ContainerState.RUNNING,
-            ),
-            (
-                {"running": None, "terminated": k8s.V1ContainerStateTerminated(exit_code=0), "waiting": None},
-                ContainerState.TERMINATED,
-            ),
-            (
-                {"running": None, "terminated": None, "waiting": k8s.V1ContainerStateWaiting()},
-                ContainerState.WAITING,
-            ),
-        ],
-    )
-    def test_define_container_state_should_execute_successfully(self, container_state, expected_state):
-        op = KubernetesPodOperatorAsync(task_id="test_task", name="test-pod", get_logs=True)
-        op.pod = _build_mock_pod(k8s.V1ContainerState(**container_state))
-        assert expected_state == op.define_container_state()
-
-    @pytest.mark.parametrize(
-        "pod",
-        (
-            _build_mock_pod(k8s.V1ContainerState(running=None, terminated=None, waiting=None)),
-            k8s.V1Pod(
-                metadata=k8s.V1ObjectMeta(name="base", namespace="default"),
-                status=k8s.V1PodStatus(container_statuses=[]),
-            ),
-        ),
-    )
-    def test_define_container_state_with_undefined_state(self, pod):
-        op = KubernetesPodOperatorAsync(task_id="test_task", name="test-pod", get_logs=True)
-        op.pod = pod
-        assert op.define_container_state() == ContainerState.UNDEFINED