From 8d77e3ea542ee6dff05d3286096d274958cb3e00 Mon Sep 17 00:00:00 2001
From: Dima Kamalov <dimakamalov@asana.com>
Date: Sat, 6 Oct 2018 10:53:26 -0700
Subject: [PATCH] [AIRFLOW-3168] add healthcheck for mysql setup in CI

Make sure mysql is available before calling it in CI setup
---
 scripts/ci/{3-setup-databases.sh => 3-setup-mysql.sh} | 9 ++++++++-
 scripts/ci/airflow_travis.cfg                         | 4 ++--
 tests/configuration.py                                | 5 +++++
 tests/executors/test_celery_executor.py               | 5 +++++
 tests/sensors/test_sql_sensor.py                      | 4 ++++
 tox.ini                                               | 4 +++-
 6 files changed, 27 insertions(+), 4 deletions(-)
 rename scripts/ci/{3-setup-databases.sh => 3-setup-mysql.sh} (79%)

diff --git a/scripts/ci/3-setup-databases.sh b/scripts/ci/3-setup-mysql.sh
similarity index 79%
rename from scripts/ci/3-setup-databases.sh
rename to scripts/ci/3-setup-mysql.sh
index 2a5cb682e0119..ca5c1aa9f1b84 100755
--- a/scripts/ci/3-setup-databases.sh
+++ b/scripts/ci/3-setup-mysql.sh
@@ -20,4 +20,11 @@ set -exuo pipefail
 
 MYSQL_HOST=mysql
 
-mysql -h ${MYSQL_HOST} -u root -e 'drop database if exists airflow; create database airflow'
+retries=3
+for ((i=0; i<retries; i++)); do
+    mysql -h ${MYSQL_HOST} -u root -e 'drop database if exists airflow; create database airflow' && exit 0
+    echo "mysql db creation did not succeed, waiting 5 seconds to retry"
+    sleep 5
+done
+
+echo "mysql db creation could not succeed" && exit 1
diff --git a/scripts/ci/airflow_travis.cfg b/scripts/ci/airflow_travis.cfg
index 6895c4d5ad05f..8adf50b9280c0 100644
--- a/scripts/ci/airflow_travis.cfg
+++ b/scripts/ci/airflow_travis.cfg
@@ -21,7 +21,7 @@ airflow_home = ~/airflow
 dags_folder = ~/airflow/dags
 base_log_folder = ~/airflow/logs
 executor = LocalExecutor
-sql_alchemy_conn = mysql://root@mysql/airflow
+sql_alchemy_conn = # overridden by tox.ini
 unit_test_mode = True
 load_examples = True
 donot_pickle = False
@@ -56,7 +56,7 @@ celery_app_name = airflow.executors.celery_executor
 worker_concurrency = 16
 worker_log_server_port = 8793
 broker_url = amqp://guest:guest@rabbitmq:5672/
-result_backend = db+mysql://root@mysql/airflow
+result_backend = # overridden by tox.ini
 flower_port = 5555
 default_queue = default
 sync_parallelism = 0
diff --git a/tests/configuration.py b/tests/configuration.py
index b50a83259add4..e94491ee4e0b7 100644
--- a/tests/configuration.py
+++ b/tests/configuration.py
@@ -191,4 +191,9 @@ def test_deprecated_options_cmd(self):
         conf.set('celery', 'celery_result_backend_cmd', '/bin/echo 99')
 
         with self.assertWarns(DeprecationWarning):
+            tmp = None
+            if 'AIRFLOW__CELERY__RESULT_BACKEND' in os.environ:
+                tmp = os.environ.pop('AIRFLOW__CELERY__RESULT_BACKEND')
             self.assertEquals(conf.getint('celery', 'result_backend'), 99)
+            if tmp:
+                os.environ['AIRFLOW__CELERY__RESULT_BACKEND'] = tmp
diff --git a/tests/executors/test_celery_executor.py b/tests/executors/test_celery_executor.py
index 5c70b5bf20d70..f1b6a429faf3d 100644
--- a/tests/executors/test_celery_executor.py
+++ b/tests/executors/test_celery_executor.py
@@ -26,11 +26,16 @@
 from airflow.executors.celery_executor import CELERY_FETCH_ERR_MSG_HEADER
 from airflow.utils.state import State
 
+from airflow import configuration
+configuration.load_test_config()
+
 # leave this it is used by the test worker
 import celery.contrib.testing.tasks  # noqa: F401
 
 
 class CeleryExecutorTest(unittest.TestCase):
+    @unittest.skipIf('sqlite' in configuration.conf.get('core', 'sql_alchemy_conn'),
+                     "sqlite is configured with SequentialExecutor")
     def test_celery_integration(self):
         executor = CeleryExecutor()
         executor.start()
diff --git a/tests/sensors/test_sql_sensor.py b/tests/sensors/test_sql_sensor.py
index 433eaa61e9146..ec909431f03a2 100644
--- a/tests/sensors/test_sql_sensor.py
+++ b/tests/sensors/test_sql_sensor.py
@@ -38,6 +38,8 @@ def setUp(self):
         }
         self.dag = DAG(TEST_DAG_ID, default_args=args)
 
+    @unittest.skipUnless('mysql' in configuration.conf.get('core', 'sql_alchemy_conn'),
+                        "this is a mysql test")
     def test_sql_sensor_mysql(self):
         t = SqlSensor(
             task_id='sql_sensor_check',
@@ -47,6 +49,8 @@ def test_sql_sensor_mysql(self):
         )
         t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
 
+    @unittest.skipUnless('postgresql' in configuration.conf.get('core', 'sql_alchemy_conn'),
+                        "this is a postgres test")
     def test_sql_sensor_postgres(self):
         t = SqlSensor(
             task_id='sql_sensor_check',
diff --git a/tox.ini b/tox.ini
index f07641f29b38a..17bfeb6ec2e78 100644
--- a/tox.ini
+++ b/tox.ini
@@ -48,7 +48,9 @@ setenv =
   KRB5_CONFIG=/etc/krb5.conf
   KRB5_KTNAME=/etc/airflow.keytab
   backend_mysql: AIRFLOW__CORE__SQL_ALCHEMY_CONN=mysql://root@mysql/airflow
+  backend_mysql: AIRFLOW__CELERY__RESULT_BACKEND=db+mysql://root@mysql/airflow
   backend_postgres: AIRFLOW__CORE__SQL_ALCHEMY_CONN=postgresql+psycopg2://postgres:airflow@postgres/airflow
+  backend_postgres: AIRFLOW__CELERY__RESULT_BACKEND=db+postgresql://postgres:airflow@postgres/airflow
   backend_sqlite: AIRFLOW__CORE__SQL_ALCHEMY_CONN=sqlite:///{homedir}/airflow.db
   backend_sqlite: AIRFLOW__CORE__EXECUTOR=SequentialExecutor
 
@@ -59,7 +61,7 @@ commands =
   pip install --progress-bar off --find-links={homedir}/.wheelhouse --no-index -e .[devel_ci]
   env_docker: {toxinidir}/scripts/ci/1-setup-env.sh
   env_docker: {toxinidir}/scripts/ci/2-setup-kdc.sh
-  env_docker: {toxinidir}/scripts/ci/3-setup-databases.sh
+  backend_mysql: {toxinidir}/scripts/ci/3-setup-mysql.sh
   {toxinidir}/scripts/ci/5-run-tests.sh []
   {toxinidir}/scripts/ci/6-check-license.sh
   codecov -e TOXENV