From 4b0901706884561e83cc4de08c8e0f4cbd688f19 Mon Sep 17 00:00:00 2001 From: Victor Chiapaikeo Date: Sat, 17 Feb 2024 08:52:30 -0500 Subject: [PATCH] Fix KeyError when KPO exits too soon --- airflow/providers/cncf/kubernetes/triggers/pod.py | 7 ++++++- tests/providers/cncf/kubernetes/triggers/test_pod.py | 9 ++++++++- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/airflow/providers/cncf/kubernetes/triggers/pod.py b/airflow/providers/cncf/kubernetes/triggers/pod.py index c9b1e62226541..8c411e4e4a345 100644 --- a/airflow/providers/cncf/kubernetes/triggers/pod.py +++ b/airflow/providers/cncf/kubernetes/triggers/pod.py @@ -245,7 +245,12 @@ async def _wait_for_container_completion(self) -> TriggerEvent: ) elif container_state == ContainerState.FAILED: return TriggerEvent( - {"status": "failed", "namespace": self.pod_namespace, "name": self.pod_name} + { + "status": "failed", + "namespace": self.pod_namespace, + "name": self.pod_name, + "message": "Container state failed", + } ) if time_get_more_logs and datetime.datetime.now(tz=datetime.timezone.utc) > time_get_more_logs: return TriggerEvent( diff --git a/tests/providers/cncf/kubernetes/triggers/test_pod.py b/tests/providers/cncf/kubernetes/triggers/test_pod.py index 228daa967da18..ef06ff3c0abe4 100644 --- a/tests/providers/cncf/kubernetes/triggers/test_pod.py +++ b/tests/providers/cncf/kubernetes/triggers/test_pod.py @@ -182,7 +182,14 @@ async def test_run_loop_return_failed_event(self, mock_hook, mock_method, mock_w ) mock_method.return_value = ContainerState.FAILED - expected_event = TriggerEvent({"status": "failed", "namespace": "default", "name": "test-pod-name"}) + expected_event = TriggerEvent( + { + "status": "failed", + "namespace": "default", + "name": "test-pod-name", + "message": "Container state failed", + } + ) actual_event = await trigger.run().asend(None) assert actual_event == expected_event