This is an automated email from the ASF dual-hosted git repository. potiuk pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push: new 6229c1c98a Fix DateTimeOperator tests in Database Isolation Mode (#41271) 6229c1c98a is described below commit 6229c1c98a2a78d00255d6b6c5b70d032cff9b80 Author: Jens Scheffler <95105677+jsche...@users.noreply.github.com> AuthorDate: Tue Aug 6 09:17:09 2024 +0200 Fix DateTimeOperator tests in Database Isolation Mode (#41271) * Fix DateTimeOperator tests in Database Isolation Mode * Review Feedback --- tests/operators/test_datetime.py | 41 ++++++++++++++++++++-------------------- 1 file changed, 20 insertions(+), 21 deletions(-) diff --git a/tests/operators/test_datetime.py b/tests/operators/test_datetime.py index 485897acc6..e4294a8856 100644 --- a/tests/operators/test_datetime.py +++ b/tests/operators/test_datetime.py @@ -23,7 +23,6 @@ import pytest import time_machine from airflow.exceptions import AirflowException -from airflow.models.dag import DAG from airflow.models.dagrun import DagRun from airflow.models.taskinstance import TaskInstance as TI from airflow.operators.datetime import BranchDateTimeOperator @@ -52,30 +51,30 @@ class TestBranchDateTimeOperator: (datetime.time(10, 0, 0), datetime.datetime(2020, 7, 7, 11, 0, 0)), ] - def setup_method(self): - self.dag = DAG( + @pytest.fixture(autouse=True) + def base_tests_setup(self, dag_maker): + with dag_maker( "branch_datetime_operator_test", default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, schedule=INTERVAL, - ) - - self.branch_1 = EmptyOperator(task_id="branch_1", dag=self.dag) - self.branch_2 = EmptyOperator(task_id="branch_2", dag=self.dag) + serialized=True, + ) as dag: + self.dag = dag + self.branch_1 = EmptyOperator(task_id="branch_1") + self.branch_2 = EmptyOperator(task_id="branch_2") - self.branch_op = BranchDateTimeOperator( - task_id="datetime_branch", - follow_task_ids_if_true="branch_1", - follow_task_ids_if_false="branch_2", - target_upper=datetime.datetime(2020, 7, 7, 11, 0, 0), - target_lower=datetime.datetime(2020, 7, 7, 10, 0, 0), - dag=self.dag, - ) + self.branch_op = BranchDateTimeOperator( + task_id="datetime_branch", + follow_task_ids_if_true="branch_1", + follow_task_ids_if_false="branch_2", + target_upper=datetime.datetime(2020, 7, 7, 11, 0, 0), + target_lower=datetime.datetime(2020, 7, 7, 10, 0, 0), + ) - self.branch_1.set_upstream(self.branch_op) - self.branch_2.set_upstream(self.branch_op) - self.dag.clear() + self.branch_1.set_upstream(self.branch_op) + self.branch_2.set_upstream(self.branch_op) - self.dr = self.dag.create_dagrun( + self.dr = dag_maker.create_dagrun( run_id="manual__", start_date=DEFAULT_DATE, execution_date=DEFAULT_DATE, @@ -231,11 +230,11 @@ class TestBranchDateTimeOperator: targets, ) @time_machine.travel("2020-12-01 09:00:00") - def test_branch_datetime_operator_use_task_logical_date(self, target_lower, target_upper): + def test_branch_datetime_operator_use_task_logical_date(self, dag_maker, target_lower, target_upper): """Check if BranchDateTimeOperator uses task execution date""" in_between_date = timezone.datetime(2020, 7, 7, 10, 30, 0) self.branch_op.use_task_logical_date = True - self.dr = self.dag.create_dagrun( + self.dr = dag_maker.create_dagrun( run_id="manual_exec_date__", start_date=in_between_date, execution_date=in_between_date,