Skip to content

Commit

Permalink
Provide arg defaults in dag_maker.create_dagrun()
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
uranusjr committed Jan 3, 2025
1 parent 960bbf9 commit 99e55bd
Showing 1 changed file with 14 additions and 10 deletions.
24 changes: 14 additions & 10 deletions tests_common/pytest_plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 99e55bd

Please sign in to comment.