prepare the tests for fetcher scheduling

Author: duprijil

dags/fetch_notices_by_date.py

@@ -15,7 +15,7 @@
     EventMessageProcessType
 from ted_sws.event_manager.services.log import log_error
 
-DAG_NAME = "fetch_notices_by_date"
+FETCHER_DAG_NAME = "fetch_notices_by_date"
 BATCH_SIZE = 2000
 WILD_CARD_DAG_KEY = "wild_card"
 TRIGGER_COMPLETE_WORKFLOW_DAG_KEY = "trigger_complete_workflow"
@@ -25,6 +25,9 @@
 FINISH_FETCH_BY_DATE_TASK_ID = "finish_fetch_by_date"
 VALIDATE_FETCHED_NOTICES_TASK_ID = "validate_fetched_notices"
 
+DAG_FETCH_DEFAULT_TIMEZONE = "UTC"
+DAG_FETCH_DEFAULT_TIMETABLE = "0 1 * * *"
+SCHEDULE_DAG_FETCH_VAR_NAME = "SCHEDULE_DAG_FETCH"
 
 @dag(default_args=DEFAULT_DAG_ARGUMENTS,
      catchup=False,
@@ -35,7 +38,7 @@ def fetch_notices_by_date():
     @event_log(TechnicalEventMessage(
         message="fetch_notice_from_ted",
         metadata=EventMessageMetadata(
-            process_type=EventMessageProcessType.DAG, process_name=DAG_NAME
+            process_type=EventMessageProcessType.DAG, process_name=FETCHER_DAG_NAME
         ))
     )
     def fetch_by_date_notice_from_ted():

tests/unit/dags/conftest.py

@@ -1,25 +1,74 @@
-# import os
-#
-# import pytest
-#
-# from airflow.models import DagBag
-# from airflow.utils import db
-# import logging
+import logging
+import os
 
+from airflow.exceptions import AirflowTimetableInvalid
+from airflow.models import DagBag
+from airflow.timetables.trigger import CronTriggerTimetable
+from airflow.utils import db
+from psutil.tests import pytest
+
+from dags.fetch_notices_by_date import FETCHER_DAG_NAME
 from tests import TESTS_PATH
 
 AIRFLOW_DAG_FOLDER = TESTS_PATH.parent.resolve() / "dags"
 
 
-# @pytest.fixture(scope="session")
+@pytest.fixture()
+def dag_bag():
+    os.environ["AIRFLOW_HOME"] = str(AIRFLOW_DAG_FOLDER)
+    os.environ["AIRFLOW__CORE__LOAD_EXAMPLES"] = "False"
+    # Initialising the Airflow DB so that it works properly with the new AIRFLOW_HOME
+    logging.disable(logging.CRITICAL)
+    db.resetdb()
+    db.initdb()
+    logging.disable(logging.NOTSET)
+    dag_bag = DagBag(dag_folder=AIRFLOW_DAG_FOLDER, include_examples=False,
+                     read_dags_from_db=False)
+    return dag_bag
+
+
+# @pytest.fixture
 # def dag_bag():
-#     os.environ["AIRFLOW_HOME"] = str(AIRFLOW_DAG_FOLDER)
-#     os.environ["AIRFLOW__CORE__LOAD_EXAMPLES"] = "False"
-#     # Initialising the Airflow DB so that it works properly with the new AIRFLOW_HOME
-#     logging.disable(logging.CRITICAL)
-#     db.resetdb()
-#     db.initdb()
-#     logging.disable(logging.NOTSET)
-#     dag_bag = DagBag(dag_folder=AIRFLOW_DAG_FOLDER, include_examples=False,
-#                      read_dags_from_db=False)
-#     return dag_bag
+#     return DagBag(dag_folder=AIRFLOW_DAG_FOLDER, include_examples=False, read_dags_from_db=False)
+
+
+@pytest.fixture
+def fetcher_dag_id():
+    return FETCHER_DAG_NAME
+
+
+@pytest.fixture
+def example_cron_table() -> str:
+    return "15 14 1 * *"
+
+
+@pytest.fixture
+def example_wrong_cron_table() -> str:
+    return "1234"
+
+
+@pytest.fixture
+def example_dag_cron_table(example_cron_table) -> CronTriggerTimetable:
+    return CronTriggerTimetable(cron=example_cron_table, timezone="UTC")
+
+
+@pytest.fixture
+def airflow_timetable_import_error_name() -> str:
+    return AirflowTimetableInvalid.__name__
+
+
+
+@pytest.fixture
+def dag_fetch_schedule_variable_name() -> str:
+    """
+    According to MM of meeting with OP from 2024.12.28
+    """
+    return "SCHEDULE_DAG_FETCH"
+
+
+@pytest.fixture
+def dag_materialised_view_update_schedule_variable_name() -> str:
+    """
+    According to MM of meeting with OP from 2024.12.28
+    """
+    return "SCHEDULE_DAG_MATERIALIZED_VIEW_UPDATE"

tests/unit/dags/test_fetch_notices_by_date.py

@@ -0,0 +1,51 @@
+from airflow import DAG
+from airflow.models import DagBag, Variable
+from airflow.timetables.trigger import CronTriggerTimetable
+
+from dags.fetch_notices_by_date import DAG_FETCH_DEFAULT_TIMETABLE, DAG_FETCH_DEFAULT_TIMEZONE, \
+    SCHEDULE_DAG_FETCH_VAR_NAME
+
+
+def test_schedule_dag_fetch_ver_name_is_correct(dag_fetch_schedule_variable_name):
+    assert SCHEDULE_DAG_FETCH_VAR_NAME == dag_fetch_schedule_variable_name
+
+
+def test_fetcher_has_default_timetable_at_the_beginning(dag_bag: DagBag, fetcher_dag_id: str):
+    fetcher_dag: DAG = dag_bag.get_dag(dag_id=fetcher_dag_id)
+    default_dag_timetable = CronTriggerTimetable(cron=DAG_FETCH_DEFAULT_TIMETABLE,
+                                                 timezone=DAG_FETCH_DEFAULT_TIMEZONE)
+
+    assert fetcher_dag is not None
+    assert fetcher_dag.schedule_interval == default_dag_timetable._expression
+
+
+def test_fetcher_gets_correct_timetable_after_reparse(dag_bag: DagBag, fetcher_dag_id: str,
+                                                      example_dag_cron_table: CronTriggerTimetable,
+                                                      airflow_timetable_import_error_name: str):
+    fetcher_dag: DAG = dag_bag.get_dag(dag_id=fetcher_dag_id)
+    assert fetcher_dag is not None
+    assert fetcher_dag.schedule_interval != example_dag_cron_table._expression
+
+    Variable.set(key=SCHEDULE_DAG_FETCH_VAR_NAME, value=example_dag_cron_table._expression)
+    dag_bag.collect_dags(only_if_updated=False)
+
+    fetcher_dag: DAG = dag_bag.get_dag(dag_id=fetcher_dag_id)
+    assert fetcher_dag is not None
+    assert fetcher_dag.schedule_interval == example_dag_cron_table._expression
+
+    assert all(airflow_timetable_import_error_name not in error for error in dag_bag.import_errors.values())
+
+
+def test_fetcher_gets_incorrect_timetable_after_reparse(dag_bag: DagBag, fetcher_dag_id: str,
+                                                        example_wrong_cron_table: str,
+                                                        airflow_timetable_import_error_name: str):
+    fetcher_dag: DAG = dag_bag.get_dag(dag_id=fetcher_dag_id)
+    assert fetcher_dag is not None
+    default_dag_timetable = CronTriggerTimetable(cron=DAG_FETCH_DEFAULT_TIMETABLE,
+                                                 timezone=DAG_FETCH_DEFAULT_TIMEZONE)
+    assert fetcher_dag.schedule_interval == default_dag_timetable._expression
+
+    Variable.set(key=SCHEDULE_DAG_FETCH_VAR_NAME, value=example_wrong_cron_table)
+    dag_bag.collect_dags(only_if_updated=False)
+
+    assert any(airflow_timetable_import_error_name in error for error in dag_bag.import_errors.values())