Pass run_after to DagRun in tests

uranusjr · uranusjr · commit ae3539f6cea9 · 2025-01-17T23:02:48.000+08:00
diff --git a/providers/tests/apache/kylin/operators/test_kylin_cube.py b/providers/tests/apache/kylin/operators/test_kylin_cube.py
@@ -176,6 +176,8 @@ def test_render_template(self, session):
                 dag_id=self.dag.dag_id,
                 run_id="kylin_test",
                 logical_date=DEFAULT_DATE,
+                data_interval=(DEFAULT_DATE, DEFAULT_DATE),
+                run_after=DEFAULT_DATE,
                 run_type=DagRunType.MANUAL,
             )
         else:
diff --git a/providers/tests/apache/spark/operators/test_spark_submit.py b/providers/tests/apache/spark/operators/test_spark_submit.py
@@ -200,6 +200,8 @@ def test_render_template(self, session):
                 dag_id=self.dag.dag_id,
                 run_id="spark_test",
                 logical_date=DEFAULT_DATE,
+                data_interval=(DEFAULT_DATE, DEFAULT_DATE),
+                run_after=DEFAULT_DATE,
                 run_type=DagRunType.MANUAL,
             )
         else:
diff --git a/providers/tests/cncf/kubernetes/operators/test_pod.py b/providers/tests/cncf/kubernetes/operators/test_pod.py
@@ -51,6 +51,7 @@
 from airflow.utils.types import DagRunType
 
 from tests_common.test_utils import db
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
 pytestmark = pytest.mark.db_test
 
@@ -92,11 +93,23 @@ def create_context(task, persist_to_db=False, map_index=None):
     else:
         dag = DAG(dag_id="dag", schedule=None, start_date=pendulum.now())
         dag.add_task(task)
-    dag_run = DagRun(
-        run_id=DagRun.generate_run_id(DagRunType.MANUAL, DEFAULT_DATE),
-        run_type=DagRunType.MANUAL,
-        dag_id=dag.dag_id,
-    )
+    now = timezone.utcnow()
+    if AIRFLOW_V_3_0_PLUS:
+        dag_run = DagRun(
+            run_id=DagRun.generate_run_id(DagRunType.MANUAL, DEFAULT_DATE),
+            run_type=DagRunType.MANUAL,
+            dag_id=dag.dag_id,
+            logical_date=now,
+            data_interval=(now, now),
+            run_after=now,
+        )
+    else:
+        dag_run = DagRun(
+            run_id=DagRun.generate_run_id(DagRunType.MANUAL, DEFAULT_DATE),
+            run_type=DagRunType.MANUAL,
+            dag_id=dag.dag_id,
+            execution_date=now,
+        )
     task_instance = TaskInstance(task=task, run_id=dag_run.run_id)
     task_instance.dag_run = dag_run
     if map_index is not None:
diff --git a/providers/tests/fab/auth_manager/api_endpoints/test_dag_run_endpoint.py b/providers/tests/fab/auth_manager/api_endpoints/test_dag_run_endpoint.py
@@ -142,18 +142,26 @@ def teardown_method(self) -> None:
     def _create_test_dag_run(self, state=DagRunState.RUNNING, extra_dag=False, commit=True, idx_start=1):
         dag_runs = []
         dags = []
-        triggered_by_kwargs = {"triggered_by": DagRunTriggeredByType.TEST} if AIRFLOW_V_3_0_PLUS else {}
+
+        def _v3_kwargs(date):
+            if AIRFLOW_V_3_0_PLUS:
+                return {
+                    "data_interval": (date, date),
+                    "logical_date": date,
+                    "run_after": date,
+                    "triggered_by": DagRunTriggeredByType.TEST,
+                }
+            return {"execution_date": date}
 
         for i in range(idx_start, idx_start + 2):
             dagrun_model = DagRun(
                 dag_id="TEST_DAG_ID",
                 run_id=f"TEST_DAG_RUN_ID_{i}",
                 run_type=DagRunType.MANUAL,
-                logical_date=timezone.parse(self.default_time) + timedelta(days=i - 1),
                 start_date=timezone.parse(self.default_time),
                 external_trigger=True,
                 state=state,
-                **triggered_by_kwargs,
+                **_v3_kwargs(timezone.parse(self.default_time) + timedelta(days=i - 1)),
             )
             dagrun_model.updated_at = timezone.parse(self.default_time)
             dag_runs.append(dagrun_model)
@@ -166,10 +174,10 @@ def _create_test_dag_run(self, state=DagRunState.RUNNING, extra_dag=False, commi
                         dag_id=f"TEST_DAG_ID_{i}",
                         run_id=f"TEST_DAG_RUN_ID_{i}",
                         run_type=DagRunType.MANUAL,
-                        logical_date=timezone.parse(self.default_time_2),
                         start_date=timezone.parse(self.default_time),
                         external_trigger=True,
                         state=state,
+                        **_v3_kwargs(timezone.parse(self.default_time_2)),
                     )
                 )
         if commit:
@@ -203,8 +211,8 @@ def test_should_return_accessible_with_tilde_as_dag_id_and_dag_level_permissions
             "external_trigger": True,
             "start_date": self.default_time,
             "conf": {},
-            "data_interval_end": None,
-            "data_interval_start": None,
+            "data_interval_end": self.default_time,
+            "data_interval_start": self.default_time,
             "last_scheduling_decision": None,
             "run_type": "manual",
             "note": None,
@@ -219,8 +227,8 @@ def test_should_return_accessible_with_tilde_as_dag_id_and_dag_level_permissions
             "external_trigger": True,
             "start_date": self.default_time,
             "conf": {},
-            "data_interval_end": None,
-            "data_interval_start": None,
+            "data_interval_end": self.default_time_2,
+            "data_interval_start": self.default_time_2,
             "last_scheduling_decision": None,
             "run_type": "manual",
             "note": None,
diff --git a/providers/tests/fab/auth_manager/api_endpoints/test_task_instance_endpoint.py b/providers/tests/fab/auth_manager/api_endpoints/test_task_instance_endpoint.py
@@ -166,6 +166,8 @@ def create_task_instances(
                     run_id=run_id,
                     dag_id=dag_id,
                     logical_date=logical_date,
+                    data_interval=(logical_date, logical_date),
+                    run_after=logical_date,
                     run_type=DagRunType.MANUAL,
                     state=dag_run_state,
                 )
diff --git a/providers/tests/fab/auth_manager/api_endpoints/test_xcom_endpoint.py b/providers/tests/fab/auth_manager/api_endpoints/test_xcom_endpoint.py
@@ -180,6 +180,8 @@ def _create_xcom_entries(self, dag_id, run_id, logical_date, task_id, mapped_ti=
                     dag_id=dag_id,
                     run_id=run_id,
                     logical_date=logical_date,
+                    data_interval=(logical_date, logical_date),
+                    run_after=logical_date,
                     start_date=logical_date,
                     run_type=DagRunType.MANUAL,
                 )
@@ -226,6 +228,7 @@ def _create_invalid_xcom_entries(self, logical_date):
                     dag_id="invalid_dag",
                     run_id="invalid_run_id",
                     logical_date=logical_date + timedelta(days=1),
+                    run_after=logical_date,
                     start_date=logical_date,
                     run_type=DagRunType.MANUAL,
                 )
@@ -243,6 +246,7 @@ def _create_invalid_xcom_entries(self, logical_date):
                     dag_id="invalid_dag",
                     run_id="not_this_run_id",
                     logical_date=logical_date,
+                    run_after=logical_date,
                     start_date=logical_date,
                     run_type=DagRunType.MANUAL,
                 )
diff --git a/providers/tests/redis/log/test_redis_task_handler.py b/providers/tests/redis/log/test_redis_task_handler.py
@@ -42,7 +42,14 @@ def ti(self):
         dag = DAG(dag_id="dag_for_testing_redis_task_handler", schedule=None, start_date=date)
         task = EmptyOperator(task_id="task_for_testing_redis_log_handler", dag=dag)
         if AIRFLOW_V_3_0_PLUS:
-            dag_run = DagRun(dag_id=dag.dag_id, logical_date=date, run_id="test", run_type="scheduled")
+            dag_run = DagRun(
+                dag_id=dag.dag_id,
+                logical_date=date,
+                data_interval=(date, date),
+                run_after=date,
+                run_id="test",
+                run_type="scheduled",
+            )
         else:
             dag_run = DagRun(dag_id=dag.dag_id, execution_date=date, run_id="test", run_type="scheduled")
 
diff --git a/tests/jobs/test_triggerer_job.py b/tests/jobs/test_triggerer_job.py
@@ -93,10 +93,13 @@ def session():
 def create_trigger_in_db(session, trigger, operator=None):
     dag_model = DagModel(dag_id="test_dag")
     dag = DAG(dag_id=dag_model.dag_id, schedule="@daily", start_date=pendulum.datetime(2023, 1, 1))
+    date = pendulum.datetime(2023, 1, 1)
     run = DagRun(
         dag_id=dag_model.dag_id,
         run_id="test_run",
-        logical_date=pendulum.datetime(2023, 1, 1),
+        logical_date=date,
+        data_interval=(date, date),
+        run_after=date,
         run_type=DagRunType.MANUAL,
     )
     trigger_orm = Trigger.from_object(trigger)
diff --git a/tests/models/test_xcom.py b/tests/models/test_xcom.py
@@ -64,6 +64,8 @@ def func(*, dag_id, task_id, logical_date):
             run_type=DagRunType.SCHEDULED,
             run_id=run_id,
             logical_date=logical_date,
+            data_interval=(logical_date, logical_date),
+            run_after=logical_date,
         )
         session.add(run)
         ti = TaskInstance(EmptyOperator(task_id=task_id), run_id=run_id)

Original file line number	Diff line number	Diff line change
`@@ -176,6 +176,8 @@ def test_render_template(self, session):`
`176`	`176`	`dag_id=self.dag.dag_id,`
`177`	`177`	`run_id="kylin_test",`
`178`	`178`	`logical_date=DEFAULT_DATE,`
	`179`	`+ data_interval=(DEFAULT_DATE, DEFAULT_DATE),`
	`180`	`+ run_after=DEFAULT_DATE,`
`179`	`181`	`run_type=DagRunType.MANUAL,`
`180`	`182`	`)`
`181`	`183`	`else:`
Original file line number	Diff line number	Diff line change
`@@ -200,6 +200,8 @@ def test_render_template(self, session):`
`200`	`200`	`dag_id=self.dag.dag_id,`
`201`	`201`	`run_id="spark_test",`
`202`	`202`	`logical_date=DEFAULT_DATE,`
	`203`	`+ data_interval=(DEFAULT_DATE, DEFAULT_DATE),`
	`204`	`+ run_after=DEFAULT_DATE,`
`203`	`205`	`run_type=DagRunType.MANUAL,`
`204`	`206`	`)`
`205`	`207`	`else:`
Original file line number	Diff line number	Diff line change
`@@ -166,6 +166,8 @@ def create_task_instances(`
`166`	`166`	`run_id=run_id,`
`167`	`167`	`dag_id=dag_id,`
`168`	`168`	`logical_date=logical_date,`
	`169`	`+ data_interval=(logical_date, logical_date),`
	`170`	`+ run_after=logical_date,`
`169`	`171`	`run_type=DagRunType.MANUAL,`
`170`	`172`	`state=dag_run_state,`
`171`	`173`	`)`
Original file line number	Diff line number	Diff line change
`@@ -64,6 +64,8 @@ def func(*, dag_id, task_id, logical_date):`
`64`	`64`	`run_type=DagRunType.SCHEDULED,`
`65`	`65`	`run_id=run_id,`
`66`	`66`	`logical_date=logical_date,`
	`67`	`+ data_interval=(logical_date, logical_date),`
	`68`	`+ run_after=logical_date,`
`67`	`69`	`)`
`68`	`70`	`session.add(run)`
`69`	`71`	`ti = TaskInstance(EmptyOperator(task_id=task_id), run_id=run_id)`