Test Composer DAG samples to verify they contain a valid DAG

Inspiration from Circle 1 of
https://medium.com/wbaa/datas-inferno-7-circles-of-data-testing-hell-with-airflow-cef4adff58d8

Beyond just asserting there are no syntax errors, we can verify that the
DAG files actually contain a DAG and that Airflow detects no cycles in
that DAG.

Author: tswast

composer/workflows/bq_copy_across_locations.py

@@ -51,9 +51,11 @@
 # Set default arguments
 # --------------------------------------------------------------------------------
 
+yesterday = datetime.datetime.now() - datetime.timedelta(days=1)
+
 default_args = {
     'owner': 'airflow',
-    'start_date': datetime.datetime.today(),
+    'start_date': yesterday,
     'depends_on_past': False,
     'email': [''],
     'email_on_failure': False,
@@ -122,17 +124,17 @@ def read_table_list(table_list_file):
 # Define a DAG (directed acyclic graph) of tasks.
 # Any task you create within the context manager is automatically added to the
 # DAG object.
-with models.DAG('bq_copy_us_to_eu_01',
-                default_args=default_args,
-                schedule_interval=None) as dag:
+with models.DAG(
+        'comoser_sample_bq_copy_across_locations',
+        default_args=default_args,
+        schedule_interval=None) as dag:
     start = dummy_operator.DummyOperator(
         task_id='start',
         trigger_rule='all_success'
     )
 
     end = dummy_operator.DummyOperator(
         task_id='end',
-
         trigger_rule='all_success'
     )
 

composer/workflows/bq_copy_across_locations_test.py

@@ -19,6 +19,8 @@
 from airflow import models
 import pytest
 
+from . import unit_testing
+
 
 @pytest.fixture(scope='module', autouse=True)
 def gcs_plugin():
@@ -35,7 +37,7 @@ def gcs_plugin():
     sys.path.remove(plugins_dir)
 
 
-def test_dag_import():
+def test_dag():
     """Test that the DAG file can be successfully imported.
 
     This tests that the DAG can be parsed, but does not run it in an Airflow
@@ -48,4 +50,5 @@ def test_dag_import():
     models.Variable.set('table_list_file_path', example_file_path)
     models.Variable.set('gcs_source_bucket', 'example-project')
     models.Variable.set('gcs_dest_bucket', 'us-central1-f')
-    from . import bq_copy_across_locations  # noqa
+    from . import bq_copy_across_locations as module
+    unit_testing.assert_has_valid_dag(module)

composer/workflows/bq_notify_test.py

@@ -14,6 +14,8 @@
 
 from airflow import models
 
+from . import unit_testing
+
 
 def test_dag_import():
     """Test that the DAG file can be successfully imported.
@@ -26,4 +28,5 @@ def test_dag_import():
     models.Variable.set('gcp_project', 'example-project')
     models.Variable.set('gce_zone', 'us-central1-f')
     models.Variable.set('email', '[email protected]')
-    from . import bq_notify  # noqa
+    from . import bq_notify as module
+    unit_testing.assert_has_valid_dag(module)

composer/workflows/connections_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from . import unit_testing
+
 
 def test_dag_import():
     """Test that the DAG file can be successfully imported.
@@ -20,4 +22,5 @@ def test_dag_import():
     environment. This is a recommended sanity check by the official Airflow
     docs: https://airflow.incubator.apache.org/tutorial.html#testing
     """
-    from . import connections  # noqa
+    from . import connections as module
+    unit_testing.assert_has_valid_dag(module)

composer/workflows/kubernetes_pod_operator_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from . import unit_testing
+
 
 def test_dag_import():
     """Test that the DAG file can be successfully imported.
@@ -20,4 +22,5 @@ def test_dag_import():
     environment. This is a recommended sanity check by the official Airflow
     docs: https://airflow.incubator.apache.org/tutorial.html#testing
     """
-    from . import kubernetes_pod_operator  # noqa
+    from . import kubernetes_pod_operator as module
+    unit_testing.assert_has_valid_dag(module)

composer/workflows/quickstart_test.py

@@ -14,6 +14,8 @@
 
 from airflow import models
 
+from . import unit_testing
+
 
 def test_dag_import():
     """Test that the DAG file can be successfully imported.
@@ -25,4 +27,5 @@ def test_dag_import():
     models.Variable.set('gcs_bucket', 'example_bucket')
     models.Variable.set('gcp_project', 'example-project')
     models.Variable.set('gce_zone', 'us-central1-f')
-    from . import quickstart  # noqa
+    from . import quickstart as module
+    unit_testing.assert_has_valid_dag(module)

composer/workflows/simple_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from . import unit_testing
+
 
 def test_dag_import():
     """Test that the DAG file can be successfully imported.
@@ -20,4 +22,5 @@ def test_dag_import():
     environment. This is a recommended sanity check by the official Airflow
     docs: https://airflow.incubator.apache.org/tutorial.html#testing
     """
-    from . import simple # noqa
+    from . import simple as module
+    unit_testing.assert_has_valid_dag(module)

composer/workflows/trigger_response_dag_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from . import unit_testing
+
 
 def test_dag_import():
     """Test that the DAG file can be successfully imported.
@@ -20,4 +22,5 @@ def test_dag_import():
     environment. This is a recommended sanity check by the official Airflow
     docs: https://airflow.incubator.apache.org/tutorial.html#testing
     """
-    from . import trigger_response_dag  # noqa
+    from . import trigger_response_dag as module
+    unit_testing.assert_has_valid_dag(module)

composer/workflows/unit_testing.py

@@ -0,0 +1,33 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for unit testing DAGs."""
+
+# [START composer_dag_unit_testing]
+from airflow import models
+
+
+def assert_has_valid_dag(module):
+    """Assert that a module contains a valid DAG."""
+
+    no_dag_found = True
+
+    for dag in vars(module).values():
+        if isinstance(dag, models.DAG):
+            no_dag_found = False
+            dag.test_cycle()  # Throws if a task cycle is found.
+
+    if no_dag_found:
+        raise AssertionError('module does not contain a valid DAG')
+# [END composer_dag_unit_testing]

composer/workflows/unit_testing_cycle.py

@@ -0,0 +1,35 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""An example DAG demonstrating a cyle in the task IDs."""
+
+import datetime
+
+from airflow import models
+from airflow.operators import dummy_operator
+
+
+yesterday = datetime.datetime.now() - datetime.timedelta(days=1)
+
+default_dag_args = {
+    'start_date': yesterday,
+}
+
+with models.DAG(
+        'composer_sample_cycle',
+        schedule_interval=datetime.timedelta(days=1),
+        default_args=default_dag_args) as dag:
+    start = dummy_operator.DummyOperator(task_id='oops_a_cycle')
+    end = dummy_operator.DummyOperator(task_id='oops_a_cycle')
+    start >> end