Fix hardcoded waiter logic in EmrCreateJobFlowOperator

The operator was incorrectly ignoring the `wait_policy` argument and always defaulting to waiting for cluster completion.

This change ensures the `wait_policy` is correctly persisted and used to select the appropriate waiter (e.g., for step completion), fixing the hardcoded behavior.

Author: henry3260

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

@@ -18,13 +18,11 @@
 from __future__ import annotations
 
 import ast
-import warnings
 from collections.abc import Sequence
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 from uuid import uuid4
 
-from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.emr import EmrContainerHook, EmrHook, EmrServerlessHook
 from airflow.providers.amazon.aws.links.emr import (
     EmrClusterLink,
@@ -657,7 +655,7 @@ class EmrCreateJobFlowOperator(AwsBaseOperator[EmrHook]):
     :param wait_for_completion: Whether to finish task immediately after creation (False) or wait for jobflow
         completion (True)
         (default: None)
-    :param wait_policy: Deprecated. Use `wait_for_completion` instead. Whether to finish the task immediately after creation (None) or:
+    :param wait_policy: Whether to finish the task immediately after creation (None) or:
         - wait for the jobflow completion (WaitPolicy.WAIT_FOR_COMPLETION)
         - wait for the jobflow completion and cluster to terminate (WaitPolicy.WAIT_FOR_STEPS_COMPLETION)
         (default: None)
@@ -701,19 +699,13 @@ def __init__(
         self.waiter_max_attempts = waiter_max_attempts or 60
         self.waiter_delay = waiter_delay or 60
         self.deferrable = deferrable
+        self.wait_policy = wait_policy
 
         if wait_policy is not None:
-            warnings.warn(
-                "`wait_policy` parameter is deprecated and will be removed in a future release; "
-                "please use `wait_for_completion` (bool) instead.",
-                AirflowProviderDeprecationWarning,
-                stacklevel=2,
-            )
-
             if wait_for_completion is not None:
                 raise ValueError(
-                    "Cannot specify both `wait_for_completion` and deprecated `wait_policy`. "
-                    "Please use `wait_for_completion` (bool)."
+                    "Cannot specify both `wait_for_completion` and `wait_policy`. "
+                    "Use `wait_policy` if you need to control the specific waiting behavior."
                 )
 
             self.wait_for_completion = wait_policy in (
@@ -758,15 +750,24 @@ def execute(self, context: Context) -> str | None:
                 log_uri=get_log_uri(emr_client=self.hook.conn, job_flow_id=self._job_flow_id),
             )
         if self.wait_for_completion:
-            waiter_name = WAITER_POLICY_NAME_MAPPING[WaitPolicy.WAIT_FOR_COMPLETION]
+            # Determine which waiter to use. Prefer explicit wait_policy when provided,
+            # otherwise default to WAIT_FOR_COMPLETION.
+            wp = self.wait_policy
+            if wp is not None:
+                waiter_name = WAITER_POLICY_NAME_MAPPING[wp]
+            else:
+                waiter_name = WAITER_POLICY_NAME_MAPPING[WaitPolicy.WAIT_FOR_COMPLETION]
 
             if self.deferrable:
+                # Pass the selected waiter_name to the trigger so deferrable mode waits
+                # according to the requested policy as well.
                 self.defer(
                     trigger=EmrCreateJobFlowTrigger(
                         job_flow_id=self._job_flow_id,
                         aws_conn_id=self.aws_conn_id,
                         waiter_delay=self.waiter_delay,
                         waiter_max_attempts=self.waiter_max_attempts,
+                        waiter_name=waiter_name,
                     ),
                     method_name="execute_complete",
                     # timeout is set to ensure that if a trigger dies, the timeout does not restart

providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py

@@ -82,10 +82,11 @@ def __init__(
         aws_conn_id: str | None = None,
         waiter_delay: int = 30,
         waiter_max_attempts: int = 60,
+        waiter_name: str = "job_flow_waiting",
     ):
         super().__init__(
             serialized_fields={"job_flow_id": job_flow_id},
-            waiter_name="job_flow_waiting",
+            waiter_name=waiter_name,
             waiter_args={"ClusterId": job_flow_id},
             failure_message="JobFlow creation failed",
             status_message="JobFlow creation in progress",

providers/amazon/tests/unit/amazon/aws/operators/test_emr_create_job_flow.py

@@ -26,7 +26,6 @@
 from botocore.waiter import Waiter
 from jinja2 import StrictUndefined
 
-from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.models import DAG, DagRun, TaskInstance
 from airflow.providers.amazon.aws.operators.emr import EmrCreateJobFlowOperator
 from airflow.providers.amazon.aws.triggers.emr import EmrCreateJobFlowTrigger
@@ -254,10 +253,25 @@ def test_create_job_flow_deferrable_no_wait(self, mocked_hook_client):
     def test_template_fields(self):
         validate_template_fields(self.operator)
 
-    def test_wait_policy_deprecation_warning(self):
-        """Test that using wait_policy raises a deprecation warning."""
-        with pytest.warns(AirflowProviderDeprecationWarning, match="`wait_policy` parameter is deprecated"):
+    def test_wait_policy_behavior(self):
+        """Test that using wait_policy sets the operator attributes correctly."""
+        op = EmrCreateJobFlowOperator(
+            task_id=TASK_ID,
+            wait_policy=WaitPolicy.WAIT_FOR_COMPLETION,
+        )
+        # wait_policy should be stored on the instance
+        assert getattr(op, "wait_policy") == WaitPolicy.WAIT_FOR_COMPLETION
+        # passing WAIT_FOR_COMPLETION should enable wait_for_completion
+        assert op.wait_for_completion is True
+
+    def test_cannot_specify_both_wait_for_completion_and_wait_policy(self):
+        """Passing both wait_for_completion and wait_policy should raise ValueError."""
+        with pytest.raises(
+            ValueError,
+            match=r"Cannot specify both `wait_for_completion` and `wait_policy`",
+        ):
             EmrCreateJobFlowOperator(
                 task_id=TASK_ID,
-                wait_policy=WaitPolicy.WAIT_FOR_COMPLETION,
+                wait_for_completion=True,
+                wait_policy=WaitPolicy.WAIT_FOR_STEPS_COMPLETION,
             )