From d1be3dcec94c57d8a4cb339b10c18d9ba6da0d6f Mon Sep 17 00:00:00 2001 From: Leondon9 Date: Tue, 25 Aug 2026 01:25:52 +0900 Subject: [PATCH] Add @task.eks_pod TaskFlow decorator for EksPodOperator The TaskFlow route to EKS today is @task.kubernetes with a connection whose kubeconfig carries an `aws eks get-token` exec plugin. The token is then minted by an exec subprocess on the worker, so the worker must ship the AWS CLI and credentials. This decorator wraps EksPodOperator instead, which builds the kubeconfig and short-lived token from the AWS connection (aws_conn_id) over boto3/STS, so the worker needs neither the CLI nor a pre-built kubeconfig. The script-injection logic is kept self-contained rather than shared with the @task.kubernetes implementation in the cncf provider, to avoid a cross-provider private import; a shared helper can follow later. --- providers/amazon/docs/operators/eks.rst | 20 ++ providers/amazon/provider.yaml | 4 + .../amazon/aws/decorators/__init__.py | 17 ++ .../providers/amazon/aws/decorators/eks.py | 172 ++++++++++++++ .../providers/amazon/get_provider_info.py | 3 + .../amazon/aws/example_eks_pod_decorator.py | 213 ++++++++++++++++++ .../unit/amazon/aws/decorators/__init__.py | 17 ++ .../unit/amazon/aws/decorators/test_eks.py | 97 ++++++++ 8 files changed, 543 insertions(+) create mode 100644 providers/amazon/src/airflow/providers/amazon/aws/decorators/__init__.py create mode 100644 providers/amazon/src/airflow/providers/amazon/aws/decorators/eks.py create mode 100644 providers/amazon/tests/system/amazon/aws/example_eks_pod_decorator.py create mode 100644 providers/amazon/tests/unit/amazon/aws/decorators/__init__.py create mode 100644 providers/amazon/tests/unit/amazon/aws/decorators/test_eks.py diff --git a/providers/amazon/docs/operators/eks.rst b/providers/amazon/docs/operators/eks.rst index 76c0f5ea4ed11..1b94e31789533 100644 --- a/providers/amazon/docs/operators/eks.rst +++ b/providers/amazon/docs/operators/eks.rst @@ -205,6 +205,26 @@ Note: An Amazon EKS Cluster with underlying compute infrastructure is required. :start-after: [START howto_operator_eks_pod_operator] :end-before: [END howto_operator_eks_pod_operator] +.. _howto/decorator:eks_pod: + +Run a Python task on an Amazon EKS Cluster from TaskFlow +======================================================== + +The ``@task.eks_pod`` decorator runs a decorated Python function inside a pod on an existing Amazon +EKS Cluster, using the TaskFlow style. It wraps +:class:`~airflow.providers.amazon.aws.operators.eks.EksPodOperator`, which builds the cluster +kubeconfig and a short-lived token from the AWS connection (``aws_conn_id``) at run time. This means +the worker does not need the ``aws`` CLI or a pre-built kubeconfig, which is what running +``@task.kubernetes`` against EKS would otherwise require. + +An Amazon EKS Cluster with underlying compute infrastructure is required. + +.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_eks_pod_decorator.py + :language: python + :dedent: 4 + :start-after: [START howto_decorator_eks_pod] + :end-before: [END howto_decorator_eks_pod] + Sensors ------- diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml index 3d26682c4d5ec..ea1a1149a0174 100644 --- a/providers/amazon/provider.yaml +++ b/providers/amazon/provider.yaml @@ -986,6 +986,10 @@ transfers: how-to-guide: /docs/apache-airflow-providers-amazon/transfer/s3_to_dynamodb.rst python-module: airflow.providers.amazon.aws.transfers.s3_to_dynamodb +task-decorators: + - class-name: airflow.providers.amazon.aws.decorators.eks.eks_pod_task + name: eks_pod + extra-links: - airflow.providers.amazon.aws.links.athena.AthenaQueryResultsLink - airflow.providers.amazon.aws.links.batch.BatchJobDefinitionLink diff --git a/providers/amazon/src/airflow/providers/amazon/aws/decorators/__init__.py b/providers/amazon/src/airflow/providers/amazon/aws/decorators/__init__.py new file mode 100644 index 0000000000000..21d298ede6ed3 --- /dev/null +++ b/providers/amazon/src/airflow/providers/amazon/aws/decorators/__init__.py @@ -0,0 +1,17 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations diff --git a/providers/amazon/src/airflow/providers/amazon/aws/decorators/eks.py b/providers/amazon/src/airflow/providers/amazon/aws/decorators/eks.py new file mode 100644 index 0000000000000..af219944faf30 --- /dev/null +++ b/providers/amazon/src/airflow/providers/amazon/aws/decorators/eks.py @@ -0,0 +1,172 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import base64 +import os +import pickle +from collections.abc import Callable, Sequence +from shlex import quote +from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING + +import dill +from kubernetes.client import models as k8s + +from airflow.providers.amazon.aws.operators.eks import EksPodOperator +from airflow.providers.cncf.kubernetes.python_kubernetes_script import write_python_script +from airflow.providers.common.compat.sdk import ( + DecoratedOperator, + TaskDecorator, + task_decorator_factory, +) + +if TYPE_CHECKING: + from airflow.sdk import Context + +_PYTHON_SCRIPT_ENV = "__PYTHON_SCRIPT" +_PYTHON_INPUT_ENV = "__PYTHON_INPUT" + + +def _generate_decoded_command(env_var: str, file: str) -> str: + return ( + f'python -c "import base64, os;' + rf"x = base64.b64decode(os.environ[\"{env_var}\"]);" + rf'f = open(\"{file}\", \"wb\"); f.write(x); f.close()"' + ) + + +def _read_file_contents(filename: str) -> str: + with open(filename, "rb") as script_file: + return base64.b64encode(script_file.read()).decode("utf-8") + + +class _EksPodDecoratedOperator(DecoratedOperator, EksPodOperator): + """Wrap a Python callable to run inside a pod on Amazon EKS via ``EksPodOperator``.""" + + custom_operator_name = "@task.eks_pod" + + # `cmds` and `arguments` are used internally by the operator + template_fields: Sequence[str] = tuple( + {"op_args", "op_kwargs", *EksPodOperator.template_fields} - {"cmds", "arguments"} + ) + + # Since we won't mutate the arguments, we should just do the shallow copy + # there are some cases we can't deepcopy the objects (e.g protobuf). + shallow_copy_attrs: Sequence[str] = ("python_callable",) + + def __init__(self, *, cluster_name: str, use_dill: bool = False, **kwargs) -> None: + self.use_dill = use_dill + + # Accept the EKS ``pod_name`` or the K8s-style ``name``, otherwise derive one from the callable. + pod_name = ( + kwargs.pop("pod_name", None) + or kwargs.pop("name", None) + or f"eks-airflow-pod-{kwargs['python_callable'].__name__}" + ) + random_name_suffix = kwargs.pop("random_name_suffix", True) + super().__init__( + cluster_name=cluster_name, + pod_name=pod_name, + random_name_suffix=random_name_suffix, + cmds=["placeholder-command"], + **kwargs, + ) + + def _generate_cmds(self) -> list[str]: + script_filename = "/tmp/script.py" + input_filename = "/tmp/script.in" + + if getattr(self, "do_xcom_push", False): + output_filename = "/airflow/xcom/return.json" + make_xcom_dir_cmd = "mkdir -p /airflow/xcom" + else: + output_filename = "/dev/null" + make_xcom_dir_cmd = ":" # shell no-op + + write_local_script_file_cmd = ( + f"{_generate_decoded_command(quote(_PYTHON_SCRIPT_ENV), quote(script_filename))}" + ) + write_local_input_file_cmd = ( + f"{_generate_decoded_command(quote(_PYTHON_INPUT_ENV), quote(input_filename))}" + ) + exec_python_cmd = f"python {script_filename} {input_filename} {output_filename}" + return [ + "bash", + "-cx", + ( + f"{write_local_script_file_cmd} && " + f"{write_local_input_file_cmd} && " + f"{make_xcom_dir_cmd} && " + f"{exec_python_cmd}" + ), + ] + + def execute(self, context: Context): + with TemporaryDirectory(prefix="venv") as tmp_dir: + pickling_library = dill if self.use_dill else pickle + script_filename = os.path.join(tmp_dir, "script.py") + input_filename = os.path.join(tmp_dir, "script.in") + + with open(input_filename, "wb") as file: + pickling_library.dump({"args": self.op_args, "kwargs": self.op_kwargs}, file) + + py_source = self.get_python_source() + jinja_context = { + "op_args": self.op_args, + "op_kwargs": self.op_kwargs, + "pickling_library": pickling_library.__name__, + "python_callable": self.python_callable.__name__, + "python_callable_source": py_source, + "string_args_global": False, + } + write_python_script(jinja_context=jinja_context, filename=script_filename) + + self.env_vars: list[k8s.V1EnvVar] = [ + *self.env_vars, + k8s.V1EnvVar(name=_PYTHON_SCRIPT_ENV, value=_read_file_contents(script_filename)), + k8s.V1EnvVar(name=_PYTHON_INPUT_ENV, value=_read_file_contents(input_filename)), + ] + + self.cmds = self._generate_cmds() + return super().execute(context) + + +def eks_pod_task( + python_callable: Callable | None = None, + multiple_outputs: bool | None = None, + **kwargs, +) -> TaskDecorator: + """ + Run a Python function in a pod on an Amazon EKS cluster. + + Wraps :class:`~airflow.providers.amazon.aws.operators.eks.EksPodOperator`, which builds the + cluster kubeconfig and a short-lived token from the AWS connection (``aws_conn_id``), so the + worker needs no ``aws`` CLI or kubeconfig. Any ``EksPodOperator`` argument is accepted via + ``kwargs``. + + :param python_callable: Function to decorate + :param multiple_outputs: if set, function return value will be + unrolled to multiple XCom values. Dict will unroll to xcom values with + keys as XCom keys. Defaults to False. + """ + return task_decorator_factory( + python_callable=python_callable, + multiple_outputs=multiple_outputs, + decorated_operator_class=_EksPodDecoratedOperator, + **kwargs, + ) diff --git a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py index ea3e1f49437cc..461d2ff3e1d75 100644 --- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py +++ b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py @@ -1155,6 +1155,9 @@ def get_provider_info(): "python-module": "airflow.providers.amazon.aws.transfers.s3_to_dynamodb", }, ], + "task-decorators": [ + {"class-name": "airflow.providers.amazon.aws.decorators.eks.eks_pod_task", "name": "eks_pod"} + ], "extra-links": [ "airflow.providers.amazon.aws.links.athena.AthenaQueryResultsLink", "airflow.providers.amazon.aws.links.batch.BatchJobDefinitionLink", diff --git a/providers/amazon/tests/system/amazon/aws/example_eks_pod_decorator.py b/providers/amazon/tests/system/amazon/aws/example_eks_pod_decorator.py new file mode 100644 index 0000000000000..4b7f35f23a548 --- /dev/null +++ b/providers/amazon/tests/system/amazon/aws/example_eks_pod_decorator.py @@ -0,0 +1,213 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from datetime import datetime + +import boto3 + +from airflow.providers.amazon.aws.hooks.eks import ClusterStates, NodegroupStates +from airflow.providers.amazon.aws.operators.eks import ( + EksCreateClusterOperator, + EksCreateNodegroupOperator, + EksDeleteClusterOperator, + EksDeleteNodegroupOperator, +) +from airflow.providers.amazon.aws.sensors.eks import EksClusterStateSensor, EksNodegroupStateSensor + +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk import DAG, chain, task +else: + # Airflow 2 path + from airflow.decorators import task # type: ignore[attr-defined,no-redef] + from airflow.models.baseoperator import chain # type: ignore[attr-defined,no-redef] + from airflow.models.dag import DAG # type: ignore[attr-defined,no-redef,assignment] + +try: + from airflow.sdk import TriggerRule +except ImportError: + # Compatibility for Airflow < 3.1 + from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined] + +from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +from system.amazon.aws.utils.k8s import get_describe_pod_operator + +DAG_ID = "example_eks_pod_decorator" + +# Externally fetched variables: +ROLE_ARN_KEY = "ROLE_ARN" +SUBNETS_KEY = "SUBNETS" + +sys_test_context_task = ( + SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).add_variable(SUBNETS_KEY, split_string=True).build() +) + + +@task +def create_launch_template(template_name: str): + # This launch template enables IMDSv2. + boto3.client("ec2").create_launch_template( + LaunchTemplateName=template_name, + LaunchTemplateData={ + "MetadataOptions": {"HttpEndpoint": "enabled", "HttpTokens": "required"}, + }, + ) + + +@task(trigger_rule=TriggerRule.ALL_DONE) +def delete_launch_template(template_name: str): + boto3.client("ec2").delete_launch_template(LaunchTemplateName=template_name) + + +with DAG( + dag_id=DAG_ID, + schedule="@once", + start_date=datetime(2021, 1, 1), + catchup=False, +) as dag: + test_context = sys_test_context_task() + env_id = test_context[ENV_ID_KEY] + + cluster_name = f"{env_id}-cluster" + nodegroup_name = f"{env_id}-nodegroup" + launch_template_name = f"{env_id}-launch-template" + + create_cluster = EksCreateClusterOperator( + task_id="create_cluster", + cluster_name=cluster_name, + cluster_role_arn=test_context[ROLE_ARN_KEY], + resources_vpc_config={"subnetIds": test_context[SUBNETS_KEY]}, + compute=None, + ) + + await_create_cluster = EksClusterStateSensor( + task_id="await_create_cluster", + cluster_name=cluster_name, + target_state=ClusterStates.ACTIVE, + ) + + create_nodegroup = EksCreateNodegroupOperator( + task_id="create_nodegroup", + cluster_name=cluster_name, + nodegroup_name=nodegroup_name, + nodegroup_subnets=test_context[SUBNETS_KEY], + nodegroup_role_arn=test_context[ROLE_ARN_KEY], + ) + # The launch template enforces IMDSv2 and is required for internal compliance when running + # these system tests on AWS infrastructure. It is not required for the decorator to work. + create_nodegroup.create_nodegroup_kwargs = {"launchTemplate": {"name": launch_template_name}} + + await_create_nodegroup = EksNodegroupStateSensor( + task_id="await_create_nodegroup", + cluster_name=cluster_name, + nodegroup_name=nodegroup_name, + target_state=NodegroupStates.ACTIVE, + poke_interval=10, + ) + + # [START howto_decorator_eks_pod] + # Credentials come from the AWS connection; the worker needs no aws CLI or kubeconfig. + @task.eks_pod( + pod_name="run-pod", + cluster_name=cluster_name, + image="python:3.12-slim", + get_logs=True, + on_finish_action="keep_pod", + ) + def run_pod(): + print("Hello from an EKS pod") + + start_pod = run_pod() + # [END howto_decorator_eks_pod] + + describe_pod = get_describe_pod_operator( + cluster_name, pod_name="{{ ti.xcom_pull(key='pod_name', task_ids='run_pod') }}" + ) + # only describe the pod if the task above failed, to help diagnose + describe_pod.trigger_rule = TriggerRule.ONE_FAILED + + await_nodegroup_stable = EksNodegroupStateSensor( + task_id="await_nodegroup_stable", + trigger_rule=TriggerRule.ALL_DONE, + cluster_name=cluster_name, + nodegroup_name=nodegroup_name, + target_state=NodegroupStates.ACTIVE, + ) + + delete_nodegroup = EksDeleteNodegroupOperator( + task_id="delete_nodegroup", + cluster_name=cluster_name, + nodegroup_name=nodegroup_name, + trigger_rule=TriggerRule.ALL_DONE, + ) + + await_delete_nodegroup = EksNodegroupStateSensor( + task_id="await_delete_nodegroup", + trigger_rule=TriggerRule.ALL_DONE, + cluster_name=cluster_name, + nodegroup_name=nodegroup_name, + target_state=NodegroupStates.NONEXISTENT, + ) + + delete_cluster = EksDeleteClusterOperator( + task_id="delete_cluster", + cluster_name=cluster_name, + trigger_rule=TriggerRule.ALL_DONE, + ) + + await_delete_cluster = EksClusterStateSensor( + task_id="await_delete_cluster", + trigger_rule=TriggerRule.ALL_DONE, + cluster_name=cluster_name, + target_state=ClusterStates.NONEXISTENT, + poke_interval=10, + ) + + chain( + # TEST SETUP + test_context, + create_launch_template(launch_template_name), + # TEST BODY + create_cluster, + await_create_cluster, + create_nodegroup, + await_create_nodegroup, + start_pod, + # TEST TEARDOWN + describe_pod, + await_nodegroup_stable, + delete_nodegroup, + await_delete_nodegroup, + delete_cluster, + await_delete_cluster, + delete_launch_template(launch_template_name), + ) + + from tests_common.test_utils.watcher import watcher + + # This test needs watcher in order to properly mark success/failure + # when "tearDown" task with trigger rule is part of the DAG + list(dag.tasks) >> watcher() + + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst) +test_run = get_test_run(dag) diff --git a/providers/amazon/tests/unit/amazon/aws/decorators/__init__.py b/providers/amazon/tests/unit/amazon/aws/decorators/__init__.py new file mode 100644 index 0000000000000..21d298ede6ed3 --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/decorators/__init__.py @@ -0,0 +1,17 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations diff --git a/providers/amazon/tests/unit/amazon/aws/decorators/test_eks.py b/providers/amazon/tests/unit/amazon/aws/decorators/test_eks.py new file mode 100644 index 0000000000000..f18ca4c37241e --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/decorators/test_eks.py @@ -0,0 +1,97 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest import mock + +import pytest + +from tests_common.test_utils.compat import timezone +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk import DAG +else: + from airflow.models.dag import DAG # type: ignore[no-redef,assignment] + +from airflow.providers.amazon.aws.decorators.eks import _EksPodDecoratedOperator, eks_pod_task + +DEFAULT_DATE = timezone.datetime(2023, 1, 1) + + +class TestEksPodDecorator: + def test_init_builds_eks_pod_operator(self): + with DAG(dag_id="test_eks_deco_init", schedule=None, start_date=DEFAULT_DATE): + + @eks_pod_task(cluster_name="my-eks", image="python:3.12-slim") + def f(): + return {"a": 1} + + task = f() + + op = task.operator + assert isinstance(op, _EksPodDecoratedOperator) + assert op.custom_operator_name == "@task.eks_pod" + assert op.task_id == "f" + assert op.cluster_name == "my-eks" + assert op.image == "python:3.12-slim" + assert op.cmds == ["placeholder-command"] + assert op.random_name_suffix is True + assert op.pod_name == "eks-airflow-pod-f" + + def test_explicit_pod_name_is_respected(self): + with DAG(dag_id="test_eks_deco_name", schedule=None, start_date=DEFAULT_DATE): + + @eks_pod_task(cluster_name="my-eks", image="python:3.12-slim", pod_name="custom-pod") + def f(): + return None + + task = f() + + assert task.operator.pod_name == "custom-pod" + + @mock.patch("airflow.providers.amazon.aws.decorators.eks.EksPodOperator.execute") + def test_execute_injects_script_then_delegates(self, mock_super_execute): + with DAG(dag_id="test_eks_deco_exec", schedule=None, start_date=DEFAULT_DATE): + + @eks_pod_task(cluster_name="my-eks", image="python:3.12-slim") + def f(): + return {"a": 1} + + task = f() + + op = task.operator + op.execute(context=mock.MagicMock()) + + env_names = {env.name for env in op.env_vars} + assert {"__PYTHON_SCRIPT", "__PYTHON_INPUT"} <= env_names + assert op.cmds[0] == "bash" + mock_super_execute.assert_called_once() + + @pytest.mark.parametrize("do_xcom_push", [True, False]) + def test_generate_cmds_handles_xcom(self, do_xcom_push): + with DAG(dag_id="test_eks_deco_xcom", schedule=None, start_date=DEFAULT_DATE): + + @eks_pod_task(cluster_name="my-eks", image="python:3.12-slim", do_xcom_push=do_xcom_push) + def f(): + return {"a": 1} + + task = f() + + cmds = task.operator._generate_cmds() + assert cmds[:2] == ["bash", "-cx"] + assert ("mkdir -p /airflow/xcom" in cmds[2]) is do_xcom_push