44from airflow .models import DagBag , Variable
55from airflow .timetables .trigger import CronTriggerTimetable
66
7- from dags .fetch_notices_by_date import DAG_FETCH_DEFAULT_TIMEZONE
87from ted_sws import DAG_FETCH_DEFAULT_TIMETABLE
98
109
@@ -44,6 +43,32 @@ def test_fetcher_change_timetable_from_env_variable_after_reparse(dag_bag: DagBa
4443 assert fetcher_dag .schedule_interval == example_dag_cron_table ._expression
4544
4645 assert all (airflow_timetable_import_error_name not in error for error in dag_bag .import_errors .values ())
46+ del os .environ [dag_fetch_schedule_variable_name ]
47+
48+
49+ def test_fetcher_has_default_timetable_if_no_variable_is_set_after_reparse (dag_bag : DagBag ,
50+ dag_fetch_schedule_variable_name : str ,
51+ fetcher_dag_id : str ,
52+ airflow_timetable_import_error_name : str ):
53+ env_var_value = os .getenv (dag_fetch_schedule_variable_name )
54+ is_env_var_set : bool = True if env_var_value is not None else False
55+ if is_env_var_set :
56+ del os .environ [dag_fetch_schedule_variable_name ]
57+ airflow_var_value = Variable .get (key = dag_fetch_schedule_variable_name , default_var = None )
58+ is_airflow_var_set : bool = True if airflow_var_value is not None else False
59+ if is_airflow_var_set :
60+ Variable .delete (key = dag_fetch_schedule_variable_name )
61+
62+ dag_bag .collect_dags ()
63+ fetcher_dag : DAG = dag_bag .get_dag (dag_id = fetcher_dag_id )
64+ assert fetcher_dag is not None
65+ assert fetcher_dag .schedule_interval == DAG_FETCH_DEFAULT_TIMETABLE
66+ assert all (airflow_timetable_import_error_name not in error for error in dag_bag .import_errors .values ())
67+
68+ if is_env_var_set :
69+ os .environ [dag_fetch_schedule_variable_name ] = env_var_value
70+ if is_airflow_var_set :
71+ Variable .set (key = dag_fetch_schedule_variable_name , value = airflow_var_value )
4772
4873
4974def test_fetcher_gets_incorrect_timetable_after_reparse (dag_bag : DagBag ,
@@ -53,9 +78,6 @@ def test_fetcher_gets_incorrect_timetable_after_reparse(dag_bag: DagBag,
5378 airflow_timetable_import_error_name : str ):
5479 fetcher_dag : DAG = dag_bag .get_dag (dag_id = fetcher_dag_id )
5580 assert fetcher_dag is not None
56- default_dag_timetable = CronTriggerTimetable (cron = DAG_FETCH_DEFAULT_TIMETABLE ,
57- timezone = DAG_FETCH_DEFAULT_TIMEZONE )
58- assert fetcher_dag .schedule_interval == default_dag_timetable ._expression
5981
6082 Variable .set (key = dag_fetch_schedule_variable_name , value = example_wrong_cron_table )
6183 dag_bag .collect_dags (only_if_updated = False )
0 commit comments