From 28e06feb4ee8efb2f77c98ce8e1991cf67a6cb1d Mon Sep 17 00:00:00 2001
From: ferruzzi <ferruzzi@amazon.com>
Date: Mon, 6 Feb 2023 10:15:23 -0800
Subject: [PATCH 1/3] Implement custom boto waiters for some EMR operators

Part of an ongoing project to migrate off the `waiter` class in the Amazon Provider Package to the custom boto waiters.
---
 airflow/providers/amazon/aws/operators/emr.py | 176 +++++++++++++-----
 airflow/providers/amazon/aws/utils/waiter.py  |   4 +
 airflow/providers/amazon/aws/waiters/emr.json |  80 ++++++++
 .../amazon/aws/operators/test_eks.py          |  12 +-
 .../aws/operators/test_emr_create_job_flow.py |  24 ++-
 .../operators/test_emr_notebook_execution.py  | 117 +++++++-----
 .../providers/amazon/aws/utils/test_waiter.py |  11 ++
 7 files changed, 311 insertions(+), 113 deletions(-)
 create mode 100644 airflow/providers/amazon/aws/waiters/emr.json

diff --git a/airflow/providers/amazon/aws/operators/emr.py b/airflow/providers/amazon/aws/operators/emr.py
index 666379530e0f1..19b8656eaa8ab 100644
--- a/airflow/providers/amazon/aws/operators/emr.py
+++ b/airflow/providers/amazon/aws/operators/emr.py
@@ -27,7 +27,8 @@
 from airflow.providers.amazon.aws.hooks.emr import EmrContainerHook, EmrHook, EmrServerlessHook
 from airflow.providers.amazon.aws.links.emr import EmrClusterLink
 from airflow.providers.amazon.aws.utils.waiter import waiter
-from airflow.utils.helpers import exactly_one
+from airflow.utils.helpers import exactly_one, prune_dict
+from airflow.utils.types import NOTSET, ArgNotSet
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -159,9 +160,11 @@ class EmrStartNotebookExecutionOperator(BaseOperator):
         group to associate with the master instance of the EMR cluster for this notebook execution.
     :param tags: Optional list of key value pair to associate with the notebook execution.
     :param waiter_countdown: Total amount of time the operator will wait for the notebook to stop.
-        Defaults to 25 * 60 seconds.
+        Defaults to 25 * 60 seconds. (Deprecated.  Please use waiter_max_attempts.)
     :param waiter_check_interval_seconds: Number of seconds between polling the state of the notebook.
-        Defaults to 60 seconds.
+        Defaults to 60 seconds. (Deprecated.  Please use waiter_delay.)
+    :param waiter_max_attempts: Maximum number of tries before failing.
+    :param waiter_delay: Number of seconds between polling the state of the notebook.
     """
 
     template_fields: Sequence[str] = (
@@ -174,6 +177,8 @@ class EmrStartNotebookExecutionOperator(BaseOperator):
         "notebook_instance_security_group_id",
         "master_instance_security_group_id",
         "tags",
+        "waiter_delay",
+        "waiter_max_attempts",
     )
 
     def __init__(
@@ -191,8 +196,23 @@ def __init__(
         aws_conn_id: str = "aws_default",
         waiter_countdown: int = 25 * 60,
         waiter_check_interval_seconds: int = 60,
+        # TODO: waiter_max_attempts and waiter_delay should default to None when the other two are deprecated.
+        waiter_max_attempts: int | ArgNotSet = NOTSET,
+        waiter_delay: int | ArgNotSet = NOTSET,
         **kwargs: Any,
     ):
+        if waiter_max_attempts is NOTSET:
+            warnings.warn(
+                "The parameter waiter_countdown has been deprecated to standardize "
+                "naming conventions.  Please use waiter_max_attempts instead.  In the "
+                "future this will default to None and defer to the waiter's default value."
+            )
+        if waiter_delay is NOTSET:
+            warnings.warn(
+                "The parameter waiter_check_interval_seconds has been deprecated to "
+                "standardize naming conventions.  Please use waiter_delay instead.  In the "
+                "future this will default to None and defer to the waiter's default value."
+            )
         super().__init__(**kwargs)
         self.editor_id = editor_id
         self.relative_path = relative_path
@@ -204,8 +224,12 @@ def __init__(
         self.wait_for_completion = wait_for_completion
         self.cluster_id = cluster_id
         self.aws_conn_id = aws_conn_id
-        self.waiter_countdown = waiter_countdown
-        self.waiter_check_interval_seconds = waiter_check_interval_seconds
+        self.waiter_max_attempts = (
+            (waiter_countdown // waiter_check_interval_seconds)
+            if waiter_max_attempts is NOTSET
+            else waiter_max_attempts
+        )
+        self.waiter_delay = waiter_check_interval_seconds if waiter_delay is NOTSET else NOTSET
         self.master_instance_security_group_id = master_instance_security_group_id
 
     def execute(self, context: Context):
@@ -233,17 +257,26 @@ def execute(self, context: Context):
         self.log.info("Notebook execution started: %s", response["NotebookExecutionId"])
         notebook_execution_id = response["NotebookExecutionId"]
         if self.wait_for_completion:
-            waiter(
-                get_state_callable=emr_hook.conn.describe_notebook_execution,
-                get_state_args={"NotebookExecutionId": notebook_execution_id},
-                parse_response=["NotebookExecution", "Status"],
-                desired_state={"RUNNING", "FINISHED"},
-                failure_states={"FAILED"},
-                object_type="notebook execution",
-                action="starting",
-                countdown=self.waiter_countdown,
-                check_interval_seconds=self.waiter_check_interval_seconds,
+            emr_hook.get_waiter("notebook_running").wait(
+                NotebookExecutionId=notebook_execution_id,
+                WaiterConfig=prune_dict(
+                    {
+                        "Delay": self.waiter_delay,
+                        "MaxAttempts": self.waiter_max_attempts,
+                    }
+                ),
             )
+
+            # The old Waiter method raised an exception if the notebook
+            # failed, adding that here.  This could maybe be deprecated
+            # later to bring it in line with how other waiters behave.
+            failure_states = {"FAILED"}
+            final_status = emr_hook.conn.describe_notebook_execution(
+                NotebookExecutionId=notebook_execution_id
+            )["NotebookExecution"]["Status"]
+            if final_status in failure_states:
+                raise AirflowException(f"Notebook Execution reached failure state {final_status}.")
+
         return notebook_execution_id
 
 
@@ -260,12 +293,18 @@ class EmrStopNotebookExecutionOperator(BaseOperator):
         to be in a STOPPED or FINISHED state. Defaults to False.
     :param aws_conn_id: aws connection to use.
     :param waiter_countdown: Total amount of time the operator will wait for the notebook to stop.
-        Defaults to 25 * 60 seconds.
+        Defaults to 25 * 60 seconds. (Deprecated.  Please use waiter_max_attempts.)
     :param waiter_check_interval_seconds: Number of seconds between polling the state of the notebook.
-        Defaults to 60 seconds.
+        Defaults to 60 seconds. (Deprecated.  Please use waiter_delay.)
+    :param waiter_max_attempts: Maximum number of tries before failing.
+    :param waiter_delay: Number of seconds between polling the state of the notebook.
     """
 
-    template_fields: Sequence[str] = ("notebook_execution_id",)
+    template_fields: Sequence[str] = (
+        "notebook_execution_id",
+        "waiter_delay",
+        "waiter_max_attempts",
+    )
 
     def __init__(
         self,
@@ -274,30 +313,47 @@ def __init__(
         aws_conn_id: str = "aws_default",
         waiter_countdown: int = 25 * 60,
         waiter_check_interval_seconds: int = 60,
+        # TODO: waiter_max_attempts and waiter_delay should default to None when the other two are deprecated.
+        waiter_max_attempts: int | ArgNotSet = NOTSET,
+        waiter_delay: int | ArgNotSet = NOTSET,
         **kwargs: Any,
     ):
+        if waiter_max_attempts is NOTSET:
+            warnings.warn(
+                "The parameter waiter_countdown has been deprecated to standardize "
+                "naming conventions.  Please use waiter_max_attempts instead.  In the "
+                "future this will default to None and defer to the waiter's default value."
+            )
+        if waiter_delay is NOTSET:
+            warnings.warn(
+                "The parameter waiter_check_interval_seconds has been deprecated to "
+                "standardize naming conventions.  Please use waiter_delay instead.  In the "
+                "future this will default to None and defer to the waiter's default value."
+            )
         super().__init__(**kwargs)
         self.notebook_execution_id = notebook_execution_id
         self.wait_for_completion = wait_for_completion
         self.aws_conn_id = aws_conn_id
-        self.waiter_countdown = waiter_countdown
-        self.waiter_check_interval_seconds = waiter_check_interval_seconds
+        self.waiter_max_attempts = (
+            (waiter_countdown // waiter_check_interval_seconds)
+            if waiter_max_attempts is NOTSET
+            else waiter_max_attempts
+        )
+        self.waiter_delay = waiter_check_interval_seconds if waiter_delay is NOTSET else waiter_delay
 
     def execute(self, context: Context) -> None:
         emr_hook = EmrHook(aws_conn_id=self.aws_conn_id)
         emr_hook.conn.stop_notebook_execution(NotebookExecutionId=self.notebook_execution_id)
 
         if self.wait_for_completion:
-            waiter(
-                get_state_callable=emr_hook.conn.describe_notebook_execution,
-                get_state_args={"NotebookExecutionId": self.notebook_execution_id},
-                parse_response=["NotebookExecution", "Status"],
-                desired_state={"STOPPED", "FINISHED"},
-                failure_states={"FAILED"},
-                object_type="notebook execution",
-                action="stopped",
-                countdown=self.waiter_countdown,
-                check_interval_seconds=self.waiter_check_interval_seconds,
+            emr_hook.get_waiter("notebook_stopped").wait(
+                NotebookExecutionId=self.notebook_execution_id,
+                WaiterConfig=prune_dict(
+                    {
+                        "Delay": self.waiter_delay,
+                        "MaxAttempts": self.waiter_max_attempts,
+                    }
+                ),
             )
 
 
@@ -527,12 +583,18 @@ class EmrCreateJobFlowOperator(BaseOperator):
     :param wait_for_completion: Whether to finish task immediately after creation (False) or wait for jobflow
         completion (True)
     :param waiter_countdown: Max. seconds to wait for jobflow completion (only in combination with
-        wait_for_completion=True, None = no limit)
+        wait_for_completion=True, None = no limit) (Deprecated.  Please use waiter_max_attempts.)
     :param waiter_check_interval_seconds: Number of seconds between polling the jobflow state. Defaults to 60
-        seconds.
+        seconds. (Deprecated.  Please use waiter_delay.)
+    :param waiter_max_attempts: Maximum number of tries before failing.
+    :param waiter_delay: Number of seconds between polling the state of the notebook.
     """
 
-    template_fields: Sequence[str] = ("job_flow_overrides",)
+    template_fields: Sequence[str] = (
+        "job_flow_overrides",
+        "waiter_delay",
+        "waiter_max_attempts",
+    )
     template_ext: Sequence[str] = (".json",)
     template_fields_renderers = {"job_flow_overrides": "json"}
     ui_color = "#f9c915"
@@ -548,16 +610,37 @@ def __init__(
         wait_for_completion: bool = False,
         waiter_countdown: int | None = None,
         waiter_check_interval_seconds: int = 60,
-        **kwargs,
+        # TODO: waiter_max_attempts and waiter_delay should default to None when the other two are deprecated.
+        waiter_max_attempts: int | ArgNotSet = NOTSET,
+        waiter_delay: int | ArgNotSet = NOTSET,
+        **kwargs: Any,
     ):
+        if waiter_max_attempts is NOTSET:
+            warnings.warn(
+                "The parameter waiter_countdown has been deprecated to standardize "
+                "naming conventions.  Please use waiter_max_attempts instead.  In the "
+                "future this will default to None and defer to the waiter's default value."
+            )
+        if waiter_delay is NOTSET:
+            warnings.warn(
+                "The parameter waiter_check_interval_seconds has been deprecated to "
+                "standardize naming conventions.  Please use waiter_delay instead.  In the "
+                "future this will default to None and defer to the waiter's default value."
+            )
         super().__init__(**kwargs)
         self.aws_conn_id = aws_conn_id
         self.emr_conn_id = emr_conn_id
         self.job_flow_overrides = job_flow_overrides or {}
         self.region_name = region_name
+        # waiter_countdown defaults to never timing out, which is not supported by waiters
+        self.waiter_countdown = waiter_countdown or 999
         self.wait_for_completion = wait_for_completion
-        self.waiter_countdown = waiter_countdown
-        self.waiter_check_interval_seconds = waiter_check_interval_seconds
+        self.waiter_max_attempts = (
+            (self.waiter_countdown // waiter_check_interval_seconds)
+            if waiter_max_attempts is NOTSET
+            else waiter_max_attempts
+        )
+        self.waiter_delay = waiter_check_interval_seconds if waiter_delay is NOTSET else waiter_delay
 
         self._job_flow_id: str | None = None
 
@@ -593,19 +676,14 @@ def execute(self, context: Context) -> str | None:
             )
 
             if self.wait_for_completion:
-                # Didn't use a boto-supplied waiter because those don't support waiting for WAITING state.
-                # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/emr.html#waiters
-                waiter(
-                    get_state_callable=self._emr_hook.get_conn().describe_cluster,
-                    get_state_args={"ClusterId": self._job_flow_id},
-                    parse_response=["Cluster", "Status", "State"],
-                    # Cluster will be in WAITING after finishing if KeepJobFlowAliveWhenNoSteps is True
-                    desired_state={"WAITING", "TERMINATED"},
-                    failure_states={"TERMINATED_WITH_ERRORS"},
-                    object_type="job flow",
-                    action="finished",
-                    countdown=self.waiter_countdown,
-                    check_interval_seconds=self.waiter_check_interval_seconds,
+                self._emr_hook.get_waiter("job_flow_waiting").wait(
+                    ClusterId=self._job_flow_id,
+                    WaiterConfig=prune_dict(
+                        {
+                            "Delay": self.waiter_delay,
+                            "MaxAttempts": self.waiter_max_attempts,
+                        }
+                    ),
                 )
 
             return self._job_flow_id
diff --git a/airflow/providers/amazon/aws/utils/waiter.py b/airflow/providers/amazon/aws/utils/waiter.py
index 2aae049838570..6817e31881727 100644
--- a/airflow/providers/amazon/aws/utils/waiter.py
+++ b/airflow/providers/amazon/aws/utils/waiter.py
@@ -41,6 +41,10 @@ def waiter(
     Will call get_state_callable until it reaches the desired_state or the failure_states.
     It will also time out if it waits longer than countdown seconds.
 
+    PLEASE NOTE:  While not yet deprecated, we are moving away from this method
+                  and encourage using the custom boto waiters as explained in
+                  https://github.com/apache/airflow/tree/main/airflow/providers/amazon/aws/waiters
+
     :param get_state_callable: A callable to run until it returns True
     :param get_state_args: Arguments to pass to get_state_callable
     :param parse_response: Dictionary keys to extract state from response of get_state_callable
diff --git a/airflow/providers/amazon/aws/waiters/emr.json b/airflow/providers/amazon/aws/waiters/emr.json
new file mode 100644
index 0000000000000..78afee6b544d1
--- /dev/null
+++ b/airflow/providers/amazon/aws/waiters/emr.json
@@ -0,0 +1,80 @@
+{
+    "version": 2,
+    "waiters": {
+        "notebook_running": {
+            "operation": "DescribeNotebookExecution",
+            "delay": 60,
+            "maxAttempts": 1500,
+            "acceptors": [
+                {
+                    "matcher": "path",
+                    "argument": "notebookexecution.status",
+                    "expected": "RUNNING",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "notebookexecution.status",
+                    "expected": "FINISHED",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "notebookexecution.status",
+                    "expected": "FAILED",
+                    "state": "failure"
+                }
+            ]
+        },
+        "notebook_stopped": {
+            "operation": "DescribeNotebookExecution",
+            "delay": 30,
+            "maxAttempts": 60,
+            "acceptors": [
+                {
+                    "matcher": "path",
+                    "argument": "notebookexecution.status",
+                    "expected": "STOPPED",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "notebookexecution.status",
+                    "expected": "FINISHED",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "notebookexecution.status",
+                    "expected": "FAILED",
+                    "state": "failure"
+                }
+            ]
+        },
+        "job_flow_waiting": {
+            "operation": "DescribeCluster",
+            "delay": 30,
+            "maxAttempts": 60,
+            "acceptors": [
+                {
+                    "matcher": "path",
+                    "argument": "cluster.status",
+                    "expected": "WAITING",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "cluster.status",
+                    "expected": "TERMINATED",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "cluster.status",
+                    "expected": "TERMINATED_WITH_ERRORS",
+                    "state": "failure"
+                }
+            ]
+        }
+    }
+}
diff --git a/tests/providers/amazon/aws/operators/test_eks.py b/tests/providers/amazon/aws/operators/test_eks.py
index 15dd50f103e39..9dee47d8b0462 100644
--- a/tests/providers/amazon/aws/operators/test_eks.py
+++ b/tests/providers/amazon/aws/operators/test_eks.py
@@ -44,6 +44,7 @@
     TASK_ID,
 )
 from tests.providers.amazon.aws.utils.eks_test_utils import convert_keys
+from tests.providers.amazon.aws.utils.test_waiter import assert_expected_waiter_type
 
 CLUSTER_NAME = "cluster1"
 NODEGROUP_NAME = "nodegroup1"
@@ -61,17 +62,6 @@
 }
 
 
-def assert_expected_waiter_type(waiter: mock.MagicMock, expected: str):
-    """
-    There does not appear to be a straight-forward way to assert the type of waiter.
-    Instead, get the class name and check if it contains the expected name.
-
-    :param waiter: A mocked Boto3 Waiter object.
-    :param expected: The expected class name of the Waiter object, for example "ClusterActive".
-    """
-    assert expected in str(type(waiter.call_args[0][0]))
-
-
 class ClusterParams(TypedDict):
     cluster_name: str
     cluster_role_arn: str
diff --git a/tests/providers/amazon/aws/operators/test_emr_create_job_flow.py b/tests/providers/amazon/aws/operators/test_emr_create_job_flow.py
index c76be1475ea80..1920f915b88db 100644
--- a/tests/providers/amazon/aws/operators/test_emr_create_job_flow.py
+++ b/tests/providers/amazon/aws/operators/test_emr_create_job_flow.py
@@ -19,13 +19,16 @@
 
 import os
 from datetime import timedelta
+from unittest import mock
 from unittest.mock import MagicMock, patch
 
+from botocore.waiter import Waiter
 from jinja2 import StrictUndefined
 
 from airflow.models import DAG, DagRun, TaskInstance
 from airflow.providers.amazon.aws.operators.emr import EmrCreateJobFlowOperator
 from airflow.utils import timezone
+from tests.providers.amazon.aws.utils.test_waiter import assert_expected_waiter_type
 from tests.test_utils import AIRFLOW_MAIN_FOLDER
 
 TASK_ID = "test_task"
@@ -34,7 +37,8 @@
 
 DEFAULT_DATE = timezone.datetime(2017, 1, 1)
 
-RUN_JOB_FLOW_SUCCESS_RETURN = {"ResponseMetadata": {"HTTPStatusCode": 200}, "JobFlowId": "j-8989898989"}
+JOB_FLOW_ID = "j-8989898989"
+RUN_JOB_FLOW_SUCCESS_RETURN = {"ResponseMetadata": {"HTTPStatusCode": 200}, "JobFlowId": JOB_FLOW_ID}
 
 TEMPLATE_SEARCHPATH = os.path.join(
     AIRFLOW_MAIN_FOLDER, "tests", "providers", "amazon", "aws", "config_templates"
@@ -158,4 +162,20 @@ def test_execute_returns_job_id(self):
         boto3_session_mock = MagicMock(return_value=emr_session_mock)
 
         with patch("boto3.session.Session", boto3_session_mock):
-            assert self.operator.execute(self.mock_context) == "j-8989898989"
+            assert self.operator.execute(self.mock_context) == JOB_FLOW_ID
+
+    @mock.patch("botocore.waiter.get_service_module_name", return_value="emr")
+    @mock.patch.object(Waiter, "wait")
+    def test_execute_with_wait(self, mock_waiter, _):
+        self.emr_client_mock.run_job_flow.return_value = RUN_JOB_FLOW_SUCCESS_RETURN
+
+        # Mock out the emr_client creator
+        emr_session_mock = MagicMock()
+        emr_session_mock.client.return_value = self.emr_client_mock
+        boto3_session_mock = MagicMock(return_value=emr_session_mock)
+        self.operator.wait_for_completion = True
+
+        with patch("boto3.session.Session", boto3_session_mock):
+            assert self.operator.execute(self.mock_context) == JOB_FLOW_ID
+            mock_waiter.assert_called_once_with(mock.ANY, ClusterId=JOB_FLOW_ID, WaiterConfig=mock.ANY)
+            assert_expected_waiter_type(mock_waiter, "job_flow_waiting")
diff --git a/tests/providers/amazon/aws/operators/test_emr_notebook_execution.py b/tests/providers/amazon/aws/operators/test_emr_notebook_execution.py
index dc8be6a2881d7..313709e18fd23 100644
--- a/tests/providers/amazon/aws/operators/test_emr_notebook_execution.py
+++ b/tests/providers/amazon/aws/operators/test_emr_notebook_execution.py
@@ -20,12 +20,15 @@
 from unittest import mock
 
 import pytest
+from botocore.waiter import Waiter
 
 from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.emr import EmrHook
 from airflow.providers.amazon.aws.operators.emr import (
     EmrStartNotebookExecutionOperator,
     EmrStopNotebookExecutionOperator,
 )
+from tests.providers.amazon.aws.utils.test_waiter import assert_expected_waiter_type
 
 PARAMS = {
     "EditorId": "test_editor",
@@ -44,8 +47,10 @@
 
 
 class TestEmrStartNotebookExecutionOperator:
-    @mock.patch("airflow.providers.amazon.aws.hooks.emr.EmrHook.conn")
-    def test_start_notebook_execution_wait_for_completion(self, mock_conn):
+    @mock.patch("botocore.waiter.get_service_module_name", return_value="emr")
+    @mock.patch.object(EmrHook, "conn")
+    @mock.patch.object(Waiter, "wait")
+    def test_start_notebook_execution_wait_for_completion(self, mock_waiter, mock_conn, _):
         test_execution_id = "test-execution-id"
         mock_conn.start_notebook_execution.return_value = {
             "NotebookExecutionId": test_execution_id,
@@ -53,7 +58,6 @@ def test_start_notebook_execution_wait_for_completion(self, mock_conn):
                 "HTTPStatusCode": 200,
             },
         }
-        mock_conn.describe_notebook_execution.return_value = {"NotebookExecution": {"Status": "FINISHED"}}
 
         op = EmrStartNotebookExecutionOperator(
             task_id="test-id",
@@ -71,7 +75,10 @@ def test_start_notebook_execution_wait_for_completion(self, mock_conn):
         op_response = op.execute(None)
 
         mock_conn.start_notebook_execution.assert_called_once_with(**PARAMS)
-        mock_conn.describe_notebook_execution.assert_called_once_with(NotebookExecutionId=test_execution_id)
+        mock_waiter.assert_called_once_with(
+            mock.ANY, NotebookExecutionId=test_execution_id, WaiterConfig=mock.ANY
+        )
+        assert_expected_waiter_type(mock_waiter, "notebook_running")
         assert op_response == test_execution_id
 
     @mock.patch("airflow.providers.amazon.aws.hooks.emr.EmrHook.conn")
@@ -129,9 +136,11 @@ def test_start_notebook_execution_http_code_fail(self, mock_conn):
 
         mock_conn.start_notebook_execution.assert_called_once_with(**PARAMS)
 
+    @mock.patch("botocore.waiter.get_service_module_name", return_value="emr")
     @mock.patch("time.sleep", return_value=None)
-    @mock.patch("airflow.providers.amazon.aws.hooks.emr.EmrHook.conn")
-    def test_start_notebook_execution_wait_for_completion_multiple_attempts(self, mock_conn, _):
+    @mock.patch.object(EmrHook, "conn")
+    @mock.patch.object(Waiter, "wait")
+    def test_start_notebook_execution_wait_for_completion_multiple_attempts(self, mock_waiter, mock_conn, *_):
         test_execution_id = "test-execution-id"
         mock_conn.start_notebook_execution.return_value = {
             "NotebookExecutionId": test_execution_id,
@@ -139,11 +148,6 @@ def test_start_notebook_execution_wait_for_completion_multiple_attempts(self, mo
                 "HTTPStatusCode": 200,
             },
         }
-        mock_conn.describe_notebook_execution.side_effect = [
-            {"NotebookExecution": {"Status": "PENDING"}},
-            {"NotebookExecution": {"Status": "PENDING"}},
-            {"NotebookExecution": {"Status": "FINISHED"}},
-        ]
 
         op = EmrStartNotebookExecutionOperator(
             task_id="test-id",
@@ -161,12 +165,16 @@ def test_start_notebook_execution_wait_for_completion_multiple_attempts(self, mo
         op_response = op.execute(None)
 
         mock_conn.start_notebook_execution.assert_called_once_with(**PARAMS)
-        mock_conn.describe_notebook_execution.assert_called_with(NotebookExecutionId=test_execution_id)
-        assert mock_conn.describe_notebook_execution.call_count == 3
+        mock_waiter.assert_called_once_with(
+            mock.ANY, NotebookExecutionId=test_execution_id, WaiterConfig=mock.ANY
+        )
+        assert_expected_waiter_type(mock_waiter, "notebook_running")
         assert op_response == test_execution_id
 
-    @mock.patch("airflow.providers.amazon.aws.hooks.emr.EmrHook.conn")
-    def test_start_notebook_execution_wait_for_completion_fail_state(self, mock_conn):
+    @mock.patch("botocore.waiter.get_service_module_name", return_value="emr")
+    @mock.patch.object(EmrHook, "conn")
+    @mock.patch.object(Waiter, "wait")
+    def test_start_notebook_execution_wait_for_completion_fail_state(self, mock_waiter, mock_conn, _):
         test_execution_id = "test-execution-id"
         mock_conn.start_notebook_execution.return_value = {
             "NotebookExecutionId": test_execution_id,
@@ -189,10 +197,15 @@ def test_start_notebook_execution_wait_for_completion_fail_state(self, mock_conn
             tags=PARAMS["Tags"],
             wait_for_completion=True,
         )
+
         with pytest.raises(AirflowException, match=r"Notebook Execution reached failure state FAILED\."):
             op.execute(None)
+
+        mock_waiter.assert_called_once_with(
+            mock.ANY, NotebookExecutionId=test_execution_id, WaiterConfig=mock.ANY
+        )
+        assert_expected_waiter_type(mock_waiter, "notebook_running")
         mock_conn.start_notebook_execution.assert_called_once_with(**PARAMS)
-        mock_conn.describe_notebook_execution.assert_called_once_with(NotebookExecutionId=test_execution_id)
 
 
 class TestStopEmrNotebookExecutionOperator:
@@ -211,10 +224,11 @@ def test_stop_notebook_execution(self, mock_conn):
         mock_conn.stop_notebook_execution.assert_called_once_with(NotebookExecutionId=test_execution_id)
         assert not mock_conn.describe_notebook_execution.called
 
-    @mock.patch("airflow.providers.amazon.aws.hooks.emr.EmrHook.conn")
-    def test_stop_notebook_execution_wait_for_completion(self, mock_conn):
+    @mock.patch("botocore.waiter.get_service_module_name", return_value="emr")
+    @mock.patch.object(EmrHook, "conn")
+    @mock.patch.object(Waiter, "wait")
+    def test_stop_notebook_execution_wait_for_completion(self, mock_waiter, mock_conn, _):
         mock_conn.stop_notebook_execution.return_value = None
-        mock_conn.describe_notebook_execution.return_value = {"NotebookExecution": {"Status": "FINISHED"}}
         test_execution_id = "test-execution-id"
 
         op = EmrStopNotebookExecutionOperator(
@@ -223,32 +237,35 @@ def test_stop_notebook_execution_wait_for_completion(self, mock_conn):
 
         op.execute(None)
         mock_conn.stop_notebook_execution.assert_called_once_with(NotebookExecutionId=test_execution_id)
-        mock_conn.describe_notebook_execution.assert_called_once_with(NotebookExecutionId=test_execution_id)
+        mock_waiter.assert_called_once_with(
+            mock.ANY, NotebookExecutionId=test_execution_id, WaiterConfig=mock.ANY
+        )
+        assert_expected_waiter_type(mock_waiter, "notebook_stopped")
 
-    @mock.patch("airflow.providers.amazon.aws.hooks.emr.EmrHook.conn")
-    def test_stop_notebook_execution_wait_for_completion_fail_state(self, mock_conn):
+    @mock.patch("botocore.waiter.get_service_module_name", return_value="emr")
+    @mock.patch.object(EmrHook, "conn")
+    @mock.patch.object(Waiter, "wait")
+    def test_stop_notebook_execution_wait_for_completion_fail_state(self, mock_waiter, mock_conn, _):
         mock_conn.stop_notebook_execution.return_value = None
-        mock_conn.describe_notebook_execution.return_value = {"NotebookExecution": {"Status": "FAILED"}}
         test_execution_id = "test-execution-id"
 
         op = EmrStopNotebookExecutionOperator(
             task_id="test-id", notebook_execution_id=test_execution_id, wait_for_completion=True
         )
 
-        with pytest.raises(AirflowException, match=r"Notebook Execution reached failure state FAILED."):
-            op.execute(None)
+        op.execute(None)
         mock_conn.stop_notebook_execution.assert_called_once_with(NotebookExecutionId=test_execution_id)
-        mock_conn.describe_notebook_execution.assert_called_once_with(NotebookExecutionId=test_execution_id)
+        mock_waiter.assert_called_once_with(
+            mock.ANY, NotebookExecutionId=test_execution_id, WaiterConfig=mock.ANY
+        )
+        assert_expected_waiter_type(mock_waiter, "notebook_stopped")
 
+    @mock.patch("botocore.waiter.get_service_module_name", return_value="emr")
     @mock.patch("time.sleep", return_value=None)
-    @mock.patch("airflow.providers.amazon.aws.hooks.emr.EmrHook.conn")
-    def test_stop_notebook_execution_wait_for_completion_multiple_attempts(self, mock_conn, _):
+    @mock.patch.object(Waiter, "wait")
+    @mock.patch.object(EmrHook, "conn")
+    def test_stop_notebook_execution_wait_for_completion_multiple_attempts(self, mock_conn, mock_waiter, *_):
         mock_conn.stop_notebook_execution.return_value = None
-        mock_conn.describe_notebook_execution.side_effect = [
-            {"NotebookExecution": {"Status": "PENDING"}},
-            {"NotebookExecution": {"Status": "PENDING"}},
-            {"NotebookExecution": {"Status": "FINISHED"}},
-        ]
         test_execution_id = "test-execution-id"
 
         op = EmrStopNotebookExecutionOperator(
@@ -257,34 +274,32 @@ def test_stop_notebook_execution_wait_for_completion_multiple_attempts(self, moc
 
         op.execute(None)
         mock_conn.stop_notebook_execution.assert_called_once_with(NotebookExecutionId=test_execution_id)
-        mock_conn.describe_notebook_execution.assert_called_with(NotebookExecutionId=test_execution_id)
-        assert mock_conn.describe_notebook_execution.call_count == 3
-
-    @mock.patch("airflow.providers.amazon.aws.operators.emr.waiter")
-    @mock.patch("airflow.providers.amazon.aws.hooks.emr.EmrHook.conn")
-    def test_stop_notebook_execution_waiter_config(self, mock_conn, mock_waiter):
-        mock_waiter.return_value = True
+        mock_waiter.assert_called_once_with(
+            mock.ANY, NotebookExecutionId=test_execution_id, WaiterConfig=mock.ANY
+        )
+        assert_expected_waiter_type(mock_waiter, "notebook_stopped")
 
+    @mock.patch("botocore.waiter.get_service_module_name", return_value="emr")
+    @mock.patch.object(Waiter, "wait")
+    @mock.patch.object(EmrHook, "conn")
+    def test_stop_notebook_execution_waiter_config(self, mock_conn, mock_waiter, _):
         test_execution_id = "test-execution-id"
+        countdown = 400
+        delay = 12
 
         op = EmrStopNotebookExecutionOperator(
             task_id="test-id",
             notebook_execution_id=test_execution_id,
             wait_for_completion=True,
-            waiter_countdown=400,
-            waiter_check_interval_seconds=12,
+            waiter_countdown=countdown,
+            waiter_check_interval_seconds=delay,
         )
 
         op.execute(None)
         mock_conn.stop_notebook_execution.assert_called_once_with(NotebookExecutionId=test_execution_id)
         mock_waiter.assert_called_once_with(
-            get_state_callable=mock_conn.describe_notebook_execution,
-            get_state_args={"NotebookExecutionId": test_execution_id},
-            parse_response=["NotebookExecution", "Status"],
-            desired_state={"STOPPED", "FINISHED"},
-            failure_states={"FAILED"},
-            object_type="notebook execution",
-            action="stopped",
-            countdown=400,
-            check_interval_seconds=12,
+            mock.ANY,
+            NotebookExecutionId=test_execution_id,
+            WaiterConfig={"Delay": delay, "MaxAttempts": countdown // delay},
         )
+        assert_expected_waiter_type(mock_waiter, "notebook_stopped")
diff --git a/tests/providers/amazon/aws/utils/test_waiter.py b/tests/providers/amazon/aws/utils/test_waiter.py
index c4815a500932d..8efe09eb53ca8 100644
--- a/tests/providers/amazon/aws/utils/test_waiter.py
+++ b/tests/providers/amazon/aws/utils/test_waiter.py
@@ -38,6 +38,17 @@ def generate_response(state: str) -> dict[str, Any]:
     }
 
 
+def assert_expected_waiter_type(waiter: mock.MagicMock, expected: str):
+    """
+    There does not appear to be a straight-forward way to assert the type of waiter.
+    Instead, get the class name and check if it contains the expected name.
+
+    :param waiter: A mocked Boto3 Waiter object.
+    :param expected: The expected class name of the Waiter object, for example "ClusterActive".
+    """
+    assert expected in str(type(waiter.call_args[0][0]))
+
+
 class TestWaiter:
     @pytest.mark.parametrize(
         "get_state_responses, fails, expected_exception, expected_num_calls",

From db52e3cb9b2874355169b4dc22ac2c4acf012e4d Mon Sep 17 00:00:00 2001
From: ferruzzi <ferruzzi@amazon.com>
Date: Wed, 1 Mar 2023 13:00:44 -0800
Subject: [PATCH 2/3] PR Suggestions

---
 airflow/providers/amazon/aws/operators/emr.py | 67 +++++++++----------
 1 file changed, 32 insertions(+), 35 deletions(-)

diff --git a/airflow/providers/amazon/aws/operators/emr.py b/airflow/providers/amazon/aws/operators/emr.py
index 19b8656eaa8ab..6bd1df4b3ecff 100644
--- a/airflow/providers/amazon/aws/operators/emr.py
+++ b/airflow/providers/amazon/aws/operators/emr.py
@@ -159,12 +159,13 @@ class EmrStartNotebookExecutionOperator(BaseOperator):
     :param: master_instance_security_group_id: Optional unique ID of an EC2 security
         group to associate with the master instance of the EMR cluster for this notebook execution.
     :param tags: Optional list of key value pair to associate with the notebook execution.
+    :param waiter_max_attempts: Maximum number of tries before failing.
+    :param waiter_delay: Number of seconds between polling the state of the notebook.
+
     :param waiter_countdown: Total amount of time the operator will wait for the notebook to stop.
         Defaults to 25 * 60 seconds. (Deprecated.  Please use waiter_max_attempts.)
     :param waiter_check_interval_seconds: Number of seconds between polling the state of the notebook.
         Defaults to 60 seconds. (Deprecated.  Please use waiter_delay.)
-    :param waiter_max_attempts: Maximum number of tries before failing.
-    :param waiter_delay: Number of seconds between polling the state of the notebook.
     """
 
     template_fields: Sequence[str] = (
@@ -194,11 +195,11 @@ def __init__(
         tags: list | None = None,
         wait_for_completion: bool = False,
         aws_conn_id: str = "aws_default",
+        # TODO: waiter_max_attempts and waiter_delay should default to None when the other two are deprecated.
+        waiter_max_attempts: int | None | ArgNotSet = NOTSET,
+        waiter_delay: int | None | ArgNotSet = NOTSET,
         waiter_countdown: int = 25 * 60,
         waiter_check_interval_seconds: int = 60,
-        # TODO: waiter_max_attempts and waiter_delay should default to None when the other two are deprecated.
-        waiter_max_attempts: int | ArgNotSet = NOTSET,
-        waiter_delay: int | ArgNotSet = NOTSET,
         **kwargs: Any,
     ):
         if waiter_max_attempts is NOTSET:
@@ -207,12 +208,14 @@ def __init__(
                 "naming conventions.  Please use waiter_max_attempts instead.  In the "
                 "future this will default to None and defer to the waiter's default value."
             )
+            waiter_max_attempts = waiter_countdown // waiter_check_interval_seconds
         if waiter_delay is NOTSET:
             warnings.warn(
                 "The parameter waiter_check_interval_seconds has been deprecated to "
                 "standardize naming conventions.  Please use waiter_delay instead.  In the "
                 "future this will default to None and defer to the waiter's default value."
             )
+            waiter_delay = waiter_check_interval_seconds
         super().__init__(**kwargs)
         self.editor_id = editor_id
         self.relative_path = relative_path
@@ -224,12 +227,8 @@ def __init__(
         self.wait_for_completion = wait_for_completion
         self.cluster_id = cluster_id
         self.aws_conn_id = aws_conn_id
-        self.waiter_max_attempts = (
-            (waiter_countdown // waiter_check_interval_seconds)
-            if waiter_max_attempts is NOTSET
-            else waiter_max_attempts
-        )
-        self.waiter_delay = waiter_check_interval_seconds if waiter_delay is NOTSET else NOTSET
+        self.waiter_max_attempts = waiter_max_attempts
+        self.waiter_delay = waiter_delay
         self.master_instance_security_group_id = master_instance_security_group_id
 
     def execute(self, context: Context):
@@ -292,12 +291,13 @@ class EmrStopNotebookExecutionOperator(BaseOperator):
     :param wait_for_completion: If True, the operator will wait for the notebook.
         to be in a STOPPED or FINISHED state. Defaults to False.
     :param aws_conn_id: aws connection to use.
+    :param waiter_max_attempts: Maximum number of tries before failing.
+    :param waiter_delay: Number of seconds between polling the state of the notebook.
+
     :param waiter_countdown: Total amount of time the operator will wait for the notebook to stop.
         Defaults to 25 * 60 seconds. (Deprecated.  Please use waiter_max_attempts.)
     :param waiter_check_interval_seconds: Number of seconds between polling the state of the notebook.
         Defaults to 60 seconds. (Deprecated.  Please use waiter_delay.)
-    :param waiter_max_attempts: Maximum number of tries before failing.
-    :param waiter_delay: Number of seconds between polling the state of the notebook.
     """
 
     template_fields: Sequence[str] = (
@@ -311,11 +311,11 @@ def __init__(
         notebook_execution_id: str,
         wait_for_completion: bool = False,
         aws_conn_id: str = "aws_default",
+        # TODO: waiter_max_attempts and waiter_delay should default to None when the other two are deprecated.
+        waiter_max_attempts: int | None | ArgNotSet = NOTSET,
+        waiter_delay: int | None | ArgNotSet = NOTSET,
         waiter_countdown: int = 25 * 60,
         waiter_check_interval_seconds: int = 60,
-        # TODO: waiter_max_attempts and waiter_delay should default to None when the other two are deprecated.
-        waiter_max_attempts: int | ArgNotSet = NOTSET,
-        waiter_delay: int | ArgNotSet = NOTSET,
         **kwargs: Any,
     ):
         if waiter_max_attempts is NOTSET:
@@ -324,22 +324,20 @@ def __init__(
                 "naming conventions.  Please use waiter_max_attempts instead.  In the "
                 "future this will default to None and defer to the waiter's default value."
             )
+            waiter_max_attempts = waiter_countdown // waiter_check_interval_seconds
         if waiter_delay is NOTSET:
             warnings.warn(
                 "The parameter waiter_check_interval_seconds has been deprecated to "
                 "standardize naming conventions.  Please use waiter_delay instead.  In the "
                 "future this will default to None and defer to the waiter's default value."
             )
+            waiter_delay = waiter_check_interval_seconds
         super().__init__(**kwargs)
         self.notebook_execution_id = notebook_execution_id
         self.wait_for_completion = wait_for_completion
         self.aws_conn_id = aws_conn_id
-        self.waiter_max_attempts = (
-            (waiter_countdown // waiter_check_interval_seconds)
-            if waiter_max_attempts is NOTSET
-            else waiter_max_attempts
-        )
-        self.waiter_delay = waiter_check_interval_seconds if waiter_delay is NOTSET else waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+        self.waiter_delay = waiter_delay
 
     def execute(self, context: Context) -> None:
         emr_hook = EmrHook(aws_conn_id=self.aws_conn_id)
@@ -582,12 +580,13 @@ class EmrCreateJobFlowOperator(BaseOperator):
     :param region_name: Region named passed to EmrHook
     :param wait_for_completion: Whether to finish task immediately after creation (False) or wait for jobflow
         completion (True)
+    :param waiter_max_attempts: Maximum number of tries before failing.
+    :param waiter_delay: Number of seconds between polling the state of the notebook.
+
     :param waiter_countdown: Max. seconds to wait for jobflow completion (only in combination with
         wait_for_completion=True, None = no limit) (Deprecated.  Please use waiter_max_attempts.)
     :param waiter_check_interval_seconds: Number of seconds between polling the jobflow state. Defaults to 60
         seconds. (Deprecated.  Please use waiter_delay.)
-    :param waiter_max_attempts: Maximum number of tries before failing.
-    :param waiter_delay: Number of seconds between polling the state of the notebook.
     """
 
     template_fields: Sequence[str] = (
@@ -608,11 +607,11 @@ def __init__(
         job_flow_overrides: str | dict[str, Any] | None = None,
         region_name: str | None = None,
         wait_for_completion: bool = False,
+        # TODO: waiter_max_attempts and waiter_delay should default to None when the other two are deprecated.
+        waiter_max_attempts: int | None | ArgNotSet = NOTSET,
+        waiter_delay: int | None | ArgNotSet = NOTSET,
         waiter_countdown: int | None = None,
         waiter_check_interval_seconds: int = 60,
-        # TODO: waiter_max_attempts and waiter_delay should default to None when the other two are deprecated.
-        waiter_max_attempts: int | ArgNotSet = NOTSET,
-        waiter_delay: int | ArgNotSet = NOTSET,
         **kwargs: Any,
     ):
         if waiter_max_attempts is NOTSET:
@@ -621,26 +620,24 @@ def __init__(
                 "naming conventions.  Please use waiter_max_attempts instead.  In the "
                 "future this will default to None and defer to the waiter's default value."
             )
+            # waiter_countdown defaults to never timing out, which is not supported
+            # by boto waiters, so we will set it here to "a very long time" for now.
+            waiter_max_attempts = (waiter_countdown or 999) // waiter_check_interval_seconds
         if waiter_delay is NOTSET:
             warnings.warn(
                 "The parameter waiter_check_interval_seconds has been deprecated to "
                 "standardize naming conventions.  Please use waiter_delay instead.  In the "
                 "future this will default to None and defer to the waiter's default value."
             )
+            waiter_delay = waiter_check_interval_seconds
         super().__init__(**kwargs)
         self.aws_conn_id = aws_conn_id
         self.emr_conn_id = emr_conn_id
         self.job_flow_overrides = job_flow_overrides or {}
         self.region_name = region_name
-        # waiter_countdown defaults to never timing out, which is not supported by waiters
-        self.waiter_countdown = waiter_countdown or 999
         self.wait_for_completion = wait_for_completion
-        self.waiter_max_attempts = (
-            (self.waiter_countdown // waiter_check_interval_seconds)
-            if waiter_max_attempts is NOTSET
-            else waiter_max_attempts
-        )
-        self.waiter_delay = waiter_check_interval_seconds if waiter_delay is NOTSET else waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+        self.waiter_delay = waiter_delay
 
         self._job_flow_id: str | None = None
 

From b0421f0d946e678ddeb5751cc80876eadb6db05f Mon Sep 17 00:00:00 2001
From: ferruzzi <ferruzzi@amazon.com>
Date: Wed, 1 Mar 2023 14:03:11 -0800
Subject: [PATCH 3/3] Add custom waiter check to unit tests

---
 tests/providers/amazon/aws/hooks/test_emr.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tests/providers/amazon/aws/hooks/test_emr.py b/tests/providers/amazon/aws/hooks/test_emr.py
index f7f0c37635a2d..0c927bdf3a399 100644
--- a/tests/providers/amazon/aws/hooks/test_emr.py
+++ b/tests/providers/amazon/aws/hooks/test_emr.py
@@ -28,6 +28,13 @@
 
 
 class TestEmrHook:
+    def test_service_waiters(self):
+        hook = EmrHook(aws_conn_id=None)
+        official_waiters = hook.conn.waiter_names
+        custom_waiters = ["job_flow_waiting", "notebook_running", "notebook_stopped"]
+
+        assert hook.list_waiters() == [*official_waiters, *custom_waiters]
+
     @mock_emr
     def test_get_conn_returns_a_boto3_connection(self):
         hook = EmrHook(aws_conn_id="aws_default", region_name="ap-southeast-2")
