2020import datetime
2121from collections import defaultdict
2222from collections .abc import Mapping
23- from contextlib import contextmanager
23+ from contextlib import contextmanager , nullcontext
2424from functools import partial , reduce
2525from typing import TYPE_CHECKING
2626from 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