From 4563df404ca2ec481153c39a1193419331a9507e Mon Sep 17 00:00:00 2001
From: Svend Vanderveken <svend.vanderveken@gmail.com>
Date: Thu, 19 Nov 2015 16:52:28 +0100
Subject: [PATCH] add Fernet key to test config

- congiguration.py now generates the test config and real airflow config with the same method
- fix warning logs in local UT execution, related to missing FERNET-KEY in unittest.cfg
- fix warning logs in Travis UT exeuction, related to missing FERNET-KEY in airflow_travis.cfg
- added one to validate config generation
---
 airflow/configuration.py      | 46 ++++++++++++++++++++++++++---------
 scripts/ci/airflow_travis.cfg |  1 +
 tests/core.py                 | 14 +++++++++++
 3 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/airflow/configuration.py b/airflow/configuration.py
index c616c3500b..f9c3e93299 100644
--- a/airflow/configuration.py
+++ b/airflow/configuration.py
@@ -4,15 +4,13 @@ from __future__ import print_function
 from __future__ import unicode_literals
 
 from future import standard_library
+
 standard_library.install_aliases()
 from builtins import str
 from configparser import ConfigParser
 import errno
 import logging
 import os
-import sys
-import textwrap
-
 
 try:
     from cryptography.fernet import Fernet
@@ -287,6 +285,7 @@ unit_test_mode = True
 load_examples = True
 donot_pickle = False
 dag_concurrency = 16
+fernet_key = {FERNET_KEY}
 
 [webserver]
 base_url = http://localhost:8080
@@ -396,24 +395,40 @@ if 'AIRFLOW_CONFIG' not in os.environ:
 else:
     AIRFLOW_CONFIG = expand_env_var(os.environ['AIRFLOW_CONFIG'])
 
+
+def parameterized_config(template):
+    """
+    Generates a configuration from the provided template + variables defined in
+    current scope
+    :param template: a config content templated with {{variables}}
+    """
+
+    FERNET_KEY = generate_fernet_key()
+    all_vars = {k: v for d in [globals(), locals()] for k, v in d.items()}
+    return template.format(**all_vars)
+
+TEST_CONFIG_FILE = AIRFLOW_HOME + '/unittests.cfg'
+if not os.path.isfile(TEST_CONFIG_FILE):
+    logging.info("Creating new airflow config file for unit tests in: " +
+                 TEST_CONFIG_FILE)
+    with open(AIRFLOW_CONFIG, 'w') as f:
+        f.write(parameterized_config(TEST_CONFIG))
+
 if not os.path.isfile(AIRFLOW_CONFIG):
     """
     These configuration options are used to generate a default configuration
     when it is missing. The right way to change your configuration is to alter
     your configuration file, not this code.
     """
-    FERNET_KEY = generate_fernet_key()
-    logging.info("Creating new config file in: " + AIRFLOW_CONFIG)
-    f = open(AIRFLOW_CONFIG, 'w')
-    f.write(DEFAULT_CONFIG.format(**locals()))
-    f.close()
+    logging.info("Creating new airflow config file in: " + AIRFLOW_CONFIG)
+    with open(AIRFLOW_CONFIG, 'w') as f:
+        f.write(parameterized_config(DEFAULT_CONFIG))
 
 TEST_CONFIG_FILE = AIRFLOW_HOME + '/unittests.cfg'
 if not os.path.isfile(TEST_CONFIG_FILE):
     logging.info("Creating new config file in: " + TEST_CONFIG_FILE)
-    f = open(TEST_CONFIG_FILE, 'w')
-    f.write(TEST_CONFIG.format(**locals()))
-    f.close()
+    with open(TEST_CONFIG_FILE, 'w') as f:
+        f.write(TEST_CONFIG.format(**locals()))
 
 logging.info("Reading the config from " + AIRFLOW_CONFIG)
 
@@ -437,9 +452,18 @@ def getboolean(section, key):
 def getfloat(section, key):
     return conf.getfloat(section, key)
 
+
 def getint(section, key):
     return conf.getint(section, key)
 
+
 def has_option(section, key):
     return conf.has_option(section, key)
 
+
+########################
+# convenience method to access config entries
+
+def get_dags_folder():
+    return os.path.expanduser(get('core', 'DAGS_FOLDER'))
+
diff --git a/scripts/ci/airflow_travis.cfg b/scripts/ci/airflow_travis.cfg
index c12fe3f5fc..682f9cb809 100644
--- a/scripts/ci/airflow_travis.cfg
+++ b/scripts/ci/airflow_travis.cfg
@@ -8,6 +8,7 @@ unit_test_mode = True
 load_examples = True
 donot_pickle = False
 parallelism = 2
+fernet_key = af7CN0q6ag5U3g08IsPsw3K45U7Xa0axgVFhoh-3zB8=
 
 [webserver]
 base_url = http://localhost:8080
diff --git a/tests/core.py b/tests/core.py
index c691170d30..506fa1a71e 100644
--- a/tests/core.py
+++ b/tests/core.py
@@ -378,6 +378,20 @@ class CoreTest(unittest.TestCase):
                                              default_var=default_value,
                                              deserialize_json=True)
 
+    def test_parameterized_config_gen(self):
+
+        cfg = configuration.parameterized_config(configuration.DEFAULT_CONFIG)
+
+        # making sure some basic building blocks are present:
+        assert "[core]" in cfg
+        assert "dags_folder" in cfg
+        assert "sql_alchemy_conn" in cfg
+        assert "fernet_key" in cfg
+
+        # making sure replacement actually happened
+        assert "{AIRFLOW_HOME}" not in cfg
+        assert "{FERNET_KEY}" not in cfg
+
 
 class CliTests(unittest.TestCase):