def test_branch_follow_true(self, _, weekday):
"""Checks if BranchDayOfWeekOperator follows true branch"""
print(datetime.datetime.now())
branch_op = BranchDayOfWeekOperator(
task_id="make_choice",
follow_task_ids_if_true=["branch_1", "branch_2"],
follow_task_ids_if_false="branch_3",
week_day=weekday,
dag=self.dag,
)
self.branch_1.set_upstream(branch_op)
self.branch_2.set_upstream(branch_op)
self.branch_3 = DummyOperator(task_id="branch_3", dag=self.dag)
self.branch_3.set_upstream(branch_op)
self.dag.clear()
dr = self.dag.create_dagrun(
run_id="manual__",
start_date=timezone.utcnow(),
execution_date=DEFAULT_DATE,
state=State.RUNNING,
)
branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
self._assert_task_ids_match_states(
dr,
{
'make_choice': State.SUCCESS,
'branch_1': State.NONE,
'branch_2': State.NONE,
'branch_3': State.SKIPPED,
},
)
def test_branch_follow_true_with_execution_date(self):
"""Checks if BranchDayOfWeekOperator follows true branch when set use_task_execution_day """
branch_op = BranchDayOfWeekOperator(
task_id="make_choice",
follow_task_ids_if_true="branch_1",
follow_task_ids_if_false="branch_2",
week_day="Wednesday",
use_task_execution_day=
True, # We compare to DEFAULT_DATE which is Wednesday
dag=self.dag,
)
self.branch_1.set_upstream(branch_op)
self.branch_2.set_upstream(branch_op)
self.dag.clear()
dr = self.dag.create_dagrun(
run_id="manual__",
start_date=timezone.utcnow(),
execution_date=DEFAULT_DATE,
state=State.RUNNING,
)
branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
self._assert_task_ids_match_states(
dr,
{
'make_choice': State.SUCCESS,
'branch_1': State.NONE,
'branch_2': State.SKIPPED,
},
)
def test_branch_xcom_push_true_branch(self):
"""Check if BranchDayOfWeekOperator push to xcom value of follow_task_ids_if_true"""
branch_op = BranchDayOfWeekOperator(
task_id="make_choice",
follow_task_ids_if_true="branch_1",
follow_task_ids_if_false="branch_2",
week_day="Monday",
dag=self.dag,
)
self.branch_1.set_upstream(branch_op)
self.branch_2.set_upstream(branch_op)
self.dag.clear()
dr = self.dag.create_dagrun(
run_id="manual__",
start_date=timezone.utcnow(),
execution_date=DEFAULT_DATE,
state=State.RUNNING,
)
branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
tis = dr.get_task_instances()
for ti in tis:
if ti.task_id == 'make_choice':
assert ti.xcom_pull(task_ids='make_choice') == 'branch_1'
def test_branch_with_no_weekday(self):
"""Check if BranchDayOfWeekOperator raises exception on missing weekday"""
with self.assertRaises(AirflowException):
BranchDayOfWeekOperator( # pylint: disable=missing-kwoa
task_id="make_choice",
follow_task_ids_if_true="branch_1",
follow_task_ids_if_false="branch_2",
dag=self.dag,
)
def test_weekday_branch_invalid_weekday_number(self):
"""Check if BranchDayOfWeekOperator raises exception on wrong value of weekday"""
invalid_week_day = 'Thsday'
with pytest.raises(
AttributeError,
match=f'Invalid Week Day passed: "{invalid_week_day}"'):
BranchDayOfWeekOperator(
task_id="make_choice",
follow_task_ids_if_true="branch_1",
follow_task_ids_if_false="branch_2",
week_day=invalid_week_day,
dag=self.dag,
)
def test_branch_with_invalid_type(self):
"""Check if BranchDayOfWeekOperator raises exception on unsupported weekday type"""
invalid_week_day = ['Monday']
with pytest.raises(
TypeError,
match='Unsupported Type for week_day parameter:'
' {}. It should be one of str, set or '
'Weekday enum type'.format(type(invalid_week_day)),
):
BranchDayOfWeekOperator(
task_id="make_choice",
follow_task_ids_if_true="branch_1",
follow_task_ids_if_false="branch_2",
week_day=invalid_week_day,
dag=self.dag,
)
from airflow.operators.dummy import DummyOperator
from airflow.operators.weekday import BranchDayOfWeekOperator
from airflow.utils.dates import days_ago
args = {
"owner": "airflow",
}
with DAG(
dag_id="example_weekday_branch_operator",
start_date=days_ago(2),
default_args=args,
tags=["example"],
schedule_interval="@daily",
) as dag:
# [START howto_operator_day_of_week_branch]
dummy_task_1 = DummyOperator(task_id='branch_true', dag=dag)
dummy_task_2 = DummyOperator(task_id='branch_false', dag=dag)
branch = BranchDayOfWeekOperator(
task_id="make_choice",
follow_task_ids_if_true="branch_true",
follow_task_ids_if_false="branch_false",
week_day="Monday",
)
# Run dummy_task_1 if branch executes on Monday
branch >> [dummy_task_1, dummy_task_2]
# [END howto_operator_day_of_week_branch]
def airflow2_good_example():
begin = DummyOperator(task_id="begin")
end = DummyOperator(task_id="end", trigger_rule=TriggerRule.NONE_FAILED)
check_day_of_week = BranchDayOfWeekOperator(
task_id="check_day_of_week",
week_day={WeekDay.SATURDAY, WeekDay.SUNDAY},
follow_task_ids_if_true="weekend",
follow_task_ids_if_false="weekday",
use_task_execution_day=True,
)
weekend = DummyOperator(task_id="weekend")
weekday = DummyOperator(task_id="weekday")
# Templated value for determining the name of the day of week based on the start date of the DagRun.
day_name = "{{ dag_run.start_date.strftime('%A').lower() }}"
# Begin weekday tasks.
with TaskGroup("weekday_activities") as weekday_activities:
which_weekday_activity_day = BranchPythonOperator(
task_id="which_weekday_activity_day",
python_callable=_get_activity,
op_args=[day_name],
)
for day, day_info in DAY_ACTIVITY_MAPPING.items():
if day_info["is_weekday"]:
day_of_week = Label(label=day)
activity = day_info["activity"]
do_activity = BashOperator(
task_id=activity.replace(" ", "_"),
bash_command=
f"echo It's {day.capitalize()} and I'm busy with {activity}.",
)
# Declaring task dependencies within the `TaskGroup` via the classic bitshift operator.
which_weekday_activity_day >> day_of_week >> do_activity
# Begin weekend tasks.
with TaskGroup("weekend_activities") as weekend_activities:
which_weekend_activity_day = BranchPythonOperator(
task_id="which_weekend_activity_day",
python_callable=_get_activity,
op_args=[day_name],
)
saturday = Label(label="saturday")
sunday = Label(label="sunday")
sleeping_in = BashOperator(
task_id="sleeping_in",
bash_command="sleep $[ ( $RANDOM % 30 ) + 1 ]s")
going_to_the_beach = _going_to_the_beach()
# Because the ``going_to_the_beach()`` function has ``multiple_outputs`` enabled, each dict key is
# accessible as their own `XCom` key.
inviting_friends = EmailOperator(
task_id="inviting_friends",
to="*****@*****.**",
subject=going_to_the_beach["subject"],
html_content=going_to_the_beach["body"],
)
# Using ``chain()`` here for list-to-list dependencies which are not supported by the bitshift
# operator and to simplify the notation for the desired dependency structure.
chain(which_weekend_activity_day, [saturday, sunday],
[going_to_the_beach, sleeping_in])
# High-level dependencies.
chain(begin, check_day_of_week, [weekday, weekend],
[weekday_activities, weekend_activities], end)