From 6a5bc721ac073f870043d292b5be955ea59115de Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Tue, 18 Oct 2022 21:32:20 -0700 Subject: [PATCH 1/5] EMR System test --- airflow/providers/amazon/aws/operators/emr.py | 23 ++++- airflow/providers/amazon/aws/sensors/emr.py | 3 +- .../operators/emr.rst | 14 ++-- .../providers/amazon/aws}/example_emr.py | 84 +++++++++++-------- 4 files changed, 77 insertions(+), 47 deletions(-) rename {airflow/providers/amazon/aws/example_dags => tests/system/providers/amazon/aws}/example_emr.py (63%) diff --git a/airflow/providers/amazon/aws/operators/emr.py b/airflow/providers/amazon/aws/operators/emr.py index 0935ed2422f1f..53cabbef7f89a 100644 --- a/airflow/providers/amazon/aws/operators/emr.py +++ b/airflow/providers/amazon/aws/operators/emr.py @@ -22,6 +22,8 @@ from typing import TYPE_CHECKING, Any, Sequence from uuid import uuid4 +import boto3 + from airflow.exceptions import AirflowException from airflow.models import BaseOperator from airflow.providers.amazon.aws.hooks.emr import EmrContainerHook, EmrHook, EmrServerlessHook @@ -50,6 +52,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 +70,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,6 +83,7 @@ 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) @@ -113,11 +118,21 @@ def execute(self, context: Context) -> list[str]: response = emr.add_job_flow_steps(JobFlowId=job_flow_id, Steps=steps) - if not response["ResponseMetadata"]["HTTPStatusCode"] == 200: + if 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"] + + self.log.info("Steps %s added to JobFlow", response["StepIds"]) + if self.wait_for_completion: + for step_id in response["StepIds"]: + boto3.client("emr").get_waiter("step_complete").wait( + ClusterId=self.job_flow_id, + StepId=step_id, + WaiterConfig={ + "Delay": 5, + "MaxAttempts": 100, + }, + ) + return response["StepIds"] class EmrEksCreateClusterOperator(BaseOperator): diff --git a/airflow/providers/amazon/aws/sensors/emr.py b/airflow/providers/amazon/aws/sensors/emr.py index 0659554762f16..a3684fa249a1d 100644 --- a/airflow/providers/amazon/aws/sensors/emr.py +++ b/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. diff --git a/docs/apache-airflow-providers-amazon/operators/emr.rst b/docs/apache-airflow-providers-amazon/operators/emr.rst index 6bb16bd494639..ee6e9698eb5ac 100644 --- a/docs/apache-airflow-providers-amazon/operators/emr.rst +++ b/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] @@ -149,7 +149,7 @@ 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 +.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_emr.py :language: python :dedent: 4 :start-after: [START howto_sensor_emr_step] diff --git a/airflow/providers/amazon/aws/example_dags/example_emr.py b/tests/system/providers/amazon/aws/example_emr.py similarity index 63% rename from airflow/providers/amazon/aws/example_dags/example_emr.py rename to tests/system/providers/amazon/aws/example_emr.py index 2f5e2d0d5e41a..a39ef928352cf 100644 --- a/airflow/providers/amazon/aws/example_dags/example_emr.py +++ b/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) From ce87fa1978c768505680ae18c33447c62360cbe0 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Wed, 26 Oct 2022 00:03:53 -0700 Subject: [PATCH 2/5] Remove Step Sensor from .rst --- docs/apache-airflow-providers-amazon/operators/emr.rst | 9 --------- 1 file changed, 9 deletions(-) diff --git a/docs/apache-airflow-providers-amazon/operators/emr.rst b/docs/apache-airflow-providers-amazon/operators/emr.rst index ee6e9698eb5ac..a9c434777d71c 100644 --- a/docs/apache-airflow-providers-amazon/operators/emr.rst +++ b/docs/apache-airflow-providers-amazon/operators/emr.rst @@ -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:: /../../tests/system/providers/amazon/aws/example_emr.py - :language: python - :dedent: 4 - :start-after: [START howto_sensor_emr_step] - :end-before: [END howto_sensor_emr_step] - Reference --------- From 7f1c9bf25de534c2e4028d5a40f7167784720ba6 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Wed, 26 Oct 2022 13:47:18 -0700 Subject: [PATCH 3/5] Add EmrStepSensor to Missing examples to exempt from tests --- tests/always/test_project_structure.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/always/test_project_structure.py b/tests/always/test_project_structure.py index 2f215d2a6f29c..438ff2f215b96 100644 --- a/tests/always/test_project_structure.py +++ b/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", } From 0903d0004da3ebfd937322710b45cfc55c6d449a Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Tue, 8 Nov 2022 22:06:20 -0800 Subject: [PATCH 4/5] Move wait_for_completion logic to EmrHook. Add unit tests for new hook method. --- airflow/providers/amazon/aws/hooks/emr.py | 22 +++++ airflow/providers/amazon/aws/operators/emr.py | 22 +---- tests/providers/amazon/aws/hooks/test_emr.py | 85 +++++++++++++++++++ 3 files changed, 108 insertions(+), 21 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/emr.py b/airflow/providers/amazon/aws/hooks/emr.py index 1365c1705d036..4f5744f6af2c4 100644 --- a/airflow/providers/amazon/aws/hooks/emr.py +++ b/airflow/providers/amazon/aws/hooks/emr.py @@ -124,6 +124,28 @@ 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 + ): + 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). diff --git a/airflow/providers/amazon/aws/operators/emr.py b/airflow/providers/amazon/aws/operators/emr.py index 53cabbef7f89a..63659e8188b9d 100644 --- a/airflow/providers/amazon/aws/operators/emr.py +++ b/airflow/providers/amazon/aws/operators/emr.py @@ -22,8 +22,6 @@ from typing import TYPE_CHECKING, Any, Sequence from uuid import uuid4 -import boto3 - from airflow.exceptions import AirflowException from airflow.models import BaseOperator from airflow.providers.amazon.aws.hooks.emr import EmrContainerHook, EmrHook, EmrServerlessHook @@ -88,8 +86,6 @@ def __init__( 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 ) @@ -116,23 +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 response["ResponseMetadata"]["HTTPStatusCode"] != 200: - raise AirflowException(f"Adding steps failed: {response}") - - self.log.info("Steps %s added to JobFlow", response["StepIds"]) - if self.wait_for_completion: - for step_id in response["StepIds"]: - boto3.client("emr").get_waiter("step_complete").wait( - ClusterId=self.job_flow_id, - StepId=step_id, - WaiterConfig={ - "Delay": 5, - "MaxAttempts": 100, - }, - ) - 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): diff --git a/tests/providers/amazon/aws/hooks/test_emr.py b/tests/providers/amazon/aws/hooks/test_emr.py index 8e91aba5e6fbc..8514168340032 100644 --- a/tests/providers/amazon/aws/hooks/test_emr.py +++ b/tests/providers/amazon/aws/hooks/test_emr.py @@ -43,6 +43,91 @@ 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 + def test_add_job_flow_steps_one_step(self): + boto3.client("emr", region_name="us-east-1") + + 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": "step_1", + } + ] + response = hook.add_job_flow_steps(job_flow_id=cluster["JobFlowId"], steps=steps) + + assert len(response) == 1 + + @mock_emr + def test_add_job_flow_steps_multiple_steps(self): + boto3.client("emr", region_name="us-east-1") + + 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": "step_1", + }, + { + "ActionOnFailure": "test_step", + "HadoopJarStep": { + "Args": ["test args"], + "Jar": "test.jar", + }, + "Name": "step_2", + }, + ] + response = hook.add_job_flow_steps(job_flow_id=cluster["JobFlowId"], steps=steps) + assert len(response) == 2 + + @mock.patch("airflow.providers.amazon.aws.hooks.emr.EmrHook.conn") + def test_add_job_flow_steps_wait_for_completion(self, mock_conn): + boto3.client("emr", region_name="us-east-1") + + 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): """ From 0ac97a6a6cac4d3e0ce7fdaa692b79256aad91ad Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Thu, 10 Nov 2022 14:24:33 -0800 Subject: [PATCH 5/5] Add doc to function, refactor tests --- airflow/providers/amazon/aws/hooks/emr.py | 9 +++- tests/providers/amazon/aws/hooks/test_emr.py | 45 ++++---------------- 2 files changed, 16 insertions(+), 38 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/emr.py b/airflow/providers/amazon/aws/hooks/emr.py index 4f5744f6af2c4..5423dd1af83a5 100644 --- a/airflow/providers/amazon/aws/hooks/emr.py +++ b/airflow/providers/amazon/aws/hooks/emr.py @@ -126,7 +126,14 @@ def create_job_flow(self, job_flow_overrides: dict[str, Any]) -> dict[str, Any]: 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: diff --git a/tests/providers/amazon/aws/hooks/test_emr.py b/tests/providers/amazon/aws/hooks/test_emr.py index 8514168340032..ea3cf73eeb42e 100644 --- a/tests/providers/amazon/aws/hooks/test_emr.py +++ b/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 @@ -44,9 +45,8 @@ 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 - def test_add_job_flow_steps_one_step(self): - boto3.client("emr", region_name="us-east-1") - + @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}} @@ -58,47 +58,18 @@ def test_add_job_flow_steps_one_step(self): "Args": ["test args"], "Jar": "test.jar", }, - "Name": "step_1", + "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) == 1 - - @mock_emr - def test_add_job_flow_steps_multiple_steps(self): - boto3.client("emr", region_name="us-east-1") - - 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": "step_1", - }, - { - "ActionOnFailure": "test_step", - "HadoopJarStep": { - "Args": ["test args"], - "Jar": "test.jar", - }, - "Name": "step_2", - }, - ] - response = hook.add_job_flow_steps(job_flow_id=cluster["JobFlowId"], steps=steps) - assert len(response) == 2 + 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): - boto3.client("emr", region_name="us-east-1") - 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",