vatsrahul1001 commented on code in PR #69243:
URL: https://github.com/apache/airflow/pull/69243#discussion_r3518236670
##########
airflow-core/tests/unit/serialization/test_serialized_objects.py:
##########
@@ -1309,6 +1309,88 @@ def
test_has_retry_policy_flag_false_when_no_policy(self):
task = deserialized.task_dict["op_no_policy"]
assert task.has_retry_policy is False
+ def test_mapped_task_retry_policy_serializes_as_flag(self):
+ """A mapped task's retry_policy must serialize as has_retry_policy,
not the object."""
+ from airflow.sdk import DAG # module-level DAG is
airflow.models.dag.DAG
+ from airflow.sdk.definitions.retry_policy import ExceptionRetryPolicy,
RetryAction, RetryRule
+
+ policy = ExceptionRetryPolicy(
+ rules=[RetryRule(exception=ValueError, action=RetryAction.FAIL,
reason="bad data")],
+ )
+
+ with DAG(dag_id="test_mapped_retry_policy_ser",
start_date=DEFAULT_DATE) as dag:
+
+ @task(retries=3, retry_policy=policy)
+ def mapped(x):
+ return x
+
+ mapped.expand(x=[1, 2, 3])
+
+ serialized = DagSerialization.serialize_dag(dag)
+ # The RetryPolicy object must never be embedded -- str(obj) leaks a
memory address.
+ assert "ExceptionRetryPolicy object at 0x" not in
json.dumps(serialized)
+
+ deserialized = DagSerialization.deserialize_dag(serialized)
+ assert deserialized.task_dict["mapped"].has_retry_policy is True
+
+ def test_mapped_task_retry_policy_serialization_is_deterministic(self):
+ """Serializing the same mapped-task-with-policy DAG twice yields
identical output.
+
+ Regression test: the RetryPolicy object was serialized via str(obj),
embedding a
+ per-process memory address, so every re-parse produced a different
serialized DAG
+ (and a spurious new DagVersion).
+ """
+ from airflow.sdk import DAG # module-level DAG is
airflow.models.dag.DAG
+ from airflow.sdk.definitions.retry_policy import ExceptionRetryPolicy,
RetryAction, RetryRule
+
+ def build():
+ policy = ExceptionRetryPolicy(
+ rules=[RetryRule(exception=ValueError,
action=RetryAction.FAIL)],
+ )
+ with DAG(dag_id="test_mapped_retry_policy_determinism",
start_date=DEFAULT_DATE) as dag:
+
+ @task(retries=3, retry_policy=policy)
+ def mapped(x):
+ return x
+
+ mapped.expand(x=[1, 2, 3])
+ return dag
+
+ assert DagSerialization.serialize_dag(build()) ==
DagSerialization.serialize_dag(build())
+
+ def test_dag_default_args_retry_policy_serializes_as_flag(self):
+ """A retry_policy in DAG default_args must not leak the object into
serialized default_args.
+
+ Regression test: serialize_dag serializes the raw default_args dict,
so a RetryPolicy
+ there hit the str(obj) fallback (memory address -> new DagVersion
every parse) even
+ though each task's has_retry_policy was set correctly.
+ """
+ from airflow.sdk import DAG # module-level DAG is
airflow.models.dag.DAG
+ from airflow.sdk.definitions.retry_policy import ExceptionRetryPolicy,
RetryAction, RetryRule
+
+ def build():
+ policy = ExceptionRetryPolicy(
+ rules=[RetryRule(exception=ValueError,
action=RetryAction.FAIL)],
+ )
+ with DAG(
+ dag_id="test_default_args_retry_policy",
+ start_date=DEFAULT_DATE,
+ default_args={"retry_policy": policy},
+ ) as dag:
+
+ @task
+ def plain():
+ return 1
+
+ plain()
+ return dag
+
+ serialized = DagSerialization.serialize_dag(build())
+ # The RetryPolicy object must never be embedded in serialized
default_args.
+ assert "ExceptionRetryPolicy object at 0x" not in
json.dumps(serialized)
+ # Deterministic across independent parses (no embedded memory address).
+ assert DagSerialization.serialize_dag(build()) ==
DagSerialization.serialize_dag(build())
+
Review Comment:
Added `test_mapped_task_no_retry_policy_flag_false` as suggested — good
call, the mapped path resolves the flag through
`_get_partial_kwargs_or_operator_default` rather than the dataclass default, so
the false case deserved its own coverage. Thanks for the review!
--
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]
