From 8ec93101391526f66d7dc495be509bb27a702673 Mon Sep 17 00:00:00 2001 From: bugraoz93 Date: Sun, 11 Aug 2024 23:45:29 +0200 Subject: [PATCH 1/2] Pass serialized parameter for dag_maker --- tests/decorators/test_python.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/decorators/test_python.py b/tests/decorators/test_python.py index fb8c75b72b4c3..efd780c639210 100644 --- a/tests/decorators/test_python.py +++ b/tests/decorators/test_python.py @@ -1000,7 +1000,7 @@ def down(a: str): def test_teardown_trigger_rule_selective_application(dag_maker, session): - with dag_maker(session=session) as dag: + with dag_maker(session=session, serialized=True) as dag: @dag.task def my_work(): @@ -1025,7 +1025,7 @@ def my_teardown(): def test_teardown_trigger_rule_override_behavior(dag_maker, session): - with dag_maker(session=session) as dag: + with dag_maker(session=session, serialized=True) as dag: @dag.task(trigger_rule=TriggerRule.ONE_SUCCESS) def my_work(): From dd9cf555db53c6f5bf9540d9fbbc609e71476766 Mon Sep 17 00:00:00 2001 From: bugraoz93 Date: Mon, 12 Aug 2024 00:37:28 +0200 Subject: [PATCH 2/2] Serialisation of object is on __exit__ moving out the dag definition out of dag_maker context --- tests/decorators/test_python.py | 80 +++++++++++++++++---------------- 1 file changed, 41 insertions(+), 39 deletions(-) diff --git a/tests/decorators/test_python.py b/tests/decorators/test_python.py index efd780c639210..9d2b9a14c82b4 100644 --- a/tests/decorators/test_python.py +++ b/tests/decorators/test_python.py @@ -1000,50 +1000,52 @@ def down(a: str): def test_teardown_trigger_rule_selective_application(dag_maker, session): - with dag_maker(session=session, serialized=True) as dag: - - @dag.task - def my_work(): - return "abc" - - @setup - @dag.task - def my_setup(): - return "abc" - - @teardown - @dag.task - def my_teardown(): - return "abc" - - work_task = my_work() - setup_task = my_setup() - teardown_task = my_teardown() + with dag_maker(session=session, serialized=True) as created_dag: + dag = created_dag + + @dag.task + def my_work(): + return "abc" + + @setup + @dag.task + def my_setup(): + return "abc" + + @teardown + @dag.task + def my_teardown(): + return "abc" + + work_task = my_work() + setup_task = my_setup() + teardown_task = my_teardown() assert work_task.operator.trigger_rule == TriggerRule.ALL_SUCCESS assert setup_task.operator.trigger_rule == TriggerRule.ALL_SUCCESS assert teardown_task.operator.trigger_rule == TriggerRule.ALL_DONE_SETUP_SUCCESS def test_teardown_trigger_rule_override_behavior(dag_maker, session): - with dag_maker(session=session, serialized=True) as dag: - - @dag.task(trigger_rule=TriggerRule.ONE_SUCCESS) - def my_work(): - return "abc" - - @setup - @dag.task(trigger_rule=TriggerRule.ONE_SUCCESS) - def my_setup(): - return "abc" - - @teardown - @dag.task(trigger_rule=TriggerRule.ONE_SUCCESS) - def my_teardown(): - return "abc" - - work_task = my_work() - setup_task = my_setup() - with pytest.raises(Exception, match="Trigger rule not configurable for teardown tasks."): - my_teardown() + with dag_maker(session=session, serialized=True) as created_dag: + dag = created_dag + + @dag.task(trigger_rule=TriggerRule.ONE_SUCCESS) + def my_work(): + return "abc" + + @setup + @dag.task(trigger_rule=TriggerRule.ONE_SUCCESS) + def my_setup(): + return "abc" + + @teardown + @dag.task(trigger_rule=TriggerRule.ONE_SUCCESS) + def my_teardown(): + return "abc" + + work_task = my_work() + setup_task = my_setup() + with pytest.raises(Exception, match="Trigger rule not configurable for teardown tasks."): + my_teardown() assert work_task.operator.trigger_rule == TriggerRule.ONE_SUCCESS assert setup_task.operator.trigger_rule == TriggerRule.ONE_SUCCESS