set AIRFLOW_HOME in noxfile_config.py

Takashi Matsuo · Takashi Matsuo · commit 851e4ccd2f3a · 2021-03-12T21:58:48.000Z
diff --git a/composer/blog/gcp-tech-blog/unit-test-dags-cloud-build/noxfile_config.py b/composer/blog/gcp-tech-blog/unit-test-dags-cloud-build/noxfile_config.py
@@ -1,4 +1,4 @@
-# Copyright 2021 Google LLC
+# Copyright 2020 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -20,13 +20,22 @@
 # The source of truth:
 # https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/noxfile_config.py
 
+import tempfile
+
+
+# Airflow creates a config file at the installation, so we want to set
+# `AIRFLOW_HOME` envvar before running pytest.
+
+_tmpdir = tempfile.TemporaryDirectory()
+
+
 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,
+    "enforce_type_hints": True,
     # 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
@@ -35,5 +44,5 @@
     # '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": {},
+    "envs": {"AIRFLOW_HOME": _tmpdir.name},
 }
diff --git a/composer/conftest.py b/composer/conftest.py
@@ -25,15 +25,16 @@
 def airflow_database():
     import airflow.utils.db
 
-    # We use separate directory for local db path per session.
-    tmp_dir = tempfile.TemporaryDirectory()
-    os.environ['AIRFLOW_HOME']  = tmp_dir.name
-    print(f"Setting AIRFLOW_HOME to {tmp_dir.name}")
+    # We use separate directory for local db path per session
+    # by setting AIRFLOW_HOME env var, which is done in noxfile_config.py.
+
+    assert('AIRFLOW_HOME' in os.environ)
+
+    airflow_home = os.environ["AIRFLOW_HOME"]
+    airflow_db = f"{airflow_home}/airflow.db"
 
     # reset both resets and initializes a new database
     airflow.utils.db.resetdb(rbac=None)  # this command will change in Airflow 2.0
 
-    yield
-
-    # cleaning up
-    tmp_dir.cleanup()
+    # Making sure we are using a data file there.
+    assert(os.path.isfile(airflow_db))
diff --git a/composer/workflows/noxfile_config.py b/composer/workflows/noxfile_config.py
@@ -20,6 +20,14 @@
 # The source of truth:
 # https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/noxfile_config.py
 
+import tempfile
+
+
+# Airflow creates a config file at the installation, so we want to set
+# `AIRFLOW_HOME` envvar before running pytest.
+
+_tmpdir = tempfile.TemporaryDirectory()
+
 TEST_CONFIG_OVERRIDE = {
     # You can opt out from the test for specific Python versions.
     # Skipping for Python 3.9 due to numpy compilation failure.
@@ -35,5 +43,5 @@
     # '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": {},
+    "envs": {"AIRFLOW_HOME": _tmpdir.name},
 }
