Skip to content

Commit 0bd0438

Browse files
committed
Cover VariableInterval resolution in the deadline session-safety test
Both existing cases mock Variable.get to raise KeyError, so neither ever resolves a variable successfully and both pass whether or not the scheduler's session is forwarded into resolve(). Commit a real Variable row before the guard opens so the lookup goes through MetastoreBackend, and assert the deadline lands at the resolved offset.
1 parent 2c9cd70 commit 0bd0438

1 file changed

Lines changed: 22 additions & 8 deletions

File tree

‎airflow-core/tests/unit/models/test_dagrun.py‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
import datetime
2121
from collections import defaultdict
2222
from collections.abc import Mapping
23-
from contextlib import contextmanager
23+
from contextlib import contextmanager, nullcontext
2424
from functools import partial, reduce
2525
from typing import TYPE_CHECKING
2626
from unittest import mock
@@ -1711,18 +1711,20 @@ def teardown_session_and_raise(*args, **kwargs):
17111711
)
17121712

17131713
@pytest.mark.parametrize(
1714-
("interval", "expect_deadline"),
1714+
("interval", "expect_deadline", "variable_value"),
17151715
[
1716-
pytest.param(datetime.timedelta(seconds=5), True, id="deadline_created"),
1717-
pytest.param(VariableInterval("missing_key"), False, id="deadline_skipped"),
1716+
pytest.param(datetime.timedelta(seconds=5), True, None, id="deadline_created"),
1717+
pytest.param(VariableInterval("missing_key"), False, None, id="deadline_skipped"),
1718+
pytest.param(VariableInterval("resolvable_key"), True, "4200", id="variable_interval_resolved"),
17181719
],
17191720
)
17201721
@mock.patch.object(Deadline, "prune_deadlines")
17211722
def test_dagrun_deadline_handling_does_not_commit(
1722-
self, _, interval, expect_deadline, session, deadline_test_dag
1723+
self, _, interval, expect_deadline, variable_value, session, deadline_test_dag
17231724
):
1724-
"""DagRuns are created under ``prohibit_commit``, so neither creating a deadline nor
1725-
skipping one whose interval will not resolve may commit the scheduler's session."""
1725+
"""DagRuns are created under ``prohibit_commit``, so creating a deadline, skipping one whose
1726+
interval will not resolve, and resolving a ``VariableInterval`` out of the metadata database
1727+
all have to leave the scheduler's session uncommitted."""
17261728
future_date = datetime.datetime(2037, 1, 1, tzinfo=datetime.timezone.utc)
17271729
scheduler_dag = deadline_test_dag(
17281730
deadline=DeadlineAlert(
@@ -1732,14 +1734,26 @@ def test_dagrun_deadline_handling_does_not_commit(
17321734
),
17331735
)
17341736

1737+
if variable_value is None:
1738+
variable_lookup = mock.patch.object(Variable, "get", side_effect=KeyError)
1739+
else:
1740+
# A real row committed before the guard opens, so resolve() has to go through
1741+
# MetastoreBackend. Drop the session forwarding and the backend opens the same scoped
1742+
# session and commits it, which the guard rejects and the lookup then reports as missing.
1743+
Variable.set(key=interval.key, value=variable_value, session=session)
1744+
session.commit()
1745+
variable_lookup = nullcontext()
1746+
17351747
# No task_states: the helper's get_task_instance() opens its own session and would trip the guard.
1736-
with prohibit_commit(session), mock.patch.object(Variable, "get", side_effect=KeyError):
1748+
with prohibit_commit(session), variable_lookup:
17371749
dag_run = self.create_dag_run(dag=scheduler_dag, session=session)
17381750
session.flush()
17391751

17401752
assert dag_run is not None
17411753
deadline = session.execute(select(Deadline)).scalars().one_or_none()
17421754
assert (deadline is not None) is expect_deadline
1755+
if variable_value is not None:
1756+
assert deadline.deadline_time == future_date + datetime.timedelta(seconds=int(variable_value))
17431757

17441758

17451759
@pytest.mark.parametrize(

0 commit comments

Comments
 (0)