diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py index 3f3438b1b7eea..3016bd5ef8cd9 100644 --- a/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/airflow/providers/cncf/kubernetes/operators/pod.py @@ -627,6 +627,10 @@ def execute_async(self, context: Context): pod_request_obj=self.pod_request_obj, context=context, ) + ti = context["ti"] + ti.xcom_push(key="pod_name", value=self.pod.metadata.name) + ti.xcom_push(key="pod_namespace", value=self.pod.metadata.namespace) + self.invoke_defer_method() def invoke_defer_method(self): @@ -663,10 +667,6 @@ def execute_complete(self, context: Context, event: dict, **kwargs): self.write_logs(pod) raise AirflowException(event["message"]) elif event["status"] == "success": - ti = context["ti"] - ti.xcom_push(key="pod_name", value=pod.metadata.name) - ti.xcom_push(key="pod_namespace", value=pod.metadata.namespace) - # fetch some logs when pod is executed successfully if self.get_logs: self.write_logs(pod) diff --git a/tests/providers/cncf/kubernetes/operators/test_pod.py b/tests/providers/cncf/kubernetes/operators/test_pod.py index 3ada6218a5329..fde4e6015985c 100644 --- a/tests/providers/cncf/kubernetes/operators/test_pod.py +++ b/tests/providers/cncf/kubernetes/operators/test_pod.py @@ -1471,12 +1471,15 @@ def run_pod_async(self, operator: KubernetesPodOperator, map_index: int = -1): ) return remote_pod_mock + @pytest.mark.parametrize("do_xcom_push", [True, False]) @patch(KUB_OP_PATH.format("build_pod_request_obj")) @patch(KUB_OP_PATH.format("get_or_create_pod")) - def test_async_create_pod_should_execute_successfully(self, mocked_pod, mocked_pod_obj): + def test_async_create_pod_should_execute_successfully(self, mocked_pod, mocked_pod_obj, do_xcom_push): """ Asserts that a task is deferred and the KubernetesCreatePodTrigger will be fired when the KubernetesPodOperator is executed in deferrable mode when deferrable=True. + + pod name and namespace are *always* pushed; do_xcom_push only controls xcom sidecar """ k = KubernetesPodOperator( @@ -1491,10 +1494,23 @@ def test_async_create_pod_should_execute_successfully(self, mocked_pod, mocked_p in_cluster=True, get_logs=True, deferrable=True, + do_xcom_push=do_xcom_push, ) k.config_file_in_dict_representation = {"a": "b"} + + mocked_pod.return_value.metadata.name = TEST_NAME + mocked_pod.return_value.metadata.namespace = TEST_NAMESPACE + + context = create_context(k) + ti_mock = MagicMock() + context["ti"] = ti_mock + with pytest.raises(TaskDeferred) as exc: - k.execute(create_context(k)) + k.execute(context) + + assert ti_mock.xcom_push.call_count == 2 + ti_mock.xcom_push.assert_any_call(key="pod_name", value=TEST_NAME) + ti_mock.xcom_push.assert_any_call(key="pod_namespace", value=TEST_NAMESPACE) assert isinstance(exc.value.trigger, KubernetesPodTrigger) @patch(KUB_OP_PATH.format("cleanup")) @@ -1655,34 +1671,6 @@ def test_async_xcom_sidecar_container_resources_default_should_execute_successfu }, ) - @pytest.mark.parametrize("do_xcom_push", [True, False]) - @patch(KUB_OP_PATH.format("post_complete_action")) - @patch(KUB_OP_PATH.format("extract_xcom")) - @patch(POD_MANAGER_CLASS) - @patch(HOOK_CLASS) - def test_async_push_xcom_check_xcom_values_should_execute_successfully( - self, mocked_hook, mock_manager, mock_extract_xcom, post_complete_action, do_xcom_push - ): - """pod name and namespace are *always* pushed; do_xcom_push only controls xcom sidecar""" - - mocked_hook.return_value.get_pod.return_value = k8s.V1Pod( - metadata=k8s.V1ObjectMeta(name=TEST_NAME, namespace=TEST_NAMESPACE) - ) - mock_manager.return_value.await_pod_completion.return_value = {} - mock_extract_xcom.return_value = "{}" - k = KubernetesPodOperator( - task_id="task", - do_xcom_push=do_xcom_push, - deferrable=True, - ) - - pod = self.run_pod_async(k) - - pod_name = XCom.get_one(run_id=self.dag_run.run_id, task_id="task", key="pod_name") - pod_namespace = XCom.get_one(run_id=self.dag_run.run_id, task_id="task", key="pod_namespace") - assert pod_name == pod.metadata.name - assert pod_namespace == pod.metadata.namespace - @pytest.mark.parametrize("get_logs", [True, False]) @patch(KUB_OP_PATH.format("post_complete_action")) @patch(KUB_OP_PATH.format("write_logs")) @@ -1780,8 +1768,6 @@ def test_async_kpo_wait_termination_before_cleanup_on_success( succeeded_state, ] - ti_mock = MagicMock() - success_event = { "status": "success", "message": TEST_SUCCESS_MESSAGE, @@ -1790,16 +1776,11 @@ def test_async_kpo_wait_termination_before_cleanup_on_success( } k = KubernetesPodOperator(task_id="task", deferrable=True, do_xcom_push=do_xcom_push) - k.execute_complete({"ti": ti_mock}, success_event) + k.execute_complete({}, success_event) # check if it gets the pod mocked_hook.return_value.get_pod.assert_called_once_with(TEST_NAME, TEST_NAMESPACE) - # check if it pushes the xcom - assert ti_mock.xcom_push.call_count == 2 - ti_mock.xcom_push.assert_any_call(key="pod_name", value=TEST_NAME) - ti_mock.xcom_push.assert_any_call(key="pod_namespace", value=TEST_NAMESPACE) - # assert that the xcom are extracted/not extracted if do_xcom_push: mock_extract_xcom.assert_called_once()