Skip to content

Commit fc318ca

Browse files
committed
add test to check default timetable
1 parent cc5c1bb commit fc318ca

1 file changed

Lines changed: 26 additions & 4 deletions

File tree

tests/unit/dags/test_fetch_notices_by_date.py

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
from airflow.models import DagBag, Variable
55
from airflow.timetables.trigger import CronTriggerTimetable
66

7-
from dags.fetch_notices_by_date import DAG_FETCH_DEFAULT_TIMEZONE
87
from 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

4974
def 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

Comments
 (0)