seanghaeli commented on code in PR #68917:
URL: https://github.com/apache/airflow/pull/68917#discussion_r4021331168
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -1610,30 +1610,136 @@ def
test_dagrun_deadline_does_not_warn_for_average_runtime_without_history(
"log_level": "warning",
} not in caplog
+ @pytest.mark.parametrize(
+ ("interval", "failure"),
+ [
+ pytest.param(
+ VariableInterval("missing_key"),
+ mock.patch.object(Variable, "get", side_effect=KeyError),
+ id="unresolvable_interval",
+ ),
+ pytest.param(
+ datetime.timedelta(hours=1),
+ mock.patch(
+
"airflow.serialization.definitions.dag.decode_deadline_alert",
+ autospec=True,
+ side_effect=ValueError("corrupt deadline alert blob"),
+ ),
+ id="undecodable_alert",
+ ),
+ pytest.param(
+ datetime.timedelta(hours=1),
+ mock.patch.object(
+ SerializedReferenceModels.FixedDatetimeDeadline,
+ "evaluate_with",
+ autospec=True,
+ side_effect=RuntimeError("evaluate_with failed"),
+ ),
+ id="unevaluable_reference",
+ ),
+ ],
+ )
+ @mock.patch("airflow._shared.observability.metrics.stats.incr")
@mock.patch.object(Deadline, "prune_deadlines")
- def test_dagrun_deadline_variable_interval_missing_variable_fails(self, _,
session, deadline_test_dag):
+ def test_dagrun_deadline_failure_does_not_abort_dagrun(
+ self, _, mock_stats_incr, interval, failure, session, deadline_test_dag
+ ):
+ future_date = datetime.datetime(2037, 1, 1,
tzinfo=datetime.timezone.utc)
- with mock.patch.object(
- Variable,
- "get",
- side_effect=KeyError,
+ scheduler_dag = deadline_test_dag(
+ deadline=DeadlineAlert(
+ reference=DeadlineReference.FIXED_DATETIME(future_date),
+ interval=interval,
+ callback=AsyncCallback(empty_callback_for_deadline),
+ ),
+ )
+
+ with (
+ conf_vars({("core", "multi_team"): "true"}),
+ mock.patch("airflow.models.dag.DagModel.get_team_name",
return_value="team_alpha"),
+ failure,
):
- future_date = datetime.datetime.now() +
datetime.timedelta(days=365)
+ dag_run = self.create_dag_run(
+ dag=scheduler_dag,
+ task_states={"task_1": TaskInstanceState.SUCCESS},
+ session=session,
+ )
+
+ assert dag_run is not None
+ assert session.execute(select(Deadline)).scalars().one_or_none() is
None
+ mock_stats_incr.assert_any_call(
+ "deadline_alerts.deadline_creation_failed",
+ tags={"dag_id": scheduler_dag.dag_id, "team_name": "team_alpha"},
+ )
- scheduler_dag = deadline_test_dag(
- deadline=DeadlineAlert(
- reference=DeadlineReference.FIXED_DATETIME(future_date),
- interval=VariableInterval("missing_key"),
- callback=AsyncCallback(empty_callback_for_deadline),
+ @mock.patch("airflow._shared.observability.metrics.stats.incr")
+ @mock.patch.object(Deadline, "prune_deadlines")
+ @mock.patch.object(Variable, "get")
+ def test_dagrun_deadline_failure_that_detaches_orm_objects_still_skips(
+ self, mock_variable_get, _, mock_stats_incr, session, deadline_test_dag
+ ):
+ """A failing interval resolution can roll back and close the caller's
session underneath
+ the handler (``create_session`` reusing the scoped session does
exactly that), detaching
+ every ORM instance. The skip path must still log and count without
touching ORM state."""
+ scheduler_dag = deadline_test_dag(
+ deadline=DeadlineAlert(
+ reference=DeadlineReference.FIXED_DATETIME(
+ datetime.datetime(2037, 1, 1, tzinfo=datetime.timezone.utc)
),
- )
+ interval=VariableInterval("missing_key"),
+ callback=AsyncCallback(empty_callback_for_deadline),
+ ),
+ )
+ # The alert rows must be persistent from an earlier transaction (as in
the live
+ # scheduler, where the dag processor committed them) so the rollback
below expires
+ # them and the close detaches them instead of returning them to
transient.
+ session.commit()
- with pytest.raises(ValueError, match="not found"):
- self.create_dag_run(
- dag=scheduler_dag,
- task_states={"task_1": TaskInstanceState.SUCCESS},
- session=session,
- )
+ def teardown_session_and_raise(*args, **kwargs):
+ session.rollback()
+ session.close()
+ raise KeyError("missing_key")
+
+ mock_variable_get.side_effect = teardown_session_and_raise
+
+ dag_run = self.create_dag_run(dag=scheduler_dag, session=session)
+
+ assert dag_run is not None
+ mock_stats_incr.assert_any_call(
+ "deadline_alerts.deadline_creation_failed",
+ tags={"dag_id": scheduler_dag.dag_id},
+ )
+
+ @pytest.mark.parametrize(
+ ("interval", "expect_deadline"),
+ [
+ pytest.param(datetime.timedelta(seconds=5), True,
id="deadline_created"),
+ pytest.param(VariableInterval("missing_key"), False,
id="deadline_skipped"),
+ ],
+ )
+ @mock.patch.object(Deadline, "prune_deadlines")
+ def test_dagrun_deadline_handling_does_not_commit(
Review Comment:
Added this + test
--
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]
