test!: Init tests for DAGs; make DagBag import without importing errors

Author: duprijil

dags/__init__.py

@@ -13,3 +13,23 @@
     "concurrency": 15,
     "execution_timeout": timedelta(days=10),
 }
+
+NOTICE_NORMALISATION_PIPELINE_TASK_ID = "notice_normalisation_pipeline"
+NOTICE_TRANSFORMATION_PIPELINE_TASK_ID = "notice_transformation_pipeline"
+NOTICE_DISTILLATION_PIPELINE_TASK_ID = "notice_distillation_pipeline"
+NOTICE_VALIDATION_PIPELINE_TASK_ID = "notice_validation_pipeline"
+NOTICE_PACKAGE_PIPELINE_TASK_ID = "notice_package_pipeline"
+NOTICE_PUBLISH_PIPELINE_TASK_ID = "notice_publish_pipeline"
+STOP_PROCESSING_TASK_ID = "stop_processing"
+BRANCH_SELECTOR_TASK_ID = 'branch_selector'
+SELECTOR_BRANCH_BEFORE_TRANSFORMATION_TASK_ID = "switch_to_transformation"
+SELECTOR_BRANCH_BEFORE_VALIDATION_TASK_ID = "switch_to_validation"
+SELECTOR_BRANCH_BEFORE_PACKAGE_TASK_ID = "switch_to_package"
+SELECTOR_BRANCH_BEFORE_PUBLISH_TASK_ID = "switch_to_publish"
+
+BRANCH_SELECTOR_MAP = {NOTICE_NORMALISATION_PIPELINE_TASK_ID: NOTICE_NORMALISATION_PIPELINE_TASK_ID,
+                       NOTICE_TRANSFORMATION_PIPELINE_TASK_ID: SELECTOR_BRANCH_BEFORE_TRANSFORMATION_TASK_ID,
+                       NOTICE_VALIDATION_PIPELINE_TASK_ID: SELECTOR_BRANCH_BEFORE_VALIDATION_TASK_ID,
+                       NOTICE_PACKAGE_PIPELINE_TASK_ID: SELECTOR_BRANCH_BEFORE_PACKAGE_TASK_ID,
+                       NOTICE_PUBLISH_PIPELINE_TASK_ID: SELECTOR_BRANCH_BEFORE_PUBLISH_TASK_ID
+                       }

dags/daily_materialized_views_update.py

@@ -15,7 +15,7 @@
 @dag(default_args=DEFAULT_DAG_ARGUMENTS,
      dag_id=DAILY_MATERIALISED_VIEWS_DAG_NAME,
      catchup=False,
-     timetable=CronTriggerTimetable(
+     schedule=CronTriggerTimetable(
          cron=config.SCHEDULE_DAG_MATERIALIZED_VIEW_UPDATE,
          timezone=DAG_DEFAULT_TIMEZONE),
      tags=['mongodb', 'daily-views-update'])

dags/fetch_notices_by_date.py

@@ -31,7 +31,7 @@
 @dag(default_args=DEFAULT_DAG_ARGUMENTS,
      dag_id=FETCHER_DAG_NAME,
      catchup=False,
-     timetable=CronTriggerTimetable(
+     schedule=CronTriggerTimetable(
          cron=config.SCHEDULE_DAG_FETCH,
          timezone=DAG_DEFAULT_TIMEZONE),
      tags=['selector', 'daily-fetch'],

dags/notice_processing_pipeline.py

@@ -4,36 +4,20 @@
 from airflow.decorators import dag
 from airflow.utils.trigger_rule import TriggerRule
 
-from dags import DEFAULT_DAG_ARGUMENTS
+from dags import DEFAULT_DAG_ARGUMENTS, NOTICE_NORMALISATION_PIPELINE_TASK_ID, STOP_PROCESSING_TASK_ID, \
+    BRANCH_SELECTOR_MAP, NOTICE_TRANSFORMATION_PIPELINE_TASK_ID, NOTICE_VALIDATION_PIPELINE_TASK_ID, \
+    NOTICE_PACKAGE_PIPELINE_TASK_ID, NOTICE_PUBLISH_PIPELINE_TASK_ID, BRANCH_SELECTOR_TASK_ID, \
+    SELECTOR_BRANCH_BEFORE_TRANSFORMATION_TASK_ID, SELECTOR_BRANCH_BEFORE_VALIDATION_TASK_ID, \
+    SELECTOR_BRANCH_BEFORE_PACKAGE_TASK_ID, SELECTOR_BRANCH_BEFORE_PUBLISH_TASK_ID, NOTICE_DISTILLATION_PIPELINE_TASK_ID
 from dags.dags_utils import get_dag_param, smart_xcom_push, smart_xcom_forward, smart_xcom_pull
 from dags.operators.DagBatchPipelineOperator import NoticeBatchPipelineOperator, NOTICE_IDS_KEY, \
     EXECUTE_ONLY_ONE_STEP_KEY, START_WITH_STEP_NAME_KEY
 from dags.pipelines.notice_batch_processor_pipelines import notices_batch_distillation_pipeline
 from dags.pipelines.notice_processor_pipelines import notice_normalisation_pipeline, notice_transformation_pipeline, \
     notice_validation_pipeline, notice_package_pipeline, notice_publish_pipeline
 
-NOTICE_NORMALISATION_PIPELINE_TASK_ID = "notice_normalisation_pipeline"
-NOTICE_TRANSFORMATION_PIPELINE_TASK_ID = "notice_transformation_pipeline"
-NOTICE_DISTILLATION_PIPELINE_TASK_ID = "notice_distillation_pipeline"
-NOTICE_VALIDATION_PIPELINE_TASK_ID = "notice_validation_pipeline"
-NOTICE_PACKAGE_PIPELINE_TASK_ID = "notice_package_pipeline"
-NOTICE_PUBLISH_PIPELINE_TASK_ID = "notice_publish_pipeline"
-STOP_PROCESSING_TASK_ID = "stop_processing"
-BRANCH_SELECTOR_TASK_ID = 'branch_selector'
-SELECTOR_BRANCH_BEFORE_TRANSFORMATION_TASK_ID = "switch_to_transformation"
-SELECTOR_BRANCH_BEFORE_VALIDATION_TASK_ID = "switch_to_validation"
-SELECTOR_BRANCH_BEFORE_PACKAGE_TASK_ID = "switch_to_package"
-SELECTOR_BRANCH_BEFORE_PUBLISH_TASK_ID = "switch_to_publish"
 DAG_NAME = "notice_processing_pipeline"
 
-BRANCH_SELECTOR_MAP = {NOTICE_NORMALISATION_PIPELINE_TASK_ID: NOTICE_NORMALISATION_PIPELINE_TASK_ID,
-                       NOTICE_TRANSFORMATION_PIPELINE_TASK_ID: SELECTOR_BRANCH_BEFORE_TRANSFORMATION_TASK_ID,
-                       NOTICE_VALIDATION_PIPELINE_TASK_ID: SELECTOR_BRANCH_BEFORE_VALIDATION_TASK_ID,
-                       NOTICE_PACKAGE_PIPELINE_TASK_ID: SELECTOR_BRANCH_BEFORE_PACKAGE_TASK_ID,
-                       NOTICE_PUBLISH_PIPELINE_TASK_ID: SELECTOR_BRANCH_BEFORE_PUBLISH_TASK_ID
-                       }
-
-
 def branch_selector(result_branch: str, xcom_forward_keys: List[str] = [NOTICE_IDS_KEY]) -> str:
     start_with_step_name = get_dag_param(key=START_WITH_STEP_NAME_KEY,
                                          default_value=NOTICE_NORMALISATION_PIPELINE_TASK_ID)

dags/reprocess_published_in_cellar_notices.py

@@ -1,10 +1,8 @@
 from airflow.decorators import dag, task
 
-from dags import DEFAULT_DAG_ARGUMENTS
+from dags import DEFAULT_DAG_ARGUMENTS, NOTICE_TRANSFORMATION_PIPELINE_TASK_ID
 from dags.dags_utils import push_dag_downstream, get_dag_param
-from dags.notice_processing_pipeline import NOTICE_TRANSFORMATION_PIPELINE_TASK_ID
-from dags.operators.DagBatchPipelineOperator import NOTICE_IDS_KEY, TriggerNoticeBatchPipelineOperator, \
-    EXECUTE_ONLY_ONE_STEP_KEY
+from dags.operators.DagBatchPipelineOperator import NOTICE_IDS_KEY, TriggerNoticeBatchPipelineOperator
 from dags.pipelines.notice_selectors_pipelines import notice_ids_selector_by_status
 from ted_sws.core.model.notice import NoticeStatus
 from ted_sws.event_manager.adapters.event_log_decorator import event_log

dags/reprocess_unpackaged_notices_from_backlog.py

@@ -1,8 +1,7 @@
 from airflow.decorators import dag, task
 
-from dags import DEFAULT_DAG_ARGUMENTS
+from dags import DEFAULT_DAG_ARGUMENTS, NOTICE_PACKAGE_PIPELINE_TASK_ID
 from dags.dags_utils import push_dag_downstream, get_dag_param
-from dags.notice_processing_pipeline import NOTICE_PACKAGE_PIPELINE_TASK_ID
 from dags.operators.DagBatchPipelineOperator import NOTICE_IDS_KEY, TriggerNoticeBatchPipelineOperator
 from dags.pipelines.notice_selectors_pipelines import notice_ids_selector_by_status
 from ted_sws.core.model.notice import NoticeStatus

dags/reprocess_unpublished_notices_from_backlog.py

@@ -1,8 +1,7 @@
 from airflow.decorators import dag, task
 
-from dags import DEFAULT_DAG_ARGUMENTS
+from dags import DEFAULT_DAG_ARGUMENTS, NOTICE_PUBLISH_PIPELINE_TASK_ID
 from dags.dags_utils import push_dag_downstream, get_dag_param
-from dags.notice_processing_pipeline import NOTICE_PUBLISH_PIPELINE_TASK_ID
 from dags.operators.DagBatchPipelineOperator import NOTICE_IDS_KEY, TriggerNoticeBatchPipelineOperator, \
     EXECUTE_ONLY_ONE_STEP_KEY
 from dags.pipelines.notice_selectors_pipelines import notice_ids_selector_by_status

dags/reprocess_untransformed_notices_from_backlog.py

@@ -1,10 +1,8 @@
 from airflow.decorators import dag, task
 
-from dags import DEFAULT_DAG_ARGUMENTS
+from dags import DEFAULT_DAG_ARGUMENTS, NOTICE_TRANSFORMATION_PIPELINE_TASK_ID
 from dags.dags_utils import push_dag_downstream, get_dag_param
-from dags.notice_processing_pipeline import NOTICE_TRANSFORMATION_PIPELINE_TASK_ID
-from dags.operators.DagBatchPipelineOperator import NOTICE_IDS_KEY, TriggerNoticeBatchPipelineOperator, \
-    EXECUTE_ONLY_ONE_STEP_KEY
+from dags.operators.DagBatchPipelineOperator import NOTICE_IDS_KEY, TriggerNoticeBatchPipelineOperator
 from dags.pipelines.notice_selectors_pipelines import notice_ids_selector_by_status
 from ted_sws.core.model.notice import NoticeStatus
 from ted_sws.event_manager.adapters.event_log_decorator import event_log
@@ -14,7 +12,8 @@
 DAG_NAME = "reprocess_untransformed_notices_from_backlog"
 
 RE_TRANSFORM_TARGET_NOTICE_STATES = [NoticeStatus.NORMALISED_METADATA, NoticeStatus.INELIGIBLE_FOR_TRANSFORMATION,
-                                     NoticeStatus.ELIGIBLE_FOR_TRANSFORMATION, NoticeStatus.PREPROCESSED_FOR_TRANSFORMATION,
+                                     NoticeStatus.ELIGIBLE_FOR_TRANSFORMATION,
+                                     NoticeStatus.PREPROCESSED_FOR_TRANSFORMATION,
                                      NoticeStatus.TRANSFORMED, NoticeStatus.DISTILLED
                                      ]
 TRIGGER_NOTICE_PROCESS_WORKFLOW_TASK_ID = "trigger_notice_process_workflow"

dags/reprocess_unvalidated_notices_from_backlog.py

@@ -1,10 +1,8 @@
 from airflow.decorators import dag, task
 
-from dags import DEFAULT_DAG_ARGUMENTS
+from dags import DEFAULT_DAG_ARGUMENTS, NOTICE_TRANSFORMATION_PIPELINE_TASK_ID
 from dags.dags_utils import push_dag_downstream, get_dag_param
-from dags.notice_processing_pipeline import NOTICE_TRANSFORMATION_PIPELINE_TASK_ID
-from dags.operators.DagBatchPipelineOperator import NOTICE_IDS_KEY, TriggerNoticeBatchPipelineOperator, \
-    EXECUTE_ONLY_ONE_STEP_KEY
+from dags.operators.DagBatchPipelineOperator import NOTICE_IDS_KEY, TriggerNoticeBatchPipelineOperator
 from dags.pipelines.notice_selectors_pipelines import notice_ids_selector_by_status
 from ted_sws.core.model.notice import NoticeStatus
 from ted_sws.event_manager.adapters.event_log_decorator import event_log

tests/unit/dags/conftest.py

@@ -1,19 +1,38 @@
+import os
+import tempfile
+
 import pytest
+from airflow.models import DagBag, Variable
 from airflow.timetables.trigger import CronTriggerTimetable
-
+from airflow.utils.db import initdb
 from dags.daily_materialized_views_update import DAILY_MATERIALISED_VIEWS_DAG_NAME
 from dags.fetch_notices_by_date import FETCHER_DAG_NAME
+from tests import AIRFLOW_DAG_FOLDER
+
+@pytest.fixture(scope="session", autouse=True)
+def setup_airflow():
+    # Setup
+    temp_db_file = tempfile.NamedTemporaryFile(mode="w+", suffix=".db")
+    os.environ["AIRFLOW__CORE__SQL_ALCHEMY_CONN"] = f"sqlite:///{temp_db_file.name}"
+    os.environ["AIRFLOW__CORE__LOAD_EXAMPLES"] = "False"
+    os.environ["AIRFLOW__CORE__UNIT_TEST_MODE"] = "True"
+    initdb()
 
+    # Run tests
+    yield
 
-# @pytest.fixture
-# def dag_bag(dag_materialised_view_update_schedule_variable_name, dag_fetch_schedule_variable_name) -> DagBag:
-#     Variable.delete(key=dag_materialised_view_update_schedule_variable_name)
-#     Variable.delete(key=dag_fetch_schedule_variable_name)
-#     return DagBag(
-#         dag_folder=AIRFLOW_DAG_FOLDER,
-#         include_examples=False,
-#         read_dags_from_db=False,
-#         collect_dags=True)
+    # Cleanup
+    temp_db_file.close()
+
+@pytest.fixture
+def dag_bag(dag_materialised_view_update_schedule_variable_name, dag_fetch_schedule_variable_name) -> DagBag:
+    Variable.delete(key=dag_materialised_view_update_schedule_variable_name)
+    Variable.delete(key=dag_fetch_schedule_variable_name)
+    return DagBag(
+        dag_folder=AIRFLOW_DAG_FOLDER,
+        include_examples=False, #Same as: os.environ["AIRFLOW__CORE__LOAD_EXAMPLES"] = "False"
+        read_dags_from_db=False,
+        collect_dags=True)
 
 
 @pytest.fixture