amoghrajesh commented on code in PR #73030:
URL: https://github.com/apache/airflow/pull/73030#discussion_r4060007571


##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -1196,6 +1196,65 @@ def execute(self, context):
     assert counted.count("operator_failures") == 1
 
 
+def test_retry_policy_fail_persists_reason(create_runtime_ti, 
mock_supervisor_comms):
+    class _AlwaysFails(BaseOperator):
+        def execute(self, context):
+            raise RuntimeError("boom")
+
+    task = _AlwaysFails(
+        task_id="fail_with_reason",
+        retry_policy=ExceptionRetryPolicy(
+            rules=[RetryRule(exception=RuntimeError, action=RetryAction.FAIL, 
reason="do not retry")]
+        ),
+    )
+    ti = create_runtime_ti(task=task, should_retry=True)
+
+    state, msg, error = run(ti, ti.get_template_context(), mock.MagicMock())
+
+    assert state == TaskInstanceState.FAILED
+    assert isinstance(msg, TaskState)
+    assert msg.retry_reason == "do not retry"
+
+
+def 
test_retry_policy_retry_exhausted_persists_combined_reason(create_runtime_ti, 
mock_supervisor_comms):
+    """A policy-chosen RETRY that hits an exhausted budget still fails, with 
both reasons recorded."""
+
+    class _AlwaysFails(BaseOperator):
+        def execute(self, context):
+            raise RuntimeError("boom")
+
+    task = _AlwaysFails(
+        task_id="retry_exhausted",
+        retry_policy=ExceptionRetryPolicy(
+            rules=[RetryRule(exception=RuntimeError, action=RetryAction.RETRY, 
reason="rate limit")]
+        ),
+    )
+    ti = create_runtime_ti(task=task, try_number=2, max_tries=2, 
should_retry=False)
+
+    state, msg, error = run(ti, ti.get_template_context(), mock.MagicMock())
+
+    assert state == TaskInstanceState.FAILED
+    assert isinstance(msg, TaskState)
+    assert msg.retry_reason == "rate limit; retries exhausted (2 of 2)"

Review Comment:
   Fixed on #73027, and taken a bit further than the suggestion: rather than 
hardcoding `try_number=3, max_tries=2`, the operator now carries `retries=2` 
and the fixture derives `max_tries` and `should_retry` itself. That keeps the 
triple inside the reachable space by construction rather than by convention, so 
it can't drift back out. Expected string is `(3 of 3)`.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to