[GitHub] [airflow] jedcunningham commented on a diff in pull request #29000: Skip KubernetesPodOperator task when it returns a provided exit code

["jedcunningham (via GitHub)" <gi...@apache.org>]

jedcunningham commented on code in PR #29000:
URL: https://github.com/apache/airflow/pull/29000#discussion_r1126713389


##########
airflow/providers/cncf/kubernetes/operators/pod.py:
##########
@@ -670,6 +675,24 @@ def cleanup(self, pod: k8s.V1Pod, remote_pod: k8s.V1Pod):
 
             error_message = get_container_termination_message(remote_pod, self.base_container_name)
             error_message = "\n" + error_message if error_message else ""
+            if self.skip_exit_code is not None:
+                container_statuses = (
+                    remote_pod.status.container_statuses if remote_pod and remote_pod.status else None
+                )
+                base_container_status = next(
+                    (x for x in container_statuses if x.name == self.base_container_name), None

Review Comment:
   If `container_statuses` is None, won't this complain that None isn't iterable?



##########
tests/providers/cncf/kubernetes/operators/test_pod.py:
##########
@@ -1217,6 +1246,70 @@ def test_async_create_pod_should_throw_exception(self, mocked_hook, mocked_clean
                 },
             )
 
+    @pytest.mark.parametrize(
+        "extra_kwargs, actual_exit_code, expected_exc, pod_status, event_status",
+        [
+            (None, 0, None, "Succeeded", "success"),
+            (None, 99, AirflowException, "Failed", "error"),
+            ({"skip_exit_code": 100}, 100, AirflowSkipException, "Failed", "error"),
+            ({"skip_exit_code": 100}, 101, AirflowException, "Failed", "error"),
+            ({"skip_exit_code": None}, 100, AirflowException, "Failed", "error"),
+        ],
+    )
+    @patch(HOOK_CLASS)
+    def test_async_create_pod_with_skip_exit_code_should_skip(
+        self,
+        mocked_hook,
+        extra_kwargs,
+        actual_exit_code,
+        expected_exc,
+        pod_status,
+        event_status,
+    ):
+        """Tests that an AirflowSkipException is raised when the pod exit with the skip_exit_code"""

Review Comment:
   ```suggestion
           """Tests that an AirflowSkipException is raised when the container exits with the skip_exit_code"""
   ```
   
   nit



##########
tests/providers/cncf/kubernetes/operators/test_pod.py:
##########
@@ -1052,6 +1052,35 @@ def test_task_id_as_name_dag_id_is_ignored(self):
         pod = k.build_pod_request_obj({})
         assert re.match(r"a-very-reasonable-task-name-[a-z0-9-]+", pod.metadata.name) is not None
 
+    @pytest.mark.parametrize(
+        "extra_kwargs, actual_exit_code, expected_exc",
+        [
+            (None, 99, AirflowException),
+            ({"skip_exit_code": 100}, 100, AirflowSkipException),
+            ({"skip_exit_code": 100}, 101, AirflowException),
+            ({"skip_exit_code": None}, 100, AirflowException),
+        ],
+    )
+    @patch(f"{POD_MANAGER_CLASS}.await_pod_completion")
+    def test_task_skip_when_pod_exit_with_certain_code(
+        self, remote_pod, extra_kwargs, actual_exit_code, expected_exc
+    ):
+        """ """

Review Comment:
   Empty doc string



-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org