System Test for EMR (AIP-47)

airflow/providers/amazon/aws/hooks/emr.py

@@ -124,6 +124,35 @@ def create_job_flow(self, job_flow_overrides: dict[str, Any]) -> dict[str, Any]:
 
         return response
 
+    def add_job_flow_steps(
+        self, job_flow_id: str, steps: list[dict] | str | None = None, wait_for_completion: bool = False
+    ) -> list[str]:
+        """
+        Add new steps to a running cluster.
+
+        :param job_flow_id: The id of the job flow to which the steps are being added
+        :param steps: A list of the steps to be executed by the job flow
+        :param wait_for_completion: If True, wait for the steps to be completed. Default is False
+        """
+        response = self.get_conn().add_job_flow_steps(JobFlowId=job_flow_id, Steps=steps)
+
+        if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
+            raise AirflowException(f"Adding steps failed: {response}")
+
+        self.log.info("Steps %s added to JobFlow", response["StepIds"])
+        if wait_for_completion:
+            waiter = self.get_conn().get_waiter("step_complete")
+            for step_id in response["StepIds"]:
+                waiter.wait(
+                    ClusterId=job_flow_id,
+                    StepId=step_id,
+                    WaiterConfig={
+                        "Delay": 5,
+                        "MaxAttempts": 100,
+                    },
+                )
+        return response["StepIds"]
+
     def test_connection(self):
         """
         Return failed state for test Amazon Elastic MapReduce Connection (untestable).

airflow/providers/amazon/aws/operators/emr.py

@@ -50,6 +50,7 @@ class EmrAddStepsOperator(BaseOperator):
     :param aws_conn_id: aws connection to uses
     :param steps: boto3 style steps or reference to a steps file (must be '.json') to
         be added to the jobflow. (templated)
+    :param wait_for_completion: If True, the operator will wait for all the steps to be completed.
     :param do_xcom_push: if True, job_flow_id is pushed to XCom with key job_flow_id.
     """
 
@@ -67,6 +68,7 @@ def __init__(
         cluster_states: list[str] | None = None,
         aws_conn_id: str = "aws_default",
         steps: list[dict] | str | None = None,
+        wait_for_completion: bool = False,
         **kwargs,
     ):
         if not (job_flow_id is None) ^ (job_flow_name is None):
@@ -79,12 +81,11 @@ def __init__(
         self.job_flow_name = job_flow_name
         self.cluster_states = cluster_states
         self.steps = steps
+        self.wait_for_completion = wait_for_completion
 
     def execute(self, context: Context) -> list[str]:
         emr_hook = EmrHook(aws_conn_id=self.aws_conn_id)
 
-        emr = emr_hook.get_conn()
-
         job_flow_id = self.job_flow_id or emr_hook.get_cluster_id_by_name(
             str(self.job_flow_name), self.cluster_states
         )
@@ -111,13 +112,7 @@ def execute(self, context: Context) -> list[str]:
         if isinstance(steps, str):
             steps = ast.literal_eval(steps)
 
-        response = emr.add_job_flow_steps(JobFlowId=job_flow_id, Steps=steps)
-
-        if not response["ResponseMetadata"]["HTTPStatusCode"] == 200:
-            raise AirflowException(f"Adding steps failed: {response}")
-        else:
-            self.log.info("Steps %s added to JobFlow", response["StepIds"])
-            return response["StepIds"]
+        return emr_hook.add_job_flow_steps(job_flow_id=job_flow_id, steps=steps, wait_for_completion=True)
 
 
 class EmrEksCreateClusterOperator(BaseOperator):

airflow/providers/amazon/aws/sensors/emr.py

@@ -21,7 +21,7 @@
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.emr import EmrContainerHook, EmrHook, EmrServerlessHook
-from airflow.sensors.base import BaseSensorOperator
+from airflow.sensors.base import BaseSensorOperator, poke_mode_only
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -378,6 +378,7 @@ def failure_message_from_response(response: dict[str, Any]) -> str | None:
         return None
 
 
+@poke_mode_only
 class EmrStepSensor(EmrBaseSensor):
     """
     Asks for the state of the step until it reaches any of the target states.

docs/apache-airflow-providers-amazon/operators/emr.rst

@@ -53,7 +53,7 @@ JobFlow configuration
 
 To create a job flow on EMR, you need to specify the configuration for the EMR cluster:
 
-.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_emr.py
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_emr.py
     :language: python
     :start-after: [START howto_operator_emr_steps_config]
     :end-before: [END howto_operator_emr_steps_config]
@@ -76,7 +76,7 @@ Create the Job Flow
 
 In the following code we are creating a new job flow using the configuration as explained above.
 
-.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_emr.py
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_emr.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_emr_create_job_flow]
@@ -90,7 +90,7 @@ Add Steps to an EMR job flow
 To add steps to an existing EMR Job flow you can use
 :class:`~airflow.providers.amazon.aws.operators.emr.EmrAddStepsOperator`.
 
-.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_emr.py
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_emr.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_emr_add_steps]
@@ -104,7 +104,7 @@ Terminate an EMR job flow
 To terminate an EMR Job Flow you can use
 :class:`~airflow.providers.amazon.aws.operators.emr.EmrTerminateJobFlowOperator`.
 
-.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_emr.py
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_emr.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_emr_terminate_job_flow]
@@ -118,7 +118,7 @@ Modify Amazon EMR container
 To modify an existing EMR container you can use
 :class:`~airflow.providers.amazon.aws.sensors.emr.EmrContainerSensor`.
 
-.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_emr.py
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_emr.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_emr_modify_cluster]
@@ -135,7 +135,7 @@ Wait on an Amazon EMR job flow state
 To monitor the state of an EMR job flow you can use
 :class:`~airflow.providers.amazon.aws.sensors.emr.EmrJobFlowSensor`.
 
-.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_emr.py
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_emr.py
     :language: python
     :dedent: 4
     :start-after: [START howto_sensor_emr_job_flow]
@@ -146,15 +146,6 @@ To monitor the state of an EMR job flow you can use
 Wait on an Amazon EMR step state
 ================================
 
-To monitor the state of a step running an existing EMR Job flow you can use
-:class:`~airflow.providers.amazon.aws.sensors.emr.EmrStepSensor`.
-
-.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_emr.py
-    :language: python
-    :dedent: 4
-    :start-after: [START howto_sensor_emr_step]
-    :end-before: [END howto_sensor_emr_step]
-
 Reference
 ---------
 

tests/always/test_project_structure.py

@@ -400,6 +400,8 @@ class TestAmazonProviderProjectStructure(ExampleCoverageTest):
         "airflow.providers.amazon.aws.transfers.exasol_to_s3.ExasolToS3Operator",
         # Glue Catalog sensor difficult to test
         "airflow.providers.amazon.aws.sensors.glue_catalog_partition.GlueCatalogPartitionSensor",
+        # EMR Step sensor difficult to test, see: https://github.com/apache/airflow/pull/27286
+        "airflow.providers.amazon.aws.sensors.emr.EmrStepSensor",
     }
 
 

tests/providers/amazon/aws/hooks/test_emr.py

@@ -17,6 +17,7 @@
 # under the License.
 from __future__ import annotations
 
+import re
 from unittest import mock
 
 import boto3
@@ -43,6 +44,61 @@ def test_create_job_flow_uses_the_emr_config_to_create_a_cluster(self):
 
         assert client.list_clusters()["Clusters"][0]["Id"] == cluster["JobFlowId"]
 
+    @mock_emr
+    @pytest.mark.parametrize("num_steps", [1, 2, 3, 4])
+    def test_add_job_flow_steps_one_step(self, num_steps):
+        hook = EmrHook(aws_conn_id="aws_default", emr_conn_id="emr_default", region_name="us-east-1")
+        cluster = hook.create_job_flow(
+            {"Name": "test_cluster", "Instances": {"KeepJobFlowAliveWhenNoSteps": False}}
+        )
+        steps = [
+            {
+                "ActionOnFailure": "test_step",
+                "HadoopJarStep": {
+                    "Args": ["test args"],
+                    "Jar": "test.jar",
+                },
+                "Name": f"step_{i}",
+            }
+            for i in range(num_steps)
+        ]
+        response = hook.add_job_flow_steps(job_flow_id=cluster["JobFlowId"], steps=steps)
+
+        assert len(response) == num_steps
+        for step_id in response:
+            assert re.match("s-[A-Z0-9]{13}$", step_id)
+
+    @mock.patch("airflow.providers.amazon.aws.hooks.emr.EmrHook.conn")
+    def test_add_job_flow_steps_wait_for_completion(self, mock_conn):
+        hook = EmrHook(aws_conn_id="aws_default", emr_conn_id="emr_default", region_name="us-east-1")
+        mock_conn.run_job_flow.return_value = {
+            "JobFlowId": "job_flow_id",
+            "ClusterArn": "cluster_arn",
+        }
+        mock_conn.add_job_flow_steps.return_value = {
+            "StepIds": [
+                "step_id",
+            ],
+            "ResponseMetadata": {"HTTPStatusCode": 200},
+        }
+
+        hook.create_job_flow({"Name": "test_cluster", "Instances": {"KeepJobFlowAliveWhenNoSteps": False}})
+
+        steps = [
+            {
+                "ActionOnFailure": "test_step",
+                "HadoopJarStep": {
+                    "Args": ["test args"],
+                    "Jar": "test.jar",
+                },
+                "Name": "step_1",
+            }
+        ]
+
+        hook.add_job_flow_steps(job_flow_id="job_flow_id", steps=steps, wait_for_completion=True)
+
+        mock_conn.get_waiter.assert_called_once_with("step_complete")
+
     @mock_emr
     def test_create_job_flow_extra_args(self):
         """

airflow/providers/amazon/aws/example_dags/example_emr.py → tests/system/providers/amazon/aws/example_emr.py

@@ -15,9 +15,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
 from __future__ import annotations
 
-import os
 from datetime import datetime
 
 from airflow import DAG
@@ -28,10 +28,11 @@
     EmrModifyClusterOperator,
     EmrTerminateJobFlowOperator,
 )
-from airflow.providers.amazon.aws.sensors.emr import EmrJobFlowSensor, EmrStepSensor
+from airflow.providers.amazon.aws.sensors.emr import EmrJobFlowSensor
+from airflow.utils.trigger_rule import TriggerRule
+from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
 
-JOB_FLOW_ROLE = os.getenv("EMR_JOB_FLOW_ROLE", "EMR_EC2_DefaultRole")
-SERVICE_ROLE = os.getenv("EMR_SERVICE_ROLE", "EMR_DefaultRole")
+DAG_ID = "example_emr"
 
 # [START howto_operator_emr_steps_config]
 SPARK_STEPS = [
@@ -63,64 +64,77 @@
         "TerminationProtected": False,
     },
     "Steps": SPARK_STEPS,
-    "JobFlowRole": JOB_FLOW_ROLE,
-    "ServiceRole": SERVICE_ROLE,
+    "JobFlowRole": "EMR_EC2_DefaultRole",
+    "ServiceRole": "EMR_DefaultRole",
 }
 # [END howto_operator_emr_steps_config]
 
+sys_test_context_task = SystemTestContextBuilder().build()
+
 with DAG(
-    dag_id="example_emr",
+    dag_id=DAG_ID,
     start_date=datetime(2021, 1, 1),
-    tags=["example"],
+    schedule="@once",
     catchup=False,
+    tags=["example"],
 ) as dag:
+    test_context = sys_test_context_task()
+    env_id = test_context[ENV_ID_KEY]
+
     # [START howto_operator_emr_create_job_flow]
-    job_flow_creator = EmrCreateJobFlowOperator(
+    create_job_flow = EmrCreateJobFlowOperator(
         task_id="create_job_flow",
         job_flow_overrides=JOB_FLOW_OVERRIDES,
     )
     # [END howto_operator_emr_create_job_flow]
 
-    job_flow_id = job_flow_creator.output
-
-    # [START howto_sensor_emr_job_flow]
-    job_sensor = EmrJobFlowSensor(task_id="check_job_flow", job_flow_id=job_flow_id)
-    # [END howto_sensor_emr_job_flow]
-
     # [START howto_operator_emr_modify_cluster]
-    cluster_modifier = EmrModifyClusterOperator(
-        task_id="modify_cluster", cluster_id=job_flow_id, step_concurrency_level=1
+    modify_cluster = EmrModifyClusterOperator(
+        task_id="modify_cluster", cluster_id=create_job_flow.output, step_concurrency_level=1
     )
     # [END howto_operator_emr_modify_cluster]
 
     # [START howto_operator_emr_add_steps]
-    step_adder = EmrAddStepsOperator(
+    add_steps = EmrAddStepsOperator(
         task_id="add_steps",
-        job_flow_id=job_flow_id,
+        job_flow_id=create_job_flow.output,
         steps=SPARK_STEPS,
+        wait_for_completion=True,
     )
     # [END howto_operator_emr_add_steps]
 
-    # [START howto_sensor_emr_step]
-    step_checker = EmrStepSensor(
-        task_id="watch_step",
-        job_flow_id=job_flow_id,
-        step_id="{{ task_instance.xcom_pull(task_ids='add_steps', key='return_value')[0] }}",
-    )
-    # [END howto_sensor_emr_step]
-
     # [START howto_operator_emr_terminate_job_flow]
-    cluster_remover = EmrTerminateJobFlowOperator(
+    remove_cluster = EmrTerminateJobFlowOperator(
         task_id="remove_cluster",
-        job_flow_id=job_flow_id,
+        job_flow_id=create_job_flow.output,
     )
     # [END howto_operator_emr_terminate_job_flow]
+    remove_cluster.trigger_rule = TriggerRule.ALL_DONE
+
+    # [START howto_sensor_emr_job_flow]
+    check_job_flow = EmrJobFlowSensor(task_id="check_job_flow", job_flow_id=create_job_flow.output)
+    # [END howto_sensor_emr_job_flow]
 
     chain(
-        job_flow_creator,
-        job_sensor,
-        cluster_modifier,
-        step_adder,
-        step_checker,
-        cluster_remover,
+        # TEST SETUP
+        test_context,
+        # TEST BODY
+        create_job_flow,
+        modify_cluster,
+        add_steps,
+        # TEST TEARDOWN
+        remove_cluster,
+        check_job_flow,
     )
+
+    from tests.system.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.system.utils import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
+test_run = get_test_run(dag)