From 99e55bd76db4766b83df4db595b43764fa29f7e8 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Fri, 3 Jan 2025 14:09:21 +0800 Subject: [PATCH] Provide arg defaults in dag_maker.create_dagrun() We can have dag_version and run_type defaults since this is just for tests and the actual values aren't that important most of the times. --- tests_common/pytest_plugin.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/tests_common/pytest_plugin.py b/tests_common/pytest_plugin.py index 43b4733e38811d..880ca9057d905a 100644 --- a/tests_common/pytest_plugin.py +++ b/tests_common/pytest_plugin.py @@ -884,33 +884,37 @@ def create_dagrun(self, *, logical_date=None, **kwargs): dag = self.dag kwargs = { + "dag_version": None, + "run_type": (run_type := DagRunType.MANUAL), "state": DagRunState.RUNNING, "start_date": self.start_date, "session": self.session, **kwargs, } - # Need to provide run_id if the user does not either provide one - # explicitly, or pass run_type for inference in dag.create_dagrun(). - if "run_id" not in kwargs and "run_type" not in kwargs: - kwargs["run_id"] = "test" - - if "run_type" not in kwargs: - kwargs["run_type"] = DagRunType.from_run_id(kwargs["run_id"]) if logical_date is None: - if kwargs["run_type"] == DagRunType.MANUAL: + if run_type == DagRunType.MANUAL: logical_date = self.start_date else: logical_date = dag.next_dagrun_info(None).logical_date logical_date = timezone.coerce_datetime(logical_date) - if "data_interval" not in kwargs: - if kwargs["run_type"] == DagRunType.MANUAL: + try: + data_interval = kwargs["data_interval"] + except KeyError: + if run_type == DagRunType.MANUAL: data_interval = dag.timetable.infer_manual_data_interval(run_after=logical_date) else: data_interval = dag.infer_automated_data_interval(logical_date) kwargs["data_interval"] = data_interval + if "run_id" not in kwargs: + kwargs["run_id"] = dag.timetable.generate_run_id( + run_type=run_type, + logical_date=logical_date, + data_interval=data_interval, + ) + if AIRFLOW_V_3_0_PLUS: kwargs.setdefault("triggered_by", DagRunTriggeredByType.TEST) kwargs["logical_date"] = logical_date