From f20391c86021dc1848f7facff3a655416cf0e8d4 Mon Sep 17 00:00:00 2001 From: Darshan Gorasiya Date: Tue, 3 Feb 2026 00:02:24 +0000 Subject: [PATCH 1/2] Fix(Amazon): Allow configuring wait_policy in EmrCreateJobFlowOperator This restores the ability to specify a WaitPolicy (e.g. WAIT_FOR_STEPS_COMPLETION) instead of forcing the default 'running' state check. Closes #61180. --- .../providers/amazon/aws/operators/emr.py | 23 ++++--------------- 1 file changed, 5 insertions(+), 18 deletions(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py index 48e8773c16a09..9402c1e05930a 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py @@ -701,25 +701,12 @@ 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)." - ) - - self.wait_for_completion = wait_policy in ( - WaitPolicy.WAIT_FOR_COMPLETION, - WaitPolicy.WAIT_FOR_STEPS_COMPLETION, - ) + if wait_for_completion is False: + raise ValueError("Cannot specify wait_policy with wait_for_completion=False") + self.wait_for_completion = True @property def _hook_parameters(self): @@ -758,7 +745,7 @@ 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] + waiter_name = WAITER_POLICY_NAME_MAPPING[self.wait_policy or WaitPolicy.WAIT_FOR_COMPLETION] if self.deferrable: self.defer( From baea7da05b1a7738f422ea7d07deefa25441361f Mon Sep 17 00:00:00 2001 From: Darshan Gorasiya Date: Tue, 3 Feb 2026 00:05:00 +0000 Subject: [PATCH 2/2] Test(Amazon): Add test for wait_policy in EmrCreateJobFlowOperator Verifies that the specified wait_policy is respected and the deprecated warning is removed. --- .../aws/operators/test_emr_create_job_flow.py | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_create_job_flow.py b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_create_job_flow.py index e49f7ba775806..941bba19d2e69 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_create_job_flow.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_create_job_flow.py @@ -254,10 +254,20 @@ 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"): - EmrCreateJobFlowOperator( - task_id=TASK_ID, - wait_policy=WaitPolicy.WAIT_FOR_COMPLETION, - ) + @mock.patch("botocore.waiter.get_service_module_name", return_value="emr") + @mock.patch.object(Waiter, "wait") + def test_execute_with_wait_policy(self, mock_waiter, _, mocked_hook_client): + mocked_hook_client.run_job_flow.return_value = RUN_JOB_FLOW_SUCCESS_RETURN + + # Test that providing wait_policy uses the correct waiter + op = EmrCreateJobFlowOperator( + task_id="test_wait_policy", + aws_conn_id="aws_default", + emr_conn_id="emr_default", + wait_policy=WaitPolicy.WAIT_FOR_STEPS_COMPLETION, + ) + + op.execute(self.mock_context) + + mock_waiter.assert_called_once_with(mock.ANY, ClusterId=JOB_FLOW_ID, WaiterConfig=mock.ANY) + assert_expected_waiter_type(mock_waiter, WAITER_POLICY_NAME_MAPPING[WaitPolicy.WAIT_FOR_STEPS_COMPLETION])