fix: broken composer tests

composer/blog/gcp-tech-blog/unit-test-dags-cloud-build/dags/example_dag_test.py

@@ -13,20 +13,20 @@
 # limitations under the License.
 
 from airflow import models
-import airflow.utils.db
 import internal_unit_testing
 import pytest
 
 # user should substitute their project ID
 PROJECT_ID = 'your-project-id'
 
 
-@pytest.fixture(autouse=True, scope="module")
-def initalize_airflow_database():
- airflow.utils.db.initdb()
+@pytest.fixture(autouse=True, scope="function")
+def set_variables(airflow_database):
+ models.Variable.set('gcp_project', PROJECT_ID)
+ yield
+ models.Variable.delete('gcp_project')
 
 
 def test_dag_import():
- models.Variable.set('gcp_project', PROJECT_ID)
  from . import example_dag
  internal_unit_testing.assert_has_valid_dag(example_dag)

composer/blog/noxfile_config.py

@@ -0,0 +1,39 @@
+# Copyright 2021 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
+#
+# http://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.
+
+# Default TEST_CONFIG_OVERRIDE for python repos.
+
+# You can copy this file into your directory, then it will be inported from
+# the noxfile.py.
+
+# The source of truth:
+# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/noxfile_config.py
+
+TEST_CONFIG_OVERRIDE = {
+ # You can opt out from the test for specific Python versions.
+ # Skipping for Python 3.9 due to numpy compilation failure.
+ "ignored_versions": ["2.7", "3.9"],
+ # Old samples are opted out of enforcing Python type hints
+ # All new samples should feature them
+ "enforce_type_hints": False,
+ # An envvar key for determining the project id to use. Change it
+ # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a
+ # build specific Cloud project. You can also use your own string
+ # to use your own Cloud project.
+ "gcloud_project_env": "GOOGLE_CLOUD_PROJECT",
+ # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT',
+ # A dictionary you want to inject into your test. Don't put any
+ # secrets here. These values will override predefined values.
+ "envs": {},
+}

composer/workflows/conftest.py → composer/conftest.py

@@ -13,10 +13,13 @@
 # limitations under the License.
 
 
-import airflow.utils.db
 import pytest
 
 
-@pytest.fixture(autouse=True, scope="session")
-def initalize_airflow_database():
- airflow.utils.db.initdb()
+# this fixture initializes the Airflow DB once per session
+# it is used by DAGs in both the blogs and workflows directories
+@pytest.fixture(scope="session")
+def airflow_database():
+ import airflow.utils.db
+ # reset both resets and initializes a new database
+ airflow.utils.db.resetdb(rbac=None) # this command will change in Airflow 2.0

composer/workflows/airflow_db_cleanup_test.py

@@ -16,7 +16,7 @@
 import internal_unit_testing
 
 
-def test_dag_import():
+def test_dag_import(airflow_database):
  """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

composer/workflows/bq_copy_across_locations_test.py

@@ -18,6 +18,21 @@
 from airflow import models
 
 import internal_unit_testing
+import pytest
+
+
+@pytest.fixture(autouse=True, scope="function")
+def set_variables(airflow_database):
+ example_file_path = os.path.join(
+ os.path.abspath(os.path.dirname(__file__)),
+ 'bq_copy_eu_to_us_sample.csv')
+ 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')
+ yield
+ models.Variable.delete('table_list_file_path')
+ models.Variable.delete('gcs_source_bucket')
+ models.Variable.delete('gcs_dest_bucket')
 
 
 def test_dag():
@@ -27,11 +42,6 @@ def test_dag():
  environment. This is a recommended confidence check by the official Airflow
  docs: https://airflow.incubator.apache.org/tutorial.html#testing
  """
- example_file_path = os.path.join(
- os.path.abspath(os.path.dirname(__file__)),
- 'bq_copy_eu_to_us_sample.csv')
- 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 as module
  internal_unit_testing.assert_has_valid_dag(module)

composer/workflows/bq_notify_test.py

@@ -15,6 +15,20 @@
 from airflow import models
 
 import internal_unit_testing
+import pytest
+
+
+@pytest.fixture(autouse=True, scope="function")
+def set_variables(airflow_database):
+ models.Variable.set('gcs_bucket', 'example_bucket')
+ models.Variable.set('gcp_project', 'example-project')
+ models.Variable.set('gce_zone', 'us-central1-f')
+ models.Variable.set('email', 'notify@example.com')
+ yield
+ models.Variable.delete('gcs_bucket')
+ models.Variable.delete('gcp_project')
+ models.Variable.delete('gce_zone')
+ models.Variable.delete('email')
 
 
 def test_dag_import():
@@ -24,9 +38,6 @@ def test_dag_import():
  environment. This is a recommended confidence check by the official Airflow
  docs: https://airflow.incubator.apache.org/tutorial.html#testing
  """
- models.Variable.set('gcs_bucket', 'example_bucket')
- models.Variable.set('gcp_project', 'example-project')
- models.Variable.set('gce_zone', 'us-central1-f')
- models.Variable.set('email', 'notify@example.com')
+
  from . import bq_notify as module
  internal_unit_testing.assert_has_valid_dag(module)

composer/workflows/dataflowtemplateoperator_tutorial_test.py

@@ -15,6 +15,20 @@
 from airflow import models
 
 import internal_unit_testing
+import pytest
+
+
+@pytest.fixture(autouse=True, scope="function")
+def set_variables(airflow_database):
+ models.Variable.set("bucket_path", "gs://example_bucket")
+ models.Variable.set("project_id", "example-project")
+ models.Variable.set("gce_zone", "us-central1-f")
+ models.Variable.set("gce_region", "us-central1-f")
+ yield
+ models.Variable.delete('bucket_path')
+ models.Variable.delete('project_id')
+ models.Variable.delete('gce_zone')
+ models.Variable.delete('gce_region')
 
 
 def test_dag_import():
@@ -24,10 +38,7 @@ def test_dag_import():
  environment. This is a recommended confidence check by the official Airflow
  docs: https://airflow.incubator.apache.org/tutorial.html#testing
  """
- models.Variable.set("bucket_path", "gs://example_bucket")
- models.Variable.set("project_id", "example-project")
- models.Variable.set("gce_zone", "us-central1-f")
- models.Variable.set("gce_region", "us-central1-f")
+
  from . import dataflowtemplateoperator_tutorial as module
 
  internal_unit_testing.assert_has_valid_dag(module)

composer/workflows/dataproc_workflow_template_instantiate_operator_tutorial_test.py

@@ -15,6 +15,14 @@
 from airflow import models
 
 import internal_unit_testing
+import pytest
+
+
+@pytest.fixture(autouse=True, scope="function")
+def set_variables(airflow_database):
+ models.Variable.set("project_id", "example-project")
+ yield
+ models.Variable.delete("project_id")
 
 
 def test_dag_import():
@@ -24,7 +32,5 @@ def test_dag_import():
  environment. This is a recommended confidence check by the official Airflow
  docs: https://airflow.incubator.apache.org/tutorial.html#testing
  """
- models.Variable.set("project_id", "example-project")
  from . import dataproc_workflow_template_instantiate_operator_tutorial as module
-
  internal_unit_testing.assert_has_valid_dag(module)

composer/workflows/hadoop_tutorial_test.py

@@ -15,6 +15,18 @@
 from airflow import models
 
 import internal_unit_testing
+import pytest
+
+
+@pytest.fixture(autouse=True, scope="function")
+def set_variables(airflow_database):
+ models.Variable.set('gcs_bucket', 'example-bucket')
+ models.Variable.set('gcp_project', 'example-project')
+ models.Variable.set('gce_zone', 'us-central1-f')
+ yield
+ models.Variable.delete('gcs_bucket')
+ models.Variable.delete('gcp_project')
+ models.Variable.delete('gce_zone')
 
 
 def test_dag_import():
@@ -24,8 +36,6 @@ def test_dag_import():
  environment. This is a recommended confidence check by the official Airflow
  docs: https://airflow.incubator.apache.org/tutorial.html#testing
  """
- models.Variable.set('gcs_bucket', 'example_bucket')
- models.Variable.set('gcp_project', 'example-project')
- models.Variable.set('gce_zone', 'us-central1-f')
+
  from . import hadoop_tutorial as module
  internal_unit_testing.assert_has_valid_dag(module)

composer/workflows/quickstart_test.py

@@ -15,6 +15,19 @@
 from airflow import models
 
 import internal_unit_testing
+import pytest
+
+
+@pytest.fixture(autouse=True, scope="function")
+def set_variables(airflow_database):
+
+ models.Variable.set('gcs_bucket', 'example_bucket')
+ models.Variable.set('gcp_project', 'example-project')
+ models.Variable.set('gce_zone', 'us-central1-f')
+ yield
+ models.Variable.delete('gcs_bucket')
+ models.Variable.delete('gcp_project')
+ models.Variable.delete('gce_zone')
 
 
 def test_dag_import():
@@ -24,8 +37,6 @@ def test_dag_import():
  environment. This is a recommended confidence check by the official Airflow
  docs: https://airflow.incubator.apache.org/tutorial.html#testing
  """
- 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 as module
  internal_unit_testing.assert_has_valid_dag(module)

composer/workflows/unit_testing_test.py

@@ -13,10 +13,19 @@
 # limitations under the License.
 
 from airflow import exceptions
+from airflow import models
+
 import internal_unit_testing
 import pytest
 
 
+@pytest.fixture(scope="function")
+def set_variables(airflow_database):
+ models.Variable.set('gcp_project', 'example-project')
+ yield
+ models.Variable.delete('gcp_project')
+
+
 def test_dag_no_dag():
  import internal_unit_testing as module # Does not contain a DAG.
  with pytest.raises(AssertionError):
@@ -30,11 +39,7 @@ def test_dag_has_cycle():
 
 
 # [START composer_dag_unit_testing_test]
-def test_dag_with_variables():
- from airflow import models
-
- # Set any Airflow variables before importing the DAG module.
- models.Variable.set('gcp_project', 'example-project')
+def test_dag_with_variables(set_variables):
 
  # Importing the module verifies that there are no syntax errors.
  from . import unit_testing_variables as module